From 3de743e64692e04df3377cf205a792f6d0cdf8d2 Mon Sep 17 00:00:00 2001 From: rkuo-danswer Date: Thu, 9 Jan 2025 18:10:59 -0800 Subject: [PATCH] possible fix for gdrive oauth in the cloud (#3642) * possible fix for gd oauth in the cloud * missed code in rename/merge --------- Co-authored-by: Richard Kuo (Danswer) --- .../onyx/connectors/google_utils/google_kv.py | 14 ++++++ backend/onyx/server/documents/connector.py | 50 +++++++++++++++++-- backend/onyx/server/documents/credential.py | 3 -- 3 files changed, 61 insertions(+), 6 deletions(-) diff --git a/backend/onyx/connectors/google_utils/google_kv.py b/backend/onyx/connectors/google_utils/google_kv.py index 96785e32544..4f714e98fb5 100644 --- a/backend/onyx/connectors/google_utils/google_kv.py +++ b/backend/onyx/connectors/google_utils/google_kv.py @@ -17,6 +17,9 @@ from onyx.configs.constants import KV_GOOGLE_DRIVE_SERVICE_ACCOUNT_KEY from onyx.connectors.google_utils.resources import get_drive_service from onyx.connectors.google_utils.resources import get_gmail_service +from onyx.connectors.google_utils.shared_constants import ( + DB_CREDENTIALS_AUTHENTICATION_METHOD, +) from onyx.connectors.google_utils.shared_constants import ( DB_CREDENTIALS_DICT_SERVICE_ACCOUNT_KEY, ) @@ -29,6 +32,9 @@ from onyx.connectors.google_utils.shared_constants import ( GOOGLE_SCOPES, ) +from onyx.connectors.google_utils.shared_constants import ( + GoogleOAuthAuthenticationMethod, +) from onyx.connectors.google_utils.shared_constants import ( MISSING_SCOPES_ERROR_STR, ) @@ -96,6 +102,7 @@ def update_credential_access_tokens( user: User, db_session: Session, source: DocumentSource, + auth_method: GoogleOAuthAuthenticationMethod, ) -> OAuthCredentials | None: app_credentials = get_google_app_cred(source) flow = InstalledAppFlow.from_client_config( @@ -119,6 +126,7 @@ def update_credential_access_tokens( new_creds_dict = { DB_CREDENTIALS_DICT_TOKEN_KEY: token_json_str, DB_CREDENTIALS_PRIMARY_ADMIN_KEY: email, + DB_CREDENTIALS_AUTHENTICATION_METHOD: auth_method.value, } if not update_credential_json(credential_id, new_creds_dict, user, db_session): @@ -129,6 +137,7 @@ def update_credential_access_tokens( def build_service_account_creds( source: DocumentSource, primary_admin_email: str | None = None, + name: str | None = None, ) -> CredentialBase: service_account_key = get_service_account_key(source=source) @@ -138,10 +147,15 @@ def build_service_account_creds( if primary_admin_email: credential_dict[DB_CREDENTIALS_PRIMARY_ADMIN_KEY] = primary_admin_email + credential_dict[ + DB_CREDENTIALS_AUTHENTICATION_METHOD + ] = GoogleOAuthAuthenticationMethod.UPLOADED.value + return CredentialBase( credential_json=credential_dict, admin_public=True, source=source, + name=name, ) diff --git a/backend/onyx/server/documents/connector.py b/backend/onyx/server/documents/connector.py index c93d5249908..5dc15efc823 100644 --- a/backend/onyx/server/documents/connector.py +++ b/backend/onyx/server/documents/connector.py @@ -54,8 +54,9 @@ upsert_service_account_key, ) from onyx.connectors.google_utils.google_kv import verify_csrf +from onyx.connectors.google_utils.shared_constants import DB_CREDENTIALS_DICT_TOKEN_KEY from onyx.connectors.google_utils.shared_constants import ( - DB_CREDENTIALS_DICT_TOKEN_KEY, + GoogleOAuthAuthenticationMethod, ) from onyx.db.connector import create_connector from onyx.db.connector import delete_connector @@ -314,6 +315,7 @@ def upsert_service_account_credential( credential_base = build_service_account_creds( DocumentSource.GOOGLE_DRIVE, primary_admin_email=service_account_credential_request.google_primary_admin, + name="Service Account (uploaded)", ) except KvKeyNotFoundError as e: raise HTTPException(status_code=400, detail=str(e)) @@ -408,6 +410,38 @@ def upload_files( return FileUploadResponse(file_paths=deduped_file_paths) +@router.get("/admin/connector") +def get_connectors_by_credential( + _: User = Depends(current_curator_or_admin_user), + db_session: Session = Depends(get_session), + credential: int | None = None, +) -> list[ConnectorSnapshot]: + """Get a list of connectors. Allow filtering by a specific credential id.""" + + connectors = fetch_connectors(db_session) + + filtered_connectors = [] + for connector in connectors: + if connector.source == DocumentSource.INGESTION_API: + # don't include INGESTION_API, as it's a system level + # connector not manageable by the user + continue + + if credential is not None: + found = False + for cc_pair in connector.credentials: + if credential == cc_pair.credential_id: + found = True + break + + if not found: + continue + + filtered_connectors.append(ConnectorSnapshot.from_connector_db_model(connector)) + + return filtered_connectors + + # Retrieves most recent failure cases for connectors that are currently failing @router.get("/admin/connector/failed-indexing-status") def get_currently_failed_indexing_status( @@ -969,7 +1003,12 @@ def gmail_callback( credential_id = int(credential_id_cookie) verify_csrf(credential_id, callback.state) credentials: Credentials | None = update_credential_access_tokens( - callback.code, credential_id, user, db_session, DocumentSource.GMAIL + callback.code, + credential_id, + user, + db_session, + DocumentSource.GMAIL, + GoogleOAuthAuthenticationMethod.UPLOADED, ) if credentials is None: raise HTTPException( @@ -995,7 +1034,12 @@ def google_drive_callback( verify_csrf(credential_id, callback.state) credentials: Credentials | None = update_credential_access_tokens( - callback.code, credential_id, user, db_session, DocumentSource.GOOGLE_DRIVE + callback.code, + credential_id, + user, + db_session, + DocumentSource.GOOGLE_DRIVE, + GoogleOAuthAuthenticationMethod.UPLOADED, ) if credentials is None: raise HTTPException( diff --git a/backend/onyx/server/documents/credential.py b/backend/onyx/server/documents/credential.py index 51d9643dc77..b68ee660cb7 100644 --- a/backend/onyx/server/documents/credential.py +++ b/backend/onyx/server/documents/credential.py @@ -9,7 +9,6 @@ from onyx.auth.users import current_user from onyx.db.credentials import alter_credential from onyx.db.credentials import cleanup_gmail_credentials -from onyx.db.credentials import cleanup_google_drive_credentials from onyx.db.credentials import create_credential from onyx.db.credentials import CREDENTIAL_PERMISSIONS_TO_IGNORE from onyx.db.credentials import delete_credential @@ -133,8 +132,6 @@ def create_credential_from_model( # Temporary fix for empty Google App credentials if credential_info.source == DocumentSource.GMAIL: cleanup_gmail_credentials(db_session=db_session) - if credential_info.source == DocumentSource.GOOGLE_DRIVE: - cleanup_google_drive_credentials(db_session=db_session) credential = create_credential(credential_info, user, db_session) return ObjectCreationIdResponse(