From 6897416fe615a3b113a064df4c5280d8a2cee5b0 Mon Sep 17 00:00:00 2001 From: Chris Weaver <25087905+Weves@users.noreply.github.com> Date: Thu, 24 Aug 2023 14:50:05 -0700 Subject: [PATCH] Support service accounts for Google Drive connector (#325) --- .../connectors/google_drive/connector.py | 100 +++- .../connectors/google_drive/connector_auth.py | 92 +++- .../connectors/google_drive/constants.py | 10 + backend/danswer/db/credentials.py | 15 + backend/danswer/server/manage.py | 106 ++++- backend/danswer/server/models.py | 18 + backend/danswer/server/utils.py | 3 +- .../connectors/google-drive/Credential.tsx | 432 ++++++++++++++++++ .../admin/connectors/google-drive/page.tsx | 259 +++-------- .../admin/connectors/table/DeleteColumn.tsx | 1 - web/src/lib/types.ts | 5 + 11 files changed, 797 insertions(+), 244 deletions(-) create mode 100644 backend/danswer/connectors/google_drive/constants.py create mode 100644 web/src/app/admin/connectors/google-drive/Credential.tsx diff --git a/backend/danswer/connectors/google_drive/connector.py b/backend/danswer/connectors/google_drive/connector.py index 1d462531c..86730adb0 100644 --- a/backend/danswer/connectors/google_drive/connector.py +++ b/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 diff --git a/backend/danswer/connectors/google_drive/connector_auth.py b/backend/danswer/connectors/google_drive/connector_auth.py index ea57a73ca..679b6fd77 100644 --- a/backend/danswer/connectors/google_drive/connector_auth.py +++ b/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) diff --git a/backend/danswer/connectors/google_drive/constants.py b/backend/danswer/connectors/google_drive/constants.py new file mode 100644 index 000000000..47dc402a3 --- /dev/null +++ b/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", +] diff --git a/backend/danswer/db/credentials.py b/backend/danswer/db/credentials.py index 30980876d..18e900dbd 100644 --- a/backend/danswer/db/credentials.py +++ b/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() diff --git a/backend/danswer/server/manage.py b/backend/danswer/server/manage.py index ab9175a59..4041d6831 100644 --- a/backend/danswer/server/manage.py +++ b/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) diff --git a/backend/danswer/server/models.py b/backend/danswer/server/models.py index f98dd1630..f79ce9072 100644 --- a/backend/danswer/server/models.py +++ b/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] diff --git a/backend/danswer/server/utils.py b/backend/danswer/server/utils.py index 4a3dc82f6..f18db93a7 100644 --- a/backend/danswer/server/utils.py +++ b/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) diff --git a/web/src/app/admin/connectors/google-drive/Credential.tsx b/web/src/app/admin/connectors/google-drive/Credential.tsx new file mode 100644 index 000000000..8ca6b1fbd --- /dev/null +++ b/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 ( + <> + { + 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); + }} + /> + + + + ); +}; + +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 ( +
+
+ Found existing service account key with the following Email: +

+ {serviceAccountCredentialData.service_account_email} +

+
+
+ If you want to update these credentials, delete the existing + credentials through the button below, and then upload a new + credentials JSON. +
+ +
+ ); + } + + if (appCredentialData?.client_id) { + return ( +
+
+ Found existing app credentials with the following Client ID: +

{appCredentialData.client_id}

+
+
+ If you want to update these credentials, delete the existing + credentials through the button below, and then upload a new + credentials JSON. +
+ +
+ ); + } + + return ( +
+

+ Follow the guide{" "} + + here + {" "} + to either (1) setup a google OAuth App in your company workspace or (2) + create a Service Account. +
+
+ Download the credentials JSON if choosing option (1) or the Service + Account key JSON if chooosing option (2), and upload it here. +

+ +
+ ); +}; + +interface DriveCredentialSectionProps { + googleDrivePublicCredential?: Credential; + googleDriveServiceAccountCredential?: Credential; + 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 ( + <> +

+ Existing credential already setup! +

+ + + ); + } + + if (serviceAccountKeyData?.service_account_email) { + return ( +
+

+ 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. +
+
+ If you want to use the service account itself, leave the{" "} + 'User email to impersonate' 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. +

+ +
+ { + 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 }) => ( +
+ +
+ +
+ + )} +
+
+
+ ); + } + + if (appCredentialData?.client_id) { + return ( +
+

+ Next, you must provide credentials via OAuth. This gives us read + access to the docs you have access to in your google drive account. +

+ +
+ ); + } + + // case where no keys have been uploaded in step 1 + return ( +

+ Please upload either a OAuth Client Credential JSON or a Google Drive + Service Account Key JSON in Step 1 before moving onto Step 2. +

+ ); +}; diff --git a/web/src/app/admin/connectors/google-drive/page.tsx b/web/src/app/admin/connectors/google-drive/page.tsx index e80503b00..ef69f3ae7 100644 --- a/web/src/app/admin/connectors/google-drive/page.tsx +++ b/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 ( - <> - { - 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); - }} - /> - - - - ); -}; +import { DriveOAuthSection, DriveJsonUploadSection } from "./Credential"; interface GoogleDriveConnectorManagementProps { - googleDrivePublicCredential: - | Credential - | undefined; + googleDrivePublicCredential?: Credential; + googleDriveServiceAccountCredential?: Credential; 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 (

Please authenticate with Google Drive as described in Step 2! Once done @@ -223,7 +152,7 @@ const GoogleDriveConnectorManagement = ({ return (

-

+

{googleDriveConnectorIndexingStatuses.length > 0 ? ( <> Checkout the{" "} @@ -239,7 +168,7 @@ const GoogleDriveConnectorManagement = ({ latest documents from Google Drive every 10 minutes.

)} -

+
{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[]>( - "/api/manage/credential", - fetcher - ); + } = useSWR[]>("/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 (
@@ -401,10 +330,17 @@ const Main = () => { ); } - const googleDrivePublicCredential = credentialsData.find( + const googleDrivePublicCredential: + | Credential + | undefined = credentialsData.find( (credential) => credential.credential_json?.google_drive_tokens && credential.public_doc ); + const googleDriveServiceAccountCredential: + | Credential + | 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 && }

- Step 1: Provide your app Credentials + Step 1: Provide your Credentials

-
- {appCredentialData?.client_id ? ( -
-
- Found existing app credentials with the following{" "} - Client ID: -

{appCredentialData.client_id}

-
-
- If you want to update these credentials, upload a new - credentials.json file below. -
- { - mutate( - "/api/manage/admin/connector/google-drive/app-credential" - ); - setPopupWithExpiration(popup); - }} - /> -
-
-
- ) : ( - <> -

- Follow the guide{" "} - - here - {" "} - to setup your google app in your company workspace. Download the - credentials.json, and upload it here. -

- { - mutate( - "/api/manage/admin/connector/google-drive/app-credential" - ); - setPopupWithExpiration(popup); - }} - /> - - )} -
+

Step 2: Authenticate with Danswer

-
- {googleDrivePublicCredential ? ( - <> -

- Existing credential already setup! -

- - - ) : ( - <> -

- Next, you must provide credentials via OAuth. This gives us read - access to the docs you have access to in your google drive - account. -

- - - )} -
+

Step 3: Start Indexing!

({ onUpdate, }: Props) { const [deleteHovered, setDeleteHovered] = useState(false); - console.log(deleteHovered); const connector = connectorIndexingStatus.connector; const credential = connectorIndexingStatus.credential; diff --git a/web/src/lib/types.ts b/web/src/lib/types.ts index abe079a40..86dd7c792 100644 --- a/web/src/lib/types.ts +++ b/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; }