G Drive connector improvements

This commit is contained in:
hagen-danswer 2024-09-27 08:08:38 -07:00
parent fbf51b70d0
commit 21a9341b1b
3 changed files with 34 additions and 25 deletions

View File

@ -4,7 +4,6 @@ from uuid import UUID
from pydantic import BaseModel
from pydantic import Field
from pydantic import model_validator
from danswer.configs.app_configs import MASK_CREDENTIAL_PREFIX
from danswer.configs.constants import DocumentSource
@ -373,16 +372,6 @@ class GoogleServiceAccountCredentialRequest(BaseModel):
google_drive_delegated_user: str | None = None # email of user to impersonate
gmail_delegated_user: str | None = None # email of user to impersonate
@model_validator(mode="after")
def check_user_delegation(self) -> "GoogleServiceAccountCredentialRequest":
if (self.google_drive_delegated_user is None) == (
self.gmail_delegated_user is None
):
raise ValueError(
"Exactly one of google_drive_delegated_user or gmail_delegated_user must be set"
)
return self
class FileUploadResponse(BaseModel):
file_paths: list[str]

View File

@ -111,7 +111,7 @@ export const DriveJsonUpload = ({
}
if (credentialFileType === "service_account") {
const response = await fetch(
const initialResponse = await fetch(
"/api/manage/admin/connector/google-drive/service-account-key",
{
method: "PUT",
@ -121,17 +121,33 @@ export const DriveJsonUpload = ({
body: credentialJsonStr,
}
);
if (response.ok) {
const followUpResponse = await fetch(
"/api/manage/admin/connector/google-drive/service-account-credential",
{
method: "PUT",
headers: {
"Content-Type": "application/json",
},
body: credentialJsonStr,
}
);
if (initialResponse.ok && followUpResponse.ok) {
setPopup({
message: "Successfully uploaded app credentials",
type: "success",
});
} else {
const errorMsg = await response.text();
} else if (initialResponse.ok) {
const errorMsg = await followUpResponse.text();
setPopup({
message: `Failed to upload app credentials - ${errorMsg}`,
type: "error",
});
} else {
const errorMsg = await initialResponse.text();
setPopup({
message: `Failed to upload service account key - ${errorMsg}`,
type: "error",
});
}
mutate(
"/api/manage/admin/connector/google-drive/service-account-key"
@ -303,7 +319,7 @@ export const DriveJsonUploadSection = ({
interface DriveCredentialSectionProps {
googleDrivePublicCredential?: Credential<GoogleDriveCredentialJson>;
googleDriveServiceAccountCredential?: Credential<GoogleDriveServiceAccountCredentialJson>;
serviceAccountKeyData?: { service_account_email: string };
serviceAccountCredentialData?: { service_account_email: string };
appCredentialData?: { client_id: string };
setPopup: (popupSpec: PopupSpec | null) => void;
refreshCredentials: () => void;
@ -313,7 +329,7 @@ interface DriveCredentialSectionProps {
export const DriveOAuthSection = ({
googleDrivePublicCredential,
googleDriveServiceAccountCredential,
serviceAccountKeyData,
serviceAccountCredentialData,
appCredentialData,
setPopup,
refreshCredentials,
@ -321,9 +337,13 @@ export const DriveOAuthSection = ({
}: DriveCredentialSectionProps) => {
const router = useRouter();
const existingCredential =
const candidateCredential =
googleDrivePublicCredential || googleDriveServiceAccountCredential;
if (existingCredential) {
if (
candidateCredential &&
"google_drive_delegated_user" in candidateCredential.credential_json
) {
return (
<>
<p className="mb-2 text-sm">
@ -339,7 +359,7 @@ export const DriveOAuthSection = ({
});
return;
}
await adminDeleteCredential(existingCredential.id);
await adminDeleteCredential(candidateCredential.id);
setPopup({
message: "Successfully revoked access to Google Drive!",
type: "success",
@ -353,7 +373,7 @@ export const DriveOAuthSection = ({
);
}
if (serviceAccountKeyData?.service_account_email) {
if (serviceAccountCredentialData?.service_account_email) {
return (
<div>
<p className="text-sm mb-6">

View File

@ -33,7 +33,7 @@ const GDriveMain = ({}: {}) => {
);
const {
data: serviceAccountKeyData,
data: serviceAccountCredentialData,
isLoading: isServiceAccountKeyLoading,
error: isServiceAccountKeyError,
} = useSWR<{ service_account_email: string }, FetchError>(
@ -59,7 +59,7 @@ const GDriveMain = ({}: {}) => {
appCredentialData ||
(isAppCredentialError && isAppCredentialError.status === 404);
const serviceAccountKeySuccessfullyFetched =
serviceAccountKeyData ||
serviceAccountCredentialData ||
(isServiceAccountKeyError && isServiceAccountKeyError.status === 404);
if (isLoadingUser) {
@ -125,7 +125,7 @@ const GDriveMain = ({}: {}) => {
<DriveJsonUploadSection
setPopup={setPopup}
appCredentialData={appCredentialData}
serviceAccountCredentialData={serviceAccountKeyData}
serviceAccountCredentialData={serviceAccountCredentialData}
isAdmin={isAdmin}
/>
@ -142,7 +142,7 @@ const GDriveMain = ({}: {}) => {
googleDriveServiceAccountCredential
}
appCredentialData={appCredentialData}
serviceAccountKeyData={serviceAccountKeyData}
serviceAccountCredentialData={serviceAccountCredentialData}
connectorExists={googleDriveConnectorIndexingStatuses.length > 0}
/>
</>