From e80a0f27166b21c05f0cb8341886ba5697c39d7c Mon Sep 17 00:00:00 2001 From: pablonyx Date: Thu, 27 Feb 2025 21:13:39 -0800 Subject: [PATCH] Improved google connector flow (#4155) * fix handling * k * k * fix function * k * k --- backend/onyx/db/credentials.py | 40 ++- backend/onyx/server/documents/credential.py | 7 +- .../[connector]/pages/gdrive/Credential.tsx | 103 +++++- .../pages/gdrive/GoogleDrivePage.tsx | 164 ++++----- .../[connector]/pages/gmail/Credential.tsx | 337 ++++++++++++------ .../[connector]/pages/gmail/GmailPage.tsx | 134 ++++--- web/src/components/Button.tsx | 3 - web/src/lib/googleConnector.ts | 120 +++++++ 8 files changed, 627 insertions(+), 281 deletions(-) create mode 100644 web/src/lib/googleConnector.ts diff --git a/backend/onyx/db/credentials.py b/backend/onyx/db/credentials.py index 40edbcef3..cbd578ff2 100644 --- a/backend/onyx/db/credentials.py +++ b/backend/onyx/db/credentials.py @@ -360,18 +360,13 @@ def backend_update_credential_json( db_session.commit() -def delete_credential( +def _delete_credential_internal( + credential: Credential, credential_id: int, - user: User | None, db_session: Session, force: bool = False, ) -> None: - credential = fetch_credential_by_id_for_user(credential_id, user, db_session) - if credential is None: - raise ValueError( - f"Credential by provided id {credential_id} does not exist or does not belong to user" - ) - + """Internal utility function to handle the actual deletion of a credential""" associated_connectors = ( db_session.query(ConnectorCredentialPair) .filter(ConnectorCredentialPair.credential_id == credential_id) @@ -416,6 +411,35 @@ def delete_credential( db_session.commit() +def delete_credential_for_user( + credential_id: int, + user: User, + db_session: Session, + force: bool = False, +) -> None: + """Delete a credential that belongs to a specific user""" + credential = fetch_credential_by_id_for_user(credential_id, user, db_session) + if credential is None: + raise ValueError( + f"Credential by provided id {credential_id} does not exist or does not belong to user" + ) + + _delete_credential_internal(credential, credential_id, db_session, force) + + +def delete_credential( + credential_id: int, + db_session: Session, + force: bool = False, +) -> None: + """Delete a credential regardless of ownership (admin function)""" + credential = fetch_credential_by_id(credential_id, db_session) + if credential is None: + raise ValueError(f"Credential by provided id {credential_id} does not exist") + + _delete_credential_internal(credential, credential_id, db_session, force) + + def create_initial_public_credential(db_session: Session) -> None: error_msg = ( "DB is not in a valid initial state." diff --git a/backend/onyx/server/documents/credential.py b/backend/onyx/server/documents/credential.py index 5060d15a3..9d8a9ead2 100644 --- a/backend/onyx/server/documents/credential.py +++ b/backend/onyx/server/documents/credential.py @@ -13,6 +13,7 @@ from onyx.db.credentials import cleanup_gmail_credentials from onyx.db.credentials import create_credential from onyx.db.credentials import CREDENTIAL_PERMISSIONS_TO_IGNORE from onyx.db.credentials import delete_credential +from onyx.db.credentials import delete_credential_for_user from onyx.db.credentials import fetch_credential_by_id_for_user from onyx.db.credentials import fetch_credentials_by_source_for_user from onyx.db.credentials import fetch_credentials_for_user @@ -88,7 +89,7 @@ def delete_credential_by_id_admin( db_session: Session = Depends(get_session), ) -> StatusResponse: """Same as the user endpoint, but can delete any credential (not just the user's own)""" - delete_credential(db_session=db_session, credential_id=credential_id, user=None) + delete_credential(db_session=db_session, credential_id=credential_id) return StatusResponse( success=True, message="Credential deleted successfully", data=credential_id ) @@ -242,7 +243,7 @@ def delete_credential_by_id( user: User = Depends(current_user), db_session: Session = Depends(get_session), ) -> StatusResponse: - delete_credential( + delete_credential_for_user( credential_id, user, db_session, @@ -259,7 +260,7 @@ def force_delete_credential_by_id( user: User = Depends(current_user), db_session: Session = Depends(get_session), ) -> StatusResponse: - delete_credential(credential_id, user, db_session, True) + delete_credential_for_user(credential_id, user, db_session, True) return StatusResponse( success=True, message="Credential deleted successfully", data=credential_id diff --git a/web/src/app/admin/connectors/[connector]/pages/gdrive/Credential.tsx b/web/src/app/admin/connectors/[connector]/pages/gdrive/Credential.tsx index 3c9f8e59f..40e2da9cb 100644 --- a/web/src/app/admin/connectors/[connector]/pages/gdrive/Credential.tsx +++ b/web/src/app/admin/connectors/[connector]/pages/gdrive/Credential.tsx @@ -1,6 +1,6 @@ import { Button } from "@/components/Button"; import { PopupSpec } from "@/components/admin/connectors/Popup"; -import { useState } from "react"; +import React, { useState, useEffect } from "react"; import { useSWRConfig } from "swr"; import * as Yup from "yup"; import { useRouter } from "next/navigation"; @@ -17,13 +17,18 @@ import { GoogleDriveCredentialJson, GoogleDriveServiceAccountCredentialJson, } from "@/lib/connectors/credentials"; +import { refreshAllGoogleData } from "@/lib/googleConnector"; +import { ValidSources } from "@/lib/types"; +import { buildSimilarCredentialInfoURL } from "@/app/admin/connector/[ccPairId]/lib"; type GoogleDriveCredentialJsonTypes = "authorized_user" | "service_account"; export const DriveJsonUpload = ({ setPopup, + onSuccess, }: { setPopup: (popupSpec: PopupSpec | null) => void; + onSuccess?: () => void; }) => { const { mutate } = useSWRConfig(); const [credentialJsonStr, setCredentialJsonStr] = useState< @@ -62,7 +67,6 @@ export const DriveJsonUpload = ({ + {isAdmin ? ( + <> +
+ If you want to update these credentials, delete the existing + credentials through the button below, and then upload a new + credentials JSON. +
+ + + ) : ( +
+ To change these credentials, please contact an administrator. +
+ )} ); } @@ -276,14 +327,14 @@ export const GmailJsonUploadSection = ({ > here {" "} - to either (1) setup a google OAuth App in your company workspace or (2) + 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. + Account key JSON if choosing option (2), and upload it here.

- + ); }; @@ -299,6 +350,34 @@ interface DriveCredentialSectionProps { user: User | null; } +async function handleRevokeAccess( + connectorExists: boolean, + setPopup: (popupSpec: PopupSpec | null) => void, + existingCredential: + | Credential + | Credential, + refreshCredentials: () => void +) { + if (connectorExists) { + const message = + "Cannot revoke the Gmail credential while any connector is still associated with the credential. " + + "Please delete all associated connectors, then try again."; + setPopup({ + message: message, + type: "error", + }); + return; + } + + await adminDeleteCredential(existingCredential.id); + setPopup({ + message: "Successfully revoked the Gmail credential!", + type: "success", + }); + + refreshCredentials(); +} + export const GmailAuthSection = ({ gmailPublicCredential, gmailServiceAccountCredential, @@ -310,31 +389,49 @@ export const GmailAuthSection = ({ user, }: DriveCredentialSectionProps) => { const router = useRouter(); + const [isAuthenticating, setIsAuthenticating] = useState(false); + const [localServiceAccountData, setLocalServiceAccountData] = useState( + serviceAccountKeyData + ); + const [localAppCredentialData, setLocalAppCredentialData] = + useState(appCredentialData); + const [localGmailPublicCredential, setLocalGmailPublicCredential] = useState( + gmailPublicCredential + ); + const [ + localGmailServiceAccountCredential, + setLocalGmailServiceAccountCredential, + ] = useState(gmailServiceAccountCredential); + + // Update local state when props change + useEffect(() => { + setLocalServiceAccountData(serviceAccountKeyData); + setLocalAppCredentialData(appCredentialData); + setLocalGmailPublicCredential(gmailPublicCredential); + setLocalGmailServiceAccountCredential(gmailServiceAccountCredential); + }, [ + serviceAccountKeyData, + appCredentialData, + gmailPublicCredential, + gmailServiceAccountCredential, + ]); const existingCredential = - gmailPublicCredential || gmailServiceAccountCredential; + localGmailPublicCredential || localGmailServiceAccountCredential; if (existingCredential) { return ( <>

- Existing credential already set up! + Uploaded and authenticated credential already exists!

- - - )} - - + } catch (error) { + setPopup({ + message: `Failed to create service account credential - ${error}`, + type: "error", + }); + } finally { + formikHelpers.setSubmitting(false); + } + }} + > + {({ isSubmitting }) => ( +
+ +
+ +
+ + )} + ); } - if (appCredentialData?.client_id) { + if (localAppCredentialData?.client_id) { return (

Next, you must provide credentials via OAuth. This gives us read - access to the docs you have access to in your gmail account. + access to the emails you have access to in your Gmail account.

); @@ -449,8 +556,8 @@ export const GmailAuthSection = ({ // case where no keys have been uploaded in step 1 return (

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

); }; diff --git a/web/src/app/admin/connectors/[connector]/pages/gmail/GmailPage.tsx b/web/src/app/admin/connectors/[connector]/pages/gmail/GmailPage.tsx index 28af99441..75b97bcba 100644 --- a/web/src/app/admin/connectors/[connector]/pages/gmail/GmailPage.tsx +++ b/web/src/app/admin/connectors/[connector]/pages/gmail/GmailPage.tsx @@ -1,10 +1,11 @@ "use client"; -import useSWR from "swr"; -import { errorHandlingFetcher } from "@/lib/fetcher"; +import React from "react"; +import { FetchError } from "@/lib/fetcher"; +import { ErrorCallout } from "@/components/ErrorCallout"; import { LoadingAnimation } from "@/components/Loading"; -import { usePopup } from "@/components/admin/connectors/Popup"; -import { CCPairBasicInfo } from "@/lib/types"; +import { PopupSpec, usePopup } from "@/components/admin/connectors/Popup"; +import { CCPairBasicInfo, ValidSources } from "@/lib/types"; import { Credential, GmailCredentialJson, @@ -14,26 +15,33 @@ import { GmailAuthSection, GmailJsonUploadSection } from "./Credential"; import { usePublicCredentials, useBasicConnectorStatus } from "@/lib/hooks"; import Title from "@/components/ui/title"; import { useUser } from "@/components/user/UserProvider"; +import { + useGoogleAppCredential, + useGoogleServiceAccountKey, + useGoogleCredentials, + useConnectorsByCredentialId, + checkCredentialsFetched, + filterUploadedCredentials, + checkConnectorsExist, + refreshAllGoogleData, +} from "@/lib/googleConnector"; export const GmailMain = () => { const { isAdmin, user } = useUser(); + const { popup, setPopup } = usePopup(); const { data: appCredentialData, isLoading: isAppCredentialLoading, error: isAppCredentialError, - } = useSWR<{ client_id: string }>( - "/api/manage/admin/connector/gmail/app-credential", - errorHandlingFetcher - ); + } = useGoogleAppCredential("gmail"); + const { data: serviceAccountKeyData, isLoading: isServiceAccountKeyLoading, error: isServiceAccountKeyError, - } = useSWR<{ service_account_email: string }>( - "/api/manage/admin/connector/gmail/service-account-key", - errorHandlingFetcher - ); + } = useGoogleServiceAccountKey("gmail"); + const { data: connectorIndexingStatuses, isLoading: isConnectorIndexingStatusesLoading, @@ -47,20 +55,45 @@ export const GmailMain = () => { refreshCredentials, } = usePublicCredentials(); - const { popup, setPopup } = usePopup(); + const { + data: gmailCredentials, + isLoading: isGmailCredentialsLoading, + error: gmailCredentialsError, + } = useGoogleCredentials(ValidSources.Gmail); - const appCredentialSuccessfullyFetched = - appCredentialData || - (isAppCredentialError && isAppCredentialError.status === 404); - const serviceAccountKeySuccessfullyFetched = - serviceAccountKeyData || - (isServiceAccountKeyError && isServiceAccountKeyError.status === 404); + const { credential_id, uploadedCredentials } = + filterUploadedCredentials(gmailCredentials); + + const { + data: gmailConnectors, + isLoading: isGmailConnectorsLoading, + error: gmailConnectorsError, + refreshConnectorsByCredentialId, + } = useConnectorsByCredentialId(credential_id); + + const { + appCredentialSuccessfullyFetched, + serviceAccountKeySuccessfullyFetched, + } = checkCredentialsFetched( + appCredentialData, + isAppCredentialError, + serviceAccountKeyData, + isServiceAccountKeyError + ); + + const handleRefresh = () => { + refreshCredentials(); + refreshConnectorsByCredentialId(); + refreshAllGoogleData(ValidSources.Gmail); + }; if ( (!appCredentialSuccessfullyFetched && isAppCredentialLoading) || (!serviceAccountKeySuccessfullyFetched && isServiceAccountKeyLoading) || (!connectorIndexingStatuses && isConnectorIndexingStatusesLoading) || - (!credentialsData && isCredentialsLoading) + (!credentialsData && isCredentialsLoading) || + (!gmailCredentials && isGmailCredentialsLoading) || + (!gmailConnectors && isGmailConnectorsLoading) ) { return (
@@ -70,19 +103,15 @@ export const GmailMain = () => { } if (credentialsError || !credentialsData) { - return ( -
-
Failed to load credentials.
-
- ); + return ; + } + + if (gmailCredentialsError || !gmailCredentials) { + return ; } if (connectorIndexingStatusesError || !connectorIndexingStatuses) { - return ( -
-
Failed to load connectors.
-
- ); + return ; } if ( @@ -90,21 +119,28 @@ export const GmailMain = () => { !serviceAccountKeySuccessfullyFetched ) { return ( -
-
- Error loading Gmail app credentials. Contact an administrator. -
-
+ ); } - const gmailPublicCredential: Credential | undefined = - credentialsData.find( - (credential) => - (credential.credential_json?.google_service_account_key || - credential.credential_json?.google_tokens) && - credential.admin_public + if (gmailConnectorsError) { + return ( + ); + } + + const connectorExistsFromCredential = checkConnectorsExist(gmailConnectors); + + const gmailPublicUploadedCredential: + | Credential + | undefined = credentialsData.find( + (credential) => + credential.credential_json?.google_tokens && + credential.admin_public && + credential.source === "gmail" && + credential.credential_json.authentication_method !== "oauth_interactive" + ); + const gmailServiceAccountCredential: | Credential | undefined = credentialsData.find( @@ -118,6 +154,13 @@ export const GmailMain = () => { (connectorIndexingStatus) => connectorIndexingStatus.source === "gmail" ); + const connectorExists = + connectorExistsFromCredential || gmailConnectorIndexingStatuses.length > 0; + + const hasUploadedCredentials = + Boolean(appCredentialData?.client_id) || + Boolean(serviceAccountKeyData?.service_account_email); + return ( <> {popup} @@ -129,21 +172,22 @@ export const GmailMain = () => { appCredentialData={appCredentialData} serviceAccountCredentialData={serviceAccountKeyData} isAdmin={isAdmin} + onSuccess={handleRefresh} /> - {isAdmin && ( + {isAdmin && hasUploadedCredentials && ( <> Step 2: Authenticate with Onyx 0} + connectorExists={connectorExists} user={user} /> diff --git a/web/src/components/Button.tsx b/web/src/components/Button.tsx index 5672469f1..ea6faadc3 100644 --- a/web/src/components/Button.tsx +++ b/web/src/components/Button.tsx @@ -3,7 +3,6 @@ interface Props { onClick?: React.MouseEventHandler; type?: "button" | "submit" | "reset"; disabled?: boolean; - fullWidth?: boolean; className?: string; } @@ -12,14 +11,12 @@ export const Button = ({ onClick, type = "submit", disabled = false, - fullWidth = false, className = "", }: Props) => { return (