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)

web/src/app/admin/connectors/google-drive/Credential.tsx

@@ -0,0 +1,432 @@
+import { Button } from "@/components/Button";
+import { PopupSpec } from "@/components/admin/connectors/Popup";
+import { useState } from "react";
+import { useSWRConfig } from "swr";
+import * as Yup from "yup";
+import { useRouter } from "next/navigation";
+import {
+  Credential,
+  GoogleDriveCredentialJson,
+  GoogleDriveServiceAccountCredentialJson,
+} from "@/lib/types";
+import { deleteCredential } from "@/lib/credential";
+import { setupGoogleDriveOAuth } from "@/lib/googleDrive";
+import { GOOGLE_DRIVE_AUTH_IS_ADMIN_COOKIE_NAME } from "@/lib/constants";
+import Cookies from "js-cookie";
+import { TextFormField } from "@/components/admin/connectors/Field";
+import { Form, Formik } from "formik";
+
+type GoogleDriveCredentialJsonTypes = "authorized_user" | "service_account";
+
+const DriveJsonUpload = ({
+  setPopup,
+}: {
+  setPopup: (popupSpec: PopupSpec | null) => void;
+}) => {
+  const { mutate } = useSWRConfig();
+  const [credentialJsonStr, setCredentialJsonStr] = useState<
+    string | undefined
+  >();
+
+  return (
+    <>
+      <input
+        className={
+          "mr-3 text-sm text-gray-900 border border-gray-300 rounded-lg " +
+          "cursor-pointer bg-gray-50 dark:text-gray-400 focus:outline-none " +
+          "dark:bg-gray-700 dark:border-gray-600 dark:placeholder-gray-400"
+        }
+        type="file"
+        accept=".json"
+        onChange={(event) => {
+          if (!event.target.files) {
+            return;
+          }
+          const file = event.target.files[0];
+          const reader = new FileReader();
+
+          reader.onload = function (loadEvent) {
+            if (!loadEvent?.target?.result) {
+              return;
+            }
+            const fileContents = loadEvent.target.result;
+            setCredentialJsonStr(fileContents as string);
+          };
+
+          reader.readAsText(file);
+        }}
+      />
+
+      <Button
+        disabled={!credentialJsonStr}
+        onClick={async () => {
+          // check if the JSON is a app credential or a service account credential
+          let credentialFileType: GoogleDriveCredentialJsonTypes;
+          try {
+            const appCredentialJson = JSON.parse(credentialJsonStr!);
+            if (appCredentialJson.web) {
+              credentialFileType = "authorized_user";
+            } else if (appCredentialJson.type === "service_account") {
+              credentialFileType = "service_account";
+            } else {
+              throw new Error(
+                "Unknown credential type, expected one of 'OAuth Web application' or 'Service Account'"
+              );
+            }
+          } catch (e) {
+            setPopup({
+              message: `Invalid file provided - ${e}`,
+              type: "error",
+            });
+            return;
+          }
+
+          if (credentialFileType === "authorized_user") {
+            const response = await fetch(
+              "/api/manage/admin/connector/google-drive/app-credential",
+              {
+                method: "PUT",
+                headers: {
+                  "Content-Type": "application/json",
+                },
+                body: credentialJsonStr,
+              }
+            );
+            if (response.ok) {
+              setPopup({
+                message: "Successfully uploaded app credentials",
+                type: "success",
+              });
+            } else {
+              const errorMsg = await response.text();
+              setPopup({
+                message: `Failed to upload app credentials - ${errorMsg}`,
+                type: "error",
+              });
+            }
+            mutate("/api/manage/admin/connector/google-drive/app-credential");
+          }
+
+          if (credentialFileType === "service_account") {
+            const response = await fetch(
+              "/api/manage/admin/connector/google-drive/service-account-key",
+              {
+                method: "PUT",
+                headers: {
+                  "Content-Type": "application/json",
+                },
+                body: credentialJsonStr,
+              }
+            );
+            if (response.ok) {
+              setPopup({
+                message: "Successfully uploaded app credentials",
+                type: "success",
+              });
+            } else {
+              const errorMsg = await response.text();
+              setPopup({
+                message: `Failed to upload app credentials - ${errorMsg}`,
+                type: "error",
+              });
+            }
+            mutate(
+              "/api/manage/admin/connector/google-drive/service-account-key"
+            );
+          }
+        }}
+      >
+        Upload
+      </Button>
+    </>
+  );
+};
+
+interface DriveJsonUploadSectionProps {
+  setPopup: (popupSpec: PopupSpec | null) => void;
+  appCredentialData?: { client_id: string };
+  serviceAccountCredentialData?: { service_account_email: string };
+}
+
+export const DriveJsonUploadSection = ({
+  setPopup,
+  appCredentialData,
+  serviceAccountCredentialData,
+}: DriveJsonUploadSectionProps) => {
+  const { mutate } = useSWRConfig();
+
+  if (serviceAccountCredentialData?.service_account_email) {
+    return (
+      <div className="mt-2 text-sm">
+        <div>
+          Found existing service account key with the following <b>Email:</b>
+          <p className="italic mt-1">
+            {serviceAccountCredentialData.service_account_email}
+          </p>
+        </div>
+        <div className="mt-4 mb-1">
+          If you want to update these credentials, delete the existing
+          credentials through the button below, and then upload a new
+          credentials JSON.
+        </div>
+        <Button
+          onClick={async () => {
+            const response = await fetch(
+              "/api/manage/admin/connector/google-drive/service-account-key",
+              {
+                method: "DELETE",
+              }
+            );
+            if (response.ok) {
+              mutate(
+                "/api/manage/admin/connector/google-drive/service-account-key"
+              );
+              setPopup({
+                message: "Successfully deleted service account key",
+                type: "success",
+              });
+            } else {
+              const errorMsg = await response.text();
+              setPopup({
+                message: `Failed to delete service account key - ${errorMsg}`,
+                type: "error",
+              });
+            }
+          }}
+        >
+          Delete
+        </Button>
+      </div>
+    );
+  }
+
+  if (appCredentialData?.client_id) {
+    return (
+      <div className="mt-2 text-sm">
+        <div>
+          Found existing app credentials with the following <b>Client ID:</b>
+          <p className="italic mt-1">{appCredentialData.client_id}</p>
+        </div>
+        <div className="mt-4 mb-1">
+          If you want to update these credentials, delete the existing
+          credentials through the button below, and then upload a new
+          credentials JSON.
+        </div>
+        <Button
+          onClick={async () => {
+            const response = await fetch(
+              "/api/manage/admin/connector/google-drive/app-credential",
+              {
+                method: "DELETE",
+              }
+            );
+            if (response.ok) {
+              mutate("/api/manage/admin/connector/google-drive/app-credential");
+              setPopup({
+                message: "Successfully deleted service account key",
+                type: "success",
+              });
+            } else {
+              const errorMsg = await response.text();
+              setPopup({
+                message: `Failed to delete app credential - ${errorMsg}`,
+                type: "error",
+              });
+            }
+          }}
+        >
+          Delete
+        </Button>
+      </div>
+    );
+  }
+
+  return (
+    <div className="mt-2">
+      <p className="text-sm mb-2">
+        Follow the guide{" "}
+        <a
+          className="text-blue-500"
+          target="_blank"
+          href="https://docs.danswer.dev/connectors/google_drive#authorization"
+        >
+          here
+        </a>{" "}
+        to either (1) setup a google OAuth App in your company workspace or (2)
+        create a Service Account.
+        <br />
+        <br />
+        Download the credentials JSON if choosing option (1) or the Service
+        Account key JSON if chooosing option (2), and upload it here.
+      </p>
+      <DriveJsonUpload setPopup={setPopup} />
+    </div>
+  );
+};
+
+interface DriveCredentialSectionProps {
+  googleDrivePublicCredential?: Credential<GoogleDriveCredentialJson>;
+  googleDriveServiceAccountCredential?: Credential<GoogleDriveServiceAccountCredentialJson>;
+  serviceAccountKeyData?: { service_account_email: string };
+  appCredentialData?: { client_id: string };
+  setPopup: (popupSpec: PopupSpec | null) => void;
+}
+
+export const DriveOAuthSection = ({
+  googleDrivePublicCredential,
+  googleDriveServiceAccountCredential,
+  serviceAccountKeyData,
+  appCredentialData,
+  setPopup,
+}: DriveCredentialSectionProps) => {
+  const { mutate } = useSWRConfig();
+  const router = useRouter();
+
+  const existingCredential =
+    googleDrivePublicCredential || googleDriveServiceAccountCredential;
+  if (existingCredential) {
+    return (
+      <>
+        <p className="mb-2 text-sm">
+          <i>Existing credential already setup!</i>
+        </p>
+        <Button
+          onClick={async () => {
+            await deleteCredential(existingCredential.id);
+            setPopup({
+              message: "Successfully revoked access to Google Drive!",
+              type: "success",
+            });
+            mutate("/api/manage/credential");
+          }}
+        >
+          Revoke Access
+        </Button>
+      </>
+    );
+  }
+
+  if (serviceAccountKeyData?.service_account_email) {
+    return (
+      <div>
+        <p className="text-sm mb-2">
+          When using a Google Drive Service Account, you can either have Danswer
+          act as the service account itself OR you can specify an account for
+          the service account to impersonate.
+          <br />
+          <br />
+          If you want to use the service account itself, leave the{" "}
+          <b>&apos;User email to impersonate&apos;</b> field blank when
+          submitting. If you do choose this option, make sure you have shared
+          the documents you want to index with the service account.
+        </p>
+
+        <div className="border-solid border-gray-600 border rounded-md p-6 mt-2 mb-4">
+          <Formik
+            initialValues={{
+              google_drive_delegated_user: "",
+            }}
+            validationSchema={Yup.object().shape({
+              google_drive_delegated_user: Yup.string().optional(),
+            })}
+            onSubmit={async (values, formikHelpers) => {
+              formikHelpers.setSubmitting(true);
+
+              const response = await fetch(
+                "/api/manage/admin/connector/google-drive/service-account-credential",
+                {
+                  method: "PUT",
+                  headers: {
+                    "Content-Type": "application/json",
+                  },
+                  body: JSON.stringify({
+                    google_drive_delegated_user:
+                      values.google_drive_delegated_user,
+                  }),
+                }
+              );
+
+              if (response.ok) {
+                setPopup({
+                  message: "Successfully created service account credential",
+                  type: "success",
+                });
+              } else {
+                const errorMsg = await response.text();
+                setPopup({
+                  message: `Failed to create service account credential - ${errorMsg}`,
+                  type: "error",
+                });
+              }
+              mutate("/api/manage/credential");
+            }}
+          >
+            {({ isSubmitting }) => (
+              <Form>
+                <TextFormField
+                  name="google_drive_delegated_user"
+                  label="[Optional] User email to impersonate:"
+                  subtext="If left blank, Danswer will use the service account itself."
+                />
+                <div className="flex">
+                  <button
+                    type="submit"
+                    disabled={isSubmitting}
+                    className={
+                      "bg-slate-500 hover:bg-slate-700 text-white " +
+                      "font-bold py-2 px-4 rounded focus:outline-none " +
+                      "focus:shadow-outline w-full max-w-sm mx-auto"
+                    }
+                  >
+                    Submit
+                  </button>
+                </div>
+              </Form>
+            )}
+          </Formik>
+        </div>
+      </div>
+    );
+  }
+
+  if (appCredentialData?.client_id) {
+    return (
+      <div className="text-sm mb-4">
+        <p className="mb-2">
+          Next, you must provide credentials via OAuth. This gives us read
+          access to the docs you have access to in your google drive account.
+        </p>
+        <Button
+          onClick={async () => {
+            const [authUrl, errorMsg] = await setupGoogleDriveOAuth({
+              isPublic: true,
+            });
+            if (authUrl) {
+              // cookie used by callback to determine where to finally redirect to
+              Cookies.set(GOOGLE_DRIVE_AUTH_IS_ADMIN_COOKIE_NAME, "true", {
+                path: "/",
+              });
+              router.push(authUrl);
+              return;
+            }
+
+            setPopup({
+              message: errorMsg,
+              type: "error",
+            });
+          }}
+        >
+          Authenticate with Google Drive
+        </Button>
+      </div>
+    );
+  }
+
+  // case where no keys have been uploaded in step 1
+  return (
+    <p className="text-sm">
+      Please upload either a OAuth Client Credential JSON or a Google Drive
+      Service Account Key JSON in Step 1 before moving onto Step 2.
+    </p>
+  );
+};

web/src/app/admin/connectors/google-drive/page.tsx

@@ -5,21 +5,17 @@ import { GoogleDriveIcon } from "@/components/icons/icons";
 import useSWR, { useSWRConfig } from "swr";
 import { fetcher } from "@/lib/fetcher";
 import { LoadingAnimation } from "@/components/Loading";
-import { useRouter } from "next/navigation";
 import { Popup, PopupSpec } from "@/components/admin/connectors/Popup";
 import { useState } from "react";
 import { HealthCheckBanner } from "@/components/health/healthcheck";
-import { Button } from "@/components/Button";
 import {
   ConnectorIndexingStatus,
   Credential,
   GoogleDriveConfig,
   GoogleDriveCredentialJson,
+  GoogleDriveServiceAccountCredentialJson,
 } from "@/lib/types";
-import { setupGoogleDriveOAuth } from "@/lib/googleDrive";
-import Cookies from "js-cookie";
-import { GOOGLE_DRIVE_AUTH_IS_ADMIN_COOKIE_NAME } from "@/lib/constants";
-import { deleteCredential, linkCredential } from "@/lib/credential";
+import { linkCredential } from "@/lib/credential";
 import { ConnectorForm } from "@/components/admin/connectors/ConnectorForm";
 import {
   BooleanFormField,
@@ -27,81 +23,11 @@ import {
 } from "@/components/admin/connectors/Field";
 import { GoogleDriveConnectorsTable } from "./GoogleDriveConnectorsTable";
 import { googleDriveConnectorNameBuilder } from "./utils";
-
-const AppCredentialUpload = ({
-  setPopup,
-}: {
-  setPopup: (popupSpec: PopupSpec | null) => void;
-}) => {
-  const [appCredentialJsonStr, setAppCredentialJsonStr] = useState<
-    string | undefined
-  >();
-
-  return (
-    <>
-      <input
-        className={
-          "mr-3 text-sm text-gray-900 border border-gray-300 rounded-lg " +
-          "cursor-pointer bg-gray-50 dark:text-gray-400 focus:outline-none " +
-          "dark:bg-gray-700 dark:border-gray-600 dark:placeholder-gray-400"
-        }
-        type="file"
-        accept=".json"
-        onChange={(event) => {
-          if (!event.target.files) {
-            return;
-          }
-          const file = event.target.files[0];
-          const reader = new FileReader();
-
-          reader.onload = function (loadEvent) {
-            if (!loadEvent?.target?.result) {
-              return;
-            }
-            const fileContents = loadEvent.target.result;
-            setAppCredentialJsonStr(fileContents as string);
-          };
-
-          reader.readAsText(file);
-        }}
-      />
-
-      <Button
-        disabled={!appCredentialJsonStr}
-        onClick={async () => {
-          const response = await fetch(
-            "/api/manage/admin/connector/google-drive/app-credential",
-            {
-              method: "PUT",
-              headers: {
-                "Content-Type": "application/json",
-              },
-              body: appCredentialJsonStr,
-            }
-          );
-          if (response.ok) {
-            setPopup({
-              message: "Successfully uploaded app credentials",
-              type: "success",
-            });
-          } else {
-            setPopup({
-              message: `Failed to upload app credentials - ${response.status}`,
-              type: "error",
-            });
-          }
-        }}
-      >
-        Upload
-      </Button>
-    </>
-  );
-};
+import { DriveOAuthSection, DriveJsonUploadSection } from "./Credential";
 
 interface GoogleDriveConnectorManagementProps {
-  googleDrivePublicCredential:
-    | Credential<GoogleDriveCredentialJson>
-    | undefined;
+  googleDrivePublicCredential?: Credential<GoogleDriveCredentialJson>;
+  googleDriveServiceAccountCredential?: Credential<GoogleDriveServiceAccountCredentialJson>;
   googleDriveConnectorIndexingStatus: ConnectorIndexingStatus<
     GoogleDriveConfig,
     GoogleDriveCredentialJson
@@ -116,6 +42,7 @@ interface GoogleDriveConnectorManagementProps {
 
 const GoogleDriveConnectorManagement = ({
   googleDrivePublicCredential,
+  googleDriveServiceAccountCredential,
   googleDriveConnectorIndexingStatus,
   googleDriveConnectorIndexingStatuses,
   credentialIsLinked,
@@ -123,7 +50,9 @@ const GoogleDriveConnectorManagement = ({
 }: GoogleDriveConnectorManagementProps) => {
   const { mutate } = useSWRConfig();
 
-  if (!googleDrivePublicCredential) {
+  const liveCredential =
+    googleDrivePublicCredential || googleDriveServiceAccountCredential;
+  if (!liveCredential) {
     return (
       <p className="text-sm">
         Please authenticate with Google Drive as described in Step 2! Once done
@@ -223,7 +152,7 @@ const GoogleDriveConnectorManagement = ({
   return (
     <div>
       <div className="text-sm">
-        <p className="my-3">
+        <div className="my-3">
           {googleDriveConnectorIndexingStatuses.length > 0 ? (
             <>
               Checkout the{" "}
@@ -239,7 +168,7 @@ const GoogleDriveConnectorManagement = ({
               latest documents from Google Drive every <b>10</b> minutes.
             </p>
           )}
-        </p>
+        </div>
       </div>
       {googleDriveConnectorIndexingStatuses.length > 0 && (
         <>
@@ -310,10 +239,7 @@ const GoogleDriveConnectorManagement = ({
           refreshFreq={10 * 60} // 10 minutes
           onSubmit={async (isSuccess, responseJson) => {
             if (isSuccess && responseJson) {
-              await linkCredential(
-                responseJson.id,
-                googleDrivePublicCredential.id
-              );
+              await linkCredential(responseJson.id, liveCredential.id);
               mutate("/api/manage/admin/connector/indexing-status");
             }
           }}
@@ -324,9 +250,6 @@ const GoogleDriveConnectorManagement = ({
 };
 
 const Main = () => {
-  const router = useRouter();
-  const { mutate } = useSWRConfig();
-
   const {
     data: appCredentialData,
     isLoading: isAppCredentialLoading,
@@ -335,6 +258,14 @@ const Main = () => {
     "/api/manage/admin/connector/google-drive/app-credential",
     fetcher
   );
+  const {
+    data: serviceAccountKeyData,
+    isLoading: isServiceAccountKeyLoading,
+    error: isServiceAccountKeyError,
+  } = useSWR<{ service_account_email: string }>(
+    "/api/manage/admin/connector/google-drive/service-account-key",
+    fetcher
+  );
   const {
     data: connectorIndexingStatuses,
     isLoading: isConnectorIndexingStatusesLoading,
@@ -347,10 +278,7 @@ const Main = () => {
     data: credentialsData,
     isLoading: isCredentialsLoading,
     error: isCredentialsError,
-  } = useSWR<Credential<GoogleDriveCredentialJson>[]>(
-    "/api/manage/credential",
-    fetcher
-  );
+  } = useSWR<Credential<any>[]>("/api/manage/credential", fetcher);
 
   const [popup, setPopup] = useState<{
     message: string;
@@ -365,6 +293,7 @@ const Main = () => {
 
   if (
     (!appCredentialData && isAppCredentialLoading) ||
+    (!serviceAccountKeyData && isServiceAccountKeyLoading) ||
     (!connectorIndexingStatuses && isConnectorIndexingStatusesLoading) ||
     (!credentialsData && isCredentialsLoading)
   ) {
@@ -391,7 +320,7 @@ const Main = () => {
     );
   }
 
-  if (isAppCredentialError) {
+  if (isAppCredentialError || isServiceAccountKeyError) {
     return (
       <div className="mx-auto">
         <div className="text-red-500">
@@ -401,10 +330,17 @@ const Main = () => {
     );
   }
 
-  const googleDrivePublicCredential = credentialsData.find(
+  const googleDrivePublicCredential:
+    | Credential<GoogleDriveCredentialJson>
+    | undefined = credentialsData.find(
     (credential) =>
       credential.credential_json?.google_drive_tokens && credential.public_doc
   );
+  const googleDriveServiceAccountCredential:
+    | Credential<GoogleDriveServiceAccountCredentialJson>
+    | undefined = credentialsData.find(
+    (credential) => credential.credential_json?.google_drive_service_account_key
+  );
   const googleDriveConnectorIndexingStatuses: ConnectorIndexingStatus<
     GoogleDriveConfig,
     GoogleDriveCredentialJson
@@ -416,127 +352,50 @@ const Main = () => {
     googleDriveConnectorIndexingStatuses[0];
 
   const credentialIsLinked =
-    googleDriveConnectorIndexingStatus !== undefined &&
-    googleDrivePublicCredential !== undefined &&
-    googleDriveConnectorIndexingStatus.connector.credential_ids.includes(
-      googleDrivePublicCredential.id
-    );
+    (googleDriveConnectorIndexingStatus !== undefined &&
+      googleDrivePublicCredential !== undefined &&
+      googleDriveConnectorIndexingStatus.connector.credential_ids.includes(
+        googleDrivePublicCredential.id
+      )) ||
+    (googleDriveConnectorIndexingStatus !== undefined &&
+      googleDriveServiceAccountCredential !== undefined &&
+      googleDriveConnectorIndexingStatus.connector.credential_ids.includes(
+        googleDriveServiceAccountCredential.id
+      ));
 
   return (
     <>
       {popup && <Popup message={popup.message} type={popup.type} />}
       <h2 className="font-bold mb-2 mt-6 ml-auto mr-auto">
-        Step 1: Provide your app Credentials
+        Step 1: Provide your Credentials
       </h2>
-      <div className="mt-2">
-        {appCredentialData?.client_id ? (
-          <div className="text-sm">
-            <div>
-              Found existing app credentials with the following{" "}
-              <b>Client ID:</b>
-              <p className="italic mt-1">{appCredentialData.client_id}</p>
-            </div>
-            <div className="mt-4">
-              If you want to update these credentials, upload a new
-              credentials.json file below.
-              <div className="mt-2">
-                <AppCredentialUpload
-                  setPopup={(popup) => {
-                    mutate(
-                      "/api/manage/admin/connector/google-drive/app-credential"
-                    );
-                    setPopupWithExpiration(popup);
-                  }}
-                />
-              </div>
-            </div>
-          </div>
-        ) : (
-          <>
-            <p className="text-sm">
-              Follow the guide{" "}
-              <a
-                className="text-blue-500"
-                target="_blank"
-                href="https://docs.danswer.dev/connectors/google_drive#authorization"
-              >
-                here
-              </a>{" "}
-              to setup your google app in your company workspace. Download the
-              credentials.json, and upload it here.
-            </p>
-            <AppCredentialUpload
-              setPopup={(popup) => {
-                mutate(
-                  "/api/manage/admin/connector/google-drive/app-credential"
-                );
-                setPopupWithExpiration(popup);
-              }}
-            />
-          </>
-        )}
-      </div>
+      <DriveJsonUploadSection
+        setPopup={setPopupWithExpiration}
+        appCredentialData={appCredentialData}
+        serviceAccountCredentialData={serviceAccountKeyData}
+      />
 
       <h2 className="font-bold mb-2 mt-6 ml-auto mr-auto">
         Step 2: Authenticate with Danswer
       </h2>
-      <div className="text-sm mb-4">
-        {googleDrivePublicCredential ? (
-          <>
-            <p className="mb-2">
-              <i>Existing credential already setup!</i>
-            </p>
-            <Button
-              onClick={async () => {
-                await deleteCredential(googleDrivePublicCredential.id);
-                setPopup({
-                  message: "Successfully revoked access to Google Drive!",
-                  type: "success",
-                });
-                mutate("/api/manage/credential");
-              }}
-            >
-              Revoke Access
-            </Button>
-          </>
-        ) : (
-          <>
-            <p className="mb-2">
-              Next, you must provide credentials via OAuth. This gives us read
-              access to the docs you have access to in your google drive
-              account.
-            </p>
-            <Button
-              onClick={async () => {
-                const [authUrl, errorMsg] = await setupGoogleDriveOAuth({
-                  isPublic: true,
-                });
-                if (authUrl) {
-                  // cookie used by callback to determine where to finally redirect to
-                  Cookies.set(GOOGLE_DRIVE_AUTH_IS_ADMIN_COOKIE_NAME, "true", {
-                    path: "/",
-                  });
-                  router.push(authUrl);
-                  return;
-                }
-
-                setPopup({
-                  message: errorMsg,
-                  type: "error",
-                });
-              }}
-            >
-              Authenticate with Google Drive
-            </Button>
-          </>
-        )}
-      </div>
+      <DriveOAuthSection
+        setPopup={setPopupWithExpiration}
+        googleDrivePublicCredential={googleDrivePublicCredential}
+        googleDriveServiceAccountCredential={
+          googleDriveServiceAccountCredential
+        }
+        appCredentialData={appCredentialData}
+        serviceAccountKeyData={serviceAccountKeyData}
+      />
 
       <h2 className="font-bold mb-2 mt-6 ml-auto mr-auto">
         Step 3: Start Indexing!
       </h2>
       <GoogleDriveConnectorManagement
         googleDrivePublicCredential={googleDrivePublicCredential}
+        googleDriveServiceAccountCredential={
+          googleDriveServiceAccountCredential
+        }
         googleDriveConnectorIndexingStatus={googleDriveConnectorIndexingStatus}
         googleDriveConnectorIndexingStatuses={
           googleDriveConnectorIndexingStatuses

web/src/components/admin/connectors/table/DeleteColumn.tsx

@@ -19,7 +19,6 @@ export function DeleteColumn<ConnectorConfigType, ConnectorCredentialType>({
   onUpdate,
 }: Props<ConnectorConfigType, ConnectorCredentialType>) {
   const [deleteHovered, setDeleteHovered] = useState<boolean>(false);
-  console.log(deleteHovered);
 
   const connector = connectorIndexingStatus.connector;
   const credential = connectorIndexingStatus.credential;

web/src/lib/types.ts

@@ -165,6 +165,11 @@ export interface GoogleDriveCredentialJson {
   google_drive_tokens: string;
 }
 
+export interface GoogleDriveServiceAccountCredentialJson {
+  google_drive_service_account_key: string;
+  google_drive_delegated_user: string;
+}
+
 export interface SlabCredentialJson {
   slab_bot_token: string;
 }