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) <rkuo@onyx.app>

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,
     )
 
 

backend/onyx/server/documents/connector.py

@@ -53,8 +53,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(
@@ -987,7 +1021,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(
@@ -1013,7 +1052,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(

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(