oauth: GoogleDrive connector dual-flow + front gating (#6373)

* GoogleDrive connector: dual-flow oauth/nango * oauth: gate google_drive * fix * Update connectors/src/connectors/google_drive/temporal/utils.ts Co-authored-by: Flavien David <[email protected]> --------- Co-authored-by: Flavien David <[email protected]>
dust-tt · Jul 22, 2024 · 0adfd85 · 0adfd85
1 parent 0a5782f
commit 0adfd85
Show file tree

Hide file tree

Showing 3 changed files with 45 additions and 69 deletions.
diff --git a/connectors/src/connectors/google_drive/index.ts b/connectors/src/connectors/google_drive/index.ts
@@ -1,5 +1,4 @@
 import type { drive_v3 } from "googleapis";
-import { google } from "googleapis";
 import type { GaxiosResponse, OAuth2Client } from "googleapis-common";
 
 import {
@@ -43,7 +42,6 @@ import type { InferAttributes, WhereOptions } from "sequelize";
 import { Op } from "sequelize";
 import { v4 as uuidv4 } from "uuid";
 
-import { googleDriveConfig } from "@connectors/connectors/google_drive/lib/config";
 import { GOOGLE_DRIVE_SHARED_WITH_ME_VIRTUAL_ID } from "@connectors/connectors/google_drive/lib/consts";
 import { getGoogleDriveObject } from "@connectors/connectors/google_drive/lib/google_drive_api";
 import {
@@ -58,7 +56,6 @@ import {
   driveObjectToDustType,
   getAuthObject,
   getDriveClient,
-  getGoogleCredentials,
 } from "@connectors/connectors/google_drive/temporal/utils";
 import { concurrentExecutor } from "@connectors/lib/async_utils";
 import { syncSucceeded } from "@connectors/lib/sync_status";
@@ -210,53 +207,16 @@ export class GoogleDriveConnectorManager extends BaseConnectorManager<null> {
     return new Ok(connector.id.toString());
   }
 
-  async clean({
-    force,
-  }: {
-    force: boolean;
-  }): Promise<Result<undefined, Error>> {
+  async clean(): Promise<Result<undefined, Error>> {
     const connector = await ConnectorResource.fetchById(this.connectorId);
     if (!connector) {
       return new Err(
         new Error(`Could not find connector with id ${this.connectorId}`)
       );
     }
 
-    const authClient = new google.auth.OAuth2(
-      googleDriveConfig.getRequiredGoogleDriveClientId(),
-      googleDriveConfig.getRequiredGoogleDriveClientSecret()
-    );
-
-    try {
-      const credentials = await getGoogleCredentials(connector.connectionId);
-
-      const revokeTokenRes = await authClient.revokeToken(
-        credentials.credentials.refresh_token
-      );
-
-      if (revokeTokenRes.status !== 200) {
-        logger.error(
-          {
-            error: revokeTokenRes.data,
-          },
-          "Could not revoke token"
-        );
-        if (!force) {
-          return new Err(new Error("Could not revoke token"));
-        }
-      }
-    } catch (err) {
-      if (!force) {
-        throw err;
-      } else {
-        logger.error(
-          {
-            err,
-          },
-          "Error revoking token"
-        );
-      }
-    }
+    // Google revocation requires refresh tokens so would have to happen in `oauth`. But Google
+    // Drive does not rely on webhooks anymore so we can just delete the connector.
 
     const res = await connector.delete();
     if (res.isErr()) {

diff --git a/connectors/src/connectors/google_drive/temporal/utils.ts b/connectors/src/connectors/google_drive/temporal/utils.ts
@@ -1,11 +1,14 @@
-import { cacheWithRedis } from "@dust-tt/types";
+import { cacheWithRedis, getOAuthConnectionAccessToken } from "@dust-tt/types";
 import type { drive_v3 } from "googleapis";
 import { google } from "googleapis";
 import { OAuth2Client } from "googleapis-common";
 
 import { googleDriveConfig } from "@connectors/connectors/google_drive/lib/config";
+import { apiConfig } from "@connectors/lib/api/config";
 import type { NangoConnectionResponse } from "@connectors/lib/nango_helpers";
 import { getConnectionFromNango } from "@connectors/lib/nango_helpers";
+import { isDualUseOAuthConnectionId } from "@connectors/lib/oauth";
+import logger from "@connectors/logger/logger";
 import type { GoogleDriveObjectType } from "@connectors/types/google_drive";
 
 export function getDocumentId(driveFileId: string): string {
@@ -125,34 +128,47 @@ export async function driveObjectToDustType(
   }
 }
 
-export async function getGoogleCredentials(
-  nangoConnectionId: string
-): Promise<NangoConnectionResponse> {
-  return getConnectionFromNango({
-    connectionId: nangoConnectionId,
-    integrationId: googleDriveConfig.getRequiredNangoGoogleDriveConnectorId(),
-    refreshToken: false,
-  });
-}
-
 export async function getAuthObject(
-  nangoConnectionId: string
+  connectionId: string
 ): Promise<OAuth2Client> {
-  const res: NangoConnectionResponse = await getConnectionFromNango({
-    connectionId: nangoConnectionId,
-    integrationId: googleDriveConfig.getRequiredNangoGoogleDriveConnectorId(),
-    refreshToken: false,
-    useCache: true,
-  });
-
   const oauth2Client = new google.auth.OAuth2();
-  oauth2Client.setCredentials({
-    access_token: res.credentials.access_token,
-    scope: res.credentials.raw.scope,
-    token_type: res.credentials.raw.token_type,
-    expiry_date: new Date(res.credentials.expires_at).getTime(),
-  });
+  if (isDualUseOAuthConnectionId(connectionId)) {
+    const tokRes = await getOAuthConnectionAccessToken({
+      config: apiConfig.getOAuthAPIConfig(),
+      logger,
+      provider: "google_drive",
+      connectionId,
+    });
+    if (tokRes.isErr()) {
+      logger.error(
+        { connectionId, error: tokRes.error },
+        "Error retrieving Google access token"
+      );
+      throw new Error("Error retrieving Google access token");
+    }
 
+    oauth2Client.setCredentials({
+      access_token: tokRes.value.access_token,
+      scope: (tokRes.value.scrubbed_raw_json as { scope: string }).scope,
+      token_type: (tokRes.value.scrubbed_raw_json as { token_type: string })
+        .token_type,
+      expiry_date: tokRes.value.access_token_expiry,
+    });
+  } else {
+    const res: NangoConnectionResponse = await getConnectionFromNango({
+      connectionId: connectionId,
+      integrationId: googleDriveConfig.getRequiredNangoGoogleDriveConnectorId(),
+      refreshToken: false,
+      useCache: true,
+    });
+
+    oauth2Client.setCredentials({
+      access_token: res.credentials.access_token,
+      scope: res.credentials.raw.scope,
+      token_type: res.credentials.raw.token_type,
+      expiry_date: new Date(res.credentials.expires_at).getTime(),
+    });
+  }
   return oauth2Client;
 }
 

diff --git a/front/pages/w/[wId]/builder/data-sources/managed.tsx b/front/pages/w/[wId]/builder/data-sources/managed.tsx
@@ -111,7 +111,7 @@ export async function setupConnection({
     // `oauth`-ready providers
     (["github", "slack"].includes(provider) ||
       // Behind flag oauth-ready providers
-      (["intercom", "notion"].includes(provider) &&
+      (["intercom", "notion", "google_drive"].includes(provider) &&
         owner.flags.includes("test_oauth_setup")))
   ) {
     // OAuth flow