More explicit credential creation flow (onyx-dot-app#2363)

* more explcit drive credential creation flow * remove logs * update naming * fix user-contributed formatting * fix (^) v2
UiPath · Oct 2, 2024 · 402ad27 · 402ad27
1 parent d4cd331
commit 402ad27
Show file tree

Hide file tree

Showing 4 changed files with 71 additions and 69 deletions.
diff --git a/backend/danswer/connectors/notion/connector.py b/backend/danswer/connectors/notion/connector.py
@@ -241,7 +241,8 @@ def _read_blocks(
                     logger.warning(
                         f"Skipping 'external_object_instance_page' ('{result_block_id}') for base block '{base_block_id}': "
                         f"Notion API does not currently support reading external blocks (as of 24/07/03) "
-                        f"(discussion: https://github.com/danswer-ai/danswer/issues/1761)")
+                        f"(discussion: https://github.com/danswer-ai/danswer/issues/1761)"
+                    )
                     continue
 
                 cur_result_text_arr = []

diff --git a/backend/danswer/server/documents/models.py b/backend/danswer/server/documents/models.py
@@ -4,6 +4,7 @@
 
 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
@@ -346,8 +347,18 @@ class GoogleServiceAccountKey(BaseModel):
 
 
 class GoogleServiceAccountCredentialRequest(BaseModel):
-    google_drive_delegated_user: str | None  # email of user to impersonate
-    gmail_delegated_user: str | None  # email of user to impersonate
+    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):

diff --git a/web/src/app/admin/connectors/[connector]/pages/gdrive/Credential.tsx b/web/src/app/admin/connectors/[connector]/pages/gdrive/Credential.tsx
@@ -17,6 +17,8 @@ import {
   GoogleDriveServiceAccountCredentialJson,
 } from "@/lib/connectors/credentials";
 
+import { Button as TremorButton } from "@tremor/react";
+
 type GoogleDriveCredentialJsonTypes = "authorized_user" | "service_account";
 
 export const DriveJsonUpload = ({
@@ -344,7 +346,7 @@ export const DriveOAuthSection = ({
   if (serviceAccountKeyData?.service_account_email) {
     return (
       <div>
-        <p className="text-sm mb-2">
+        <p className="text-sm mb-6">
           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.
@@ -356,70 +358,59 @@ export const DriveOAuthSection = ({
           the documents you want to index with the service account.
         </p>
 
-        <Card>
-          <Formik
-            initialValues={{
-              google_drive_delegated_user: "",
-            }}
-            validationSchema={Yup.object().shape({
-              google_drive_delegated_user: Yup.string().optional(),
-            })}
-            onSubmit={async (values, formikHelpers) => {
-              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",
-                });
+        <Formik
+          initialValues={{
+            google_drive_delegated_user: "",
+          }}
+          validationSchema={Yup.object().shape({
+            google_drive_delegated_user: Yup.string().optional(),
+          })}
+          onSubmit={async (values, formikHelpers) => {
+            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,
+                }),
               }
-              refreshCredentials();
-            }}
-          >
-            {({ isSubmitting }) => (
-              <Form>
-                <TextFormField
-                  name="google_drive_delegated_user"
-                  label="[Optional] User email to impersonate:"
-                  subtext="If left blank, Danswer will use the service account itself."
-                />
-                <div className="flex">
-                  <button
-                    type="submit"
-                    disabled={isSubmitting}
-                    className={
-                      "bg-slate-500 hover:bg-slate-700 text-white " +
-                      "font-bold py-2 px-4 rounded focus:outline-none " +
-                      "focus:shadow-outline w-full max-w-sm mx-auto"
-                    }
-                  >
-                    Submit
-                  </button>
-                </div>
-              </Form>
-            )}
-          </Formik>
-        </Card>
+            );
+
+            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",
+              });
+            }
+            refreshCredentials();
+          }}
+        >
+          {({ isSubmitting }) => (
+            <Form>
+              <TextFormField
+                name="google_drive_delegated_user"
+                label="[Optional] User email to impersonate:"
+                subtext="If left blank, Danswer will use the service account itself."
+              />
+              <div className="flex">
+                <TremorButton type="submit" disabled={isSubmitting}>
+                  Create Credential
+                </TremorButton>
+              </div>
+            </Form>
+          )}
+        </Formik>
       </div>
     );
   }

diff --git a/web/src/app/admin/connectors/[connector]/pages/gdrive/GoogleDrivePage.tsx b/web/src/app/admin/connectors/[connector]/pages/gdrive/GoogleDrivePage.tsx
@@ -8,8 +8,6 @@ import { ErrorCallout } from "@/components/ErrorCallout";
 import { LoadingAnimation } from "@/components/Loading";
 import { usePopup } from "@/components/admin/connectors/Popup";
 import { ConnectorIndexingStatus } from "@/lib/types";
-import { getCurrentUser } from "@/lib/user";
-import { User, UserRole } from "@/lib/types";
 import { usePublicCredentials } from "@/lib/hooks";
 import { Title } from "@tremor/react";
 import { DriveJsonUploadSection, DriveOAuthSection } from "./Credential";
@@ -109,6 +107,7 @@ const GDriveMain = ({}: {}) => {
     | undefined = credentialsData.find(
     (credential) => credential.credential_json?.google_drive_service_account_key
   );
+
   const googleDriveConnectorIndexingStatuses: ConnectorIndexingStatus<
     GoogleDriveConfig,
     GoogleDriveCredentialJson