Support service accounts for Google Drive connector (#325)

backend/danswer/connectors/google_drive/connector.py

@@ -5,9 +5,10 @@ from collections.abc import Generator
 from collections.abc import Sequence
 from itertools import chain
 from typing import Any
+from typing import cast
 
 import docx2txt  # type:ignore
-from google.oauth2.credentials import Credentials  # type: ignore
+from google.auth.credentials import Credentials  # type: ignore
 from googleapiclient import discovery  # type: ignore
 from PyPDF2 import PdfReader
 
@@ -16,8 +17,19 @@ from danswer.configs.app_configs import GOOGLE_DRIVE_FOLLOW_SHORTCUTS
 from danswer.configs.app_configs import GOOGLE_DRIVE_INCLUDE_SHARED
 from danswer.configs.app_configs import INDEX_BATCH_SIZE
 from danswer.configs.constants import DocumentSource
-from danswer.connectors.google_drive.connector_auth import DB_CREDENTIALS_DICT_KEY
-from danswer.connectors.google_drive.connector_auth import get_drive_tokens
+from danswer.connectors.google_drive.connector_auth import (
+    get_google_drive_creds_for_authorized_user,
+)
+from danswer.connectors.google_drive.connector_auth import (
+    get_google_drive_creds_for_service_account,
+)
+from danswer.connectors.google_drive.constants import (
+    DB_CREDENTIALS_DICT_DELEGATED_USER_KEY,
+)
+from danswer.connectors.google_drive.constants import (
+    DB_CREDENTIALS_DICT_SERVICE_ACCOUNT_KEY,
+)
+from danswer.connectors.google_drive.constants import DB_CREDENTIALS_DICT_TOKEN_KEY
 from danswer.connectors.interfaces import GenerateDocumentsOutput
 from danswer.connectors.interfaces import LoadConnector
 from danswer.connectors.interfaces import PollConnector
@@ -31,10 +43,6 @@ logger = setup_logger()
 
 # allow 10 minutes for modifiedTime to get propogated
 DRIVE_START_TIME_OFFSET = 60 * 10
-SCOPES = [
-    "https://www.googleapis.com/auth/drive.readonly",
-    "https://www.googleapis.com/auth/drive.metadata.readonly",
-]
 SUPPORTED_DRIVE_DOC_TYPES = [
     "application/vnd.google-apps.document",
     "application/vnd.google-apps.spreadsheet",
@@ -335,16 +343,51 @@ class GoogleDriveConnector(LoadConnector, PollConnector):
 
         return folder_ids
 
-    def load_credentials(self, credentials: dict[str, Any]) -> dict[str, Any] | None:
-        access_token_json_str = credentials[DB_CREDENTIALS_DICT_KEY]
-        creds = get_drive_tokens(token_json_str=access_token_json_str)
+    def load_credentials(self, credentials: dict[str, Any]) -> dict[str, str] | None:
+        """Checks for two different types of credentials.
+        (1) A credential which holds a token acquired via a user going thorugh
+        the Google OAuth flow.
+        (2) A credential which holds a service account key JSON file, which
+        can then be used to impersonate any user in the workspace.
+        """
+        creds = None
+        new_creds_dict = None
+        if DB_CREDENTIALS_DICT_TOKEN_KEY in credentials:
+            access_token_json_str = cast(
+                str, credentials[DB_CREDENTIALS_DICT_TOKEN_KEY]
+            )
+            creds = get_google_drive_creds_for_authorized_user(
+                token_json_str=access_token_json_str
+            )
+
+            # tell caller to update token stored in DB if it has changed
+            # (e.g. the token has been refreshed)
+            new_creds_json_str = creds.to_json() if creds else ""
+            if new_creds_json_str != access_token_json_str:
+                new_creds_dict = {DB_CREDENTIALS_DICT_TOKEN_KEY: new_creds_json_str}
+
+        if DB_CREDENTIALS_DICT_SERVICE_ACCOUNT_KEY in credentials:
+            service_account_key_json_str = credentials[
+                DB_CREDENTIALS_DICT_SERVICE_ACCOUNT_KEY
+            ]
+            creds = get_google_drive_creds_for_service_account(
+                service_account_key_json_str=service_account_key_json_str
+            )
+
+            # "Impersonate" a user if one is specified
+            delegated_user_email = cast(
+                str | None, credentials.get(DB_CREDENTIALS_DICT_DELEGATED_USER_KEY)
+            )
+            if delegated_user_email:
+                creds = creds.with_subject(delegated_user_email) if creds else None
+
         if creds is None:
-            raise PermissionError("Unable to access Google Drive.")
+            raise PermissionError(
+                "Unable to access Google Drive - unknown credential structure."
+            )
+
         self.creds = creds
-        new_creds_json_str = creds.to_json()
-        if new_creds_json_str != access_token_json_str:
-            return {DB_CREDENTIALS_DICT_KEY: new_creds_json_str}
-        return None
+        return new_creds_dict
 
     def _fetch_docs_from_drive(
         self,
@@ -417,3 +460,30 @@ class GoogleDriveConnector(LoadConnector, PollConnector):
         yield from self._fetch_docs_from_drive(
             max(start - DRIVE_START_TIME_OFFSET, 0, 0), end
         )
+
+
+if __name__ == "__main__":
+    import json
+    import os
+
+    service_account_json_path = os.environ.get("GOOGLE_SERVICE_ACCOUNT_KEY_JSON_PATH")
+    if not service_account_json_path:
+        raise ValueError(
+            "Please set GOOGLE_SERVICE_ACCOUNT_KEY_JSON_PATH environment variable"
+        )
+    with open(service_account_json_path) as f:
+        creds = json.load(f)
+
+    credentials_dict = {
+        DB_CREDENTIALS_DICT_SERVICE_ACCOUNT_KEY: json.dumps(creds),
+    }
+    delegated_user = os.environ.get("GOOGLE_DRIVE_DELEGATED_USER")
+    if delegated_user:
+        credentials_dict[DB_CREDENTIALS_DICT_DELEGATED_USER_KEY] = delegated_user
+
+    connector = GoogleDriveConnector()
+    connector.load_credentials(credentials_dict)
+    document_batch_generator = connector.load_from_state()
+    for document_batch in document_batch_generator:
+        print(document_batch)
+        break

backend/danswer/connectors/google_drive/connector_auth.py

@@ -1,45 +1,48 @@
 import json
+from enum import Enum
 from typing import cast
 from urllib.parse import parse_qs
 from urllib.parse import ParseResult
 from urllib.parse import urlparse
 
 from google.auth.transport.requests import Request  # type: ignore
-from google.oauth2.credentials import Credentials  # type: ignore
+from google.oauth2.credentials import Credentials as OAuthCredentials  # type: ignore
+from google.oauth2.service_account import Credentials as ServiceAccountCredentials  # type: ignore
 from google_auth_oauthlib.flow import InstalledAppFlow  # type: ignore
 from sqlalchemy.orm import Session
 
 from danswer.configs.app_configs import WEB_DOMAIN
+from danswer.connectors.google_drive.constants import CRED_KEY
+from danswer.connectors.google_drive.constants import (
+    DB_CREDENTIALS_DICT_DELEGATED_USER_KEY,
+)
+from danswer.connectors.google_drive.constants import (
+    DB_CREDENTIALS_DICT_SERVICE_ACCOUNT_KEY,
+)
+from danswer.connectors.google_drive.constants import DB_CREDENTIALS_DICT_TOKEN_KEY
+from danswer.connectors.google_drive.constants import GOOGLE_DRIVE_CRED_KEY
+from danswer.connectors.google_drive.constants import GOOGLE_DRIVE_SERVICE_ACCOUNT_KEY
+from danswer.connectors.google_drive.constants import SCOPES
 from danswer.db.credentials import update_credential_json
 from danswer.db.models import User
 from danswer.dynamic_configs import get_dynamic_config_store
+from danswer.server.models import CredentialBase
 from danswer.server.models import GoogleAppCredentials
+from danswer.server.models import GoogleServiceAccountKey
 from danswer.utils.logger import setup_logger
 
 logger = setup_logger()
 
-DB_CREDENTIALS_DICT_KEY = "google_drive_tokens"
-CRED_KEY = "credential_id_{}"
-GOOGLE_DRIVE_CRED_KEY = "google_drive_app_credential"
-SCOPES = ["https://www.googleapis.com/auth/drive.readonly"]
-
 
 def _build_frontend_google_drive_redirect() -> str:
     return f"{WEB_DOMAIN}/admin/connectors/google-drive/auth/callback"
 
 
-def get_drive_tokens(
-    *, creds: Credentials | None = None, token_json_str: str | None = None
-) -> Credentials | None:
-    if creds is None and token_json_str is None:
-        return None
-
-    if token_json_str is not None:
-        creds_json = json.loads(token_json_str)
-        creds = Credentials.from_authorized_user_info(creds_json, SCOPES)
-
-    if not creds:
-        return None
+def get_google_drive_creds_for_authorized_user(
+    token_json_str: str,
+) -> OAuthCredentials | None:
+    creds_json = json.loads(token_json_str)
+    creds = OAuthCredentials.from_authorized_user_info(creds_json, SCOPES)
     if creds.valid:
         return creds
 
@@ -52,9 +55,22 @@ def get_drive_tokens(
         except Exception as e:
             logger.exception(f"Failed to refresh google drive access token due to: {e}")
             return None
+
     return None
 
 
+def get_google_drive_creds_for_service_account(
+    service_account_key_json_str: str,
+) -> ServiceAccountCredentials | None:
+    service_account_key = json.loads(service_account_key_json_str)
+    creds = ServiceAccountCredentials.from_service_account_info(
+        service_account_key, scopes=SCOPES
+    )
+    if not creds.valid or not creds.expired:
+        creds.refresh(Request())
+    return creds if creds.valid else None
+
+
 def verify_csrf(credential_id: int, state: str) -> None:
     csrf = get_dynamic_config_store().load(CRED_KEY.format(str(credential_id)))
     if csrf != state:
@@ -85,7 +101,7 @@ def update_credential_access_tokens(
     credential_id: int,
     user: User,
     db_session: Session,
-) -> Credentials | None:
+) -> OAuthCredentials | None:
     app_credentials = get_google_app_cred()
     flow = InstalledAppFlow.from_client_config(
         app_credentials.dict(),
@@ -95,13 +111,30 @@ def update_credential_access_tokens(
     flow.fetch_token(code=auth_code)
     creds = flow.credentials
     token_json_str = creds.to_json()
-    new_creds_dict = {DB_CREDENTIALS_DICT_KEY: token_json_str}
+    new_creds_dict = {DB_CREDENTIALS_DICT_TOKEN_KEY: token_json_str}
 
     if not update_credential_json(credential_id, new_creds_dict, user, db_session):
         return None
     return creds
 
 
+def build_service_account_creds(
+    delegated_user_email: str | None = None,
+) -> CredentialBase:
+    service_account_key = get_service_account_key()
+
+    credential_dict = {
+        DB_CREDENTIALS_DICT_SERVICE_ACCOUNT_KEY: service_account_key.json(),
+    }
+    if delegated_user_email:
+        credential_dict[DB_CREDENTIALS_DICT_DELEGATED_USER_KEY] = delegated_user_email
+
+    return CredentialBase(
+        credential_json=credential_dict,
+        public_doc=True,
+    )
+
+
 def get_google_app_cred() -> GoogleAppCredentials:
     creds_str = str(get_dynamic_config_store().load(GOOGLE_DRIVE_CRED_KEY))
     return GoogleAppCredentials(**json.loads(creds_str))
@@ -109,3 +142,22 @@ def get_google_app_cred() -> GoogleAppCredentials:
 
 def upsert_google_app_cred(app_credentials: GoogleAppCredentials) -> None:
     get_dynamic_config_store().store(GOOGLE_DRIVE_CRED_KEY, app_credentials.json())
+
+
+def delete_google_app_cred() -> None:
+    get_dynamic_config_store().delete(GOOGLE_DRIVE_CRED_KEY)
+
+
+def get_service_account_key() -> GoogleServiceAccountKey:
+    creds_str = str(get_dynamic_config_store().load(GOOGLE_DRIVE_SERVICE_ACCOUNT_KEY))
+    return GoogleServiceAccountKey(**json.loads(creds_str))
+
+
+def upsert_service_account_key(service_account_key: GoogleServiceAccountKey) -> None:
+    get_dynamic_config_store().store(
+        GOOGLE_DRIVE_SERVICE_ACCOUNT_KEY, service_account_key.json()
+    )
+
+
+def delete_service_account_key() -> None:
+    get_dynamic_config_store().delete(GOOGLE_DRIVE_SERVICE_ACCOUNT_KEY)

backend/danswer/connectors/google_drive/constants.py

@@ -0,0 +1,10 @@
+DB_CREDENTIALS_DICT_TOKEN_KEY = "google_drive_tokens"
+DB_CREDENTIALS_DICT_SERVICE_ACCOUNT_KEY = "google_drive_service_account_key"
+DB_CREDENTIALS_DICT_DELEGATED_USER_KEY = "google_drive_delegated_user"
+CRED_KEY = "credential_id_{}"
+GOOGLE_DRIVE_CRED_KEY = "google_drive_app_credential"
+GOOGLE_DRIVE_SERVICE_ACCOUNT_KEY = "google_drive_service_account_key"
+SCOPES = [
+    "https://www.googleapis.com/auth/drive.readonly",
+    "https://www.googleapis.com/auth/drive.metadata.readonly",
+]

backend/danswer/db/credentials.py

@@ -1,9 +1,13 @@
 from typing import Any
 
+from sqlalchemy import delete
 from sqlalchemy import select
 from sqlalchemy.orm import Session
 from sqlalchemy.sql.expression import or_
 
+from danswer.connectors.google_drive.constants import (
+    DB_CREDENTIALS_DICT_SERVICE_ACCOUNT_KEY,
+)
 from danswer.db.engine import get_sqlalchemy_engine
 from danswer.db.models import Credential
 from danswer.db.models import User
@@ -137,3 +141,14 @@ def create_initial_public_credential() -> None:
         )
         db_session.add(credential)
         db_session.commit()
+
+
+def delete_google_drive_service_account_credentials(
+    user: User | None, db_session: Session
+) -> None:
+    credentials = fetch_credentials(user, db_session)
+    for credential in credentials:
+        if credential.credential_json.get(DB_CREDENTIALS_DICT_SERVICE_ACCOUNT_KEY):
+            db_session.delete(credential)
+
+    db_session.commit()

backend/danswer/server/manage.py

@@ -17,17 +17,26 @@ from danswer.auth.users import current_admin_user
 from danswer.auth.users import current_user
 from danswer.configs.app_configs import DISABLE_GENERATIVE_AI
 from danswer.configs.app_configs import GENERATIVE_MODEL_ACCESS_CHECK_FREQ
-from danswer.configs.app_configs import MASK_CREDENTIAL_PREFIX
 from danswer.configs.constants import GEN_AI_API_KEY_STORAGE_KEY
 from danswer.connectors.file.utils import write_temp_files
-from danswer.connectors.google_drive.connector_auth import DB_CREDENTIALS_DICT_KEY
+from danswer.connectors.google_drive.connector_auth import build_service_account_creds
+from danswer.connectors.google_drive.connector_auth import (
+    DB_CREDENTIALS_DICT_SERVICE_ACCOUNT_KEY,
+)
+from danswer.connectors.google_drive.connector_auth import DB_CREDENTIALS_DICT_TOKEN_KEY
+from danswer.connectors.google_drive.connector_auth import delete_google_app_cred
+from danswer.connectors.google_drive.connector_auth import delete_service_account_key
 from danswer.connectors.google_drive.connector_auth import get_auth_url
-from danswer.connectors.google_drive.connector_auth import get_drive_tokens
 from danswer.connectors.google_drive.connector_auth import get_google_app_cred
+from danswer.connectors.google_drive.connector_auth import (
+    get_google_drive_creds_for_authorized_user,
+)
+from danswer.connectors.google_drive.connector_auth import get_service_account_key
 from danswer.connectors.google_drive.connector_auth import (
     update_credential_access_tokens,
 )
 from danswer.connectors.google_drive.connector_auth import upsert_google_app_cred
+from danswer.connectors.google_drive.connector_auth import upsert_service_account_key
 from danswer.connectors.google_drive.connector_auth import verify_csrf
 from danswer.db.connector import create_connector
 from danswer.db.connector import delete_connector
@@ -41,6 +50,7 @@ from danswer.db.connector_credential_pair import get_connector_credential_pairs
 from danswer.db.connector_credential_pair import remove_credential_from_connector
 from danswer.db.credentials import create_credential
 from danswer.db.credentials import delete_credential
+from danswer.db.credentials import delete_google_drive_service_account_credentials
 from danswer.db.credentials import fetch_credential_by_id
 from danswer.db.credentials import fetch_credentials
 from danswer.db.credentials import update_credential
@@ -71,6 +81,8 @@ from danswer.server.models import DeletionAttemptSnapshot
 from danswer.server.models import FileUploadResponse
 from danswer.server.models import GDriveCallback
 from danswer.server.models import GoogleAppCredentials
+from danswer.server.models import GoogleServiceAccountCredentialRequest
+from danswer.server.models import GoogleServiceAccountKey
 from danswer.server.models import IndexAttemptSnapshot
 from danswer.server.models import ObjectCreationIdResponse
 from danswer.server.models import RunConnectorRequest
@@ -117,7 +129,7 @@ def check_google_app_credentials_exist(
 
 
 @router.put("/admin/connector/google-drive/app-credential")
-def update_google_app_credentials(
+def upsert_google_app_credentials(
     app_credentials: GoogleAppCredentials, _: User = Depends(current_admin_user)
 ) -> StatusResponse:
     try:
@@ -130,6 +142,84 @@ def update_google_app_credentials(
     )
 
 
+@router.delete("/admin/connector/google-drive/app-credential")
+def delete_google_app_credentials(
+    _: User = Depends(current_admin_user),
+) -> StatusResponse:
+    try:
+        delete_google_app_cred()
+    except ConfigNotFoundError as e:
+        raise HTTPException(status_code=400, detail=str(e))
+
+    return StatusResponse(
+        success=True, message="Successfully deleted Google App Credentials"
+    )
+
+
+@router.get("/admin/connector/google-drive/service-account-key")
+def check_google_service_account_key_exist(
+    _: User = Depends(current_admin_user),
+) -> dict[str, str]:
+    try:
+        return {"service_account_email": get_service_account_key().client_email}
+    except ConfigNotFoundError:
+        raise HTTPException(
+            status_code=404, detail="Google Service Account Key not found"
+        )
+
+
+@router.put("/admin/connector/google-drive/service-account-key")
+def upsert_google_service_account_key(
+    service_account_key: GoogleServiceAccountKey, _: User = Depends(current_admin_user)
+) -> StatusResponse:
+    try:
+        upsert_service_account_key(service_account_key)
+    except ValueError as e:
+        raise HTTPException(status_code=400, detail=str(e))
+
+    return StatusResponse(
+        success=True, message="Successfully saved Google Service Account Key"
+    )
+
+
+@router.delete("/admin/connector/google-drive/service-account-key")
+def delete_google_service_account_key(
+    _: User = Depends(current_admin_user),
+) -> StatusResponse:
+    try:
+        delete_service_account_key()
+    except ConfigNotFoundError as e:
+        raise HTTPException(status_code=400, detail=str(e))
+
+    return StatusResponse(
+        success=True, message="Successfully deleted Google Service Account Key"
+    )
+
+
+@router.put("/admin/connector/google-drive/service-account-credential")
+def upsert_service_account_credential(
+    service_account_credential_request: GoogleServiceAccountCredentialRequest,
+    user: User = Depends(current_admin_user),
+    db_session: Session = Depends(get_session),
+) -> ObjectCreationIdResponse:
+    """Special API which allows the creation of a credential for a service account.
+    Combines the input with the saved service account key to create an entry in the
+    `Credential` table."""
+    try:
+        credential_base = build_service_account_creds(
+            delegated_user_email=service_account_credential_request.google_drive_delegated_user
+        )
+        print(credential_base)
+    except ConfigNotFoundError as e:
+        raise HTTPException(status_code=400, detail=str(e))
+
+    # first delete all existing service account credentials
+    delete_google_drive_service_account_credentials(user, db_session)
+    return create_credential(
+        credential_data=credential_base, user=user, db_session=db_session
+    )
+
+
 @router.get("/admin/connector/google-drive/check-auth/{credential_id}")
 def check_drive_tokens(
     credential_id: int,
@@ -139,11 +229,13 @@ def check_drive_tokens(
     db_credentials = fetch_credential_by_id(credential_id, user, db_session)
     if (
         not db_credentials
-        or DB_CREDENTIALS_DICT_KEY not in db_credentials.credential_json
+        or DB_CREDENTIALS_DICT_TOKEN_KEY not in db_credentials.credential_json
     ):
         return AuthStatus(authenticated=False)
-    token_json_str = str(db_credentials.credential_json[DB_CREDENTIALS_DICT_KEY])
-    google_drive_creds = get_drive_tokens(token_json_str=token_json_str)
+    token_json_str = str(db_credentials.credential_json[DB_CREDENTIALS_DICT_TOKEN_KEY])
+    google_drive_creds = get_google_drive_creds_for_authorized_user(
+        token_json_str=token_json_str
+    )
     if google_drive_creds is None:
         return AuthStatus(authenticated=False)
     return AuthStatus(authenticated=True)

backend/danswer/server/models.py

@@ -58,6 +58,24 @@ class GoogleAppCredentials(BaseModel):
     web: GoogleAppWebCredentials
 
 
+class GoogleServiceAccountKey(BaseModel):
+    type: str
+    project_id: str
+    private_key_id: str
+    private_key: str
+    client_email: str
+    client_id: str
+    auth_uri: str
+    token_uri: str
+    auth_provider_x509_cert_url: str
+    client_x509_cert_url: str
+    universe_domain: str
+
+
+class GoogleServiceAccountCredentialRequest(BaseModel):
+    google_drive_delegated_user: str | None  # email of user to impersonate
+
+
 class FileUploadResponse(BaseModel):
     file_paths: list[str]
 

backend/danswer/server/utils.py

@@ -10,7 +10,8 @@ def mask_credential_dict(credential_dict: dict[str, Any]) -> dict[str, str]:
     for key, val in credential_dict.items():
         if not isinstance(val, str):
             raise ValueError(
-                "Unable to mask credentials of type other than string, cannot process request."
+                f"Unable to mask credentials of type other than string, cannot process request."
+                f"Recieved type: {type(val)}"
             )
 
         masked_creds[key] = mask_string(val)