[Util] script to check google drive file content (#2904)

* [Util] script to check google drive file content * migration to CLI * clean * clean 2 * clean 3
dust-tt · Dec 18, 2023 · 9f112e7 · 9f112e7
1 parent 4f482de
commit 9f112e7
Show file tree

Hide file tree

Showing 2 changed files with 46 additions and 4 deletions.
diff --git a/connectors/src/admin/cli.ts b/connectors/src/admin/cli.ts
@@ -9,7 +9,12 @@ import {
   STOP_CONNECTOR_BY_TYPE,
   SYNC_CONNECTOR_BY_TYPE,
 } from "@connectors/connectors";
-import { getDocumentId } from "@connectors/connectors/google_drive/temporal/activities";
+import {
+  getAuthObject,
+  getDocumentId,
+  getDriveClient,
+  MIME_TYPES_TO_EXPORT,
+} from "@connectors/connectors/google_drive/temporal/activities";
 import {
   launchGoogleDriveIncrementalSyncWorkflow,
   launchGoogleDriveRenewWebhooksWorkflow,
@@ -277,6 +282,44 @@ const google = async (command: string, args: parseArgs.ParsedArgs) => {
       }
       return;
     }
+    case "check-file": {
+      if (!args.connectorId) {
+        throw new Error("Missing --connectorId argument");
+      }
+      if (!args.fileId) {
+        throw new Error("Missing --fileId argument");
+      }
+      if (
+        !args.fileType ||
+        (args.fileType !== "document" && args.fileType !== "presentation")
+      ) {
+        throw new Error(
+          `Invalid or missing --fileType argument: ${args.fileType}`
+        );
+      }
+      console.log(`Checking gdrive file`);
+      const connector = await Connector.findByPk(args.connectorId);
+      if (!connector) {
+        throw new Error(`Connector ${args.connectorId} not found`);
+      }
+      const drive = await getDriveClient(
+        await getAuthObject(connector.connectionId)
+      );
+      const res = await drive.files.export({
+        fileId: args.fileId,
+        mimeType:
+          MIME_TYPES_TO_EXPORT[
+            args.fileType === "document"
+              ? "application/vnd.google-apps.document"
+              : "application/vnd.google-apps.presentation"
+          ],
+      });
+      console.log(`Status: ${res.status}`);
+      console.log(`Type: ${typeof res.data}`);
+      console.log(`Content:`);
+      console.log(res.data);
+      return;
+    }
     case "restart-google-webhooks": {
       await throwOnError(launchGoogleDriveRenewWebhooksWorkflow());
       return;

diff --git a/connectors/src/connectors/google_drive/temporal/activities.ts b/connectors/src/connectors/google_drive/temporal/activities.ts
@@ -22,6 +22,7 @@ import { drive_v3 } from "googleapis";
 import { OAuth2Client } from "googleapis-common";
 import { CreationAttributes, literal, Op } from "sequelize";
 
+import { registerWebhook } from "@connectors/connectors/google_drive/lib";
 import { dataSourceConfigFromConnector } from "@connectors/lib/api/data_source_config";
 import { dpdf2text } from "@connectors/lib/dpdf2text";
 import { ExternalOauthTokenError } from "@connectors/lib/error";
@@ -36,12 +37,10 @@ import {
 import { getConnectionFromNango } from "@connectors/lib/nango_helpers";
 import logger from "@connectors/logger/logger";
 
-import { registerWebhook } from "../lib";
-
 const FILES_SYNC_CONCURRENCY = 10;
 const FILES_GC_CONCURRENCY = 5;
 
-const MIME_TYPES_TO_EXPORT: { [key: string]: string } = {
+export const MIME_TYPES_TO_EXPORT: { [key: string]: string } = {
   "application/vnd.google-apps.document": "text/plain",
   "application/vnd.google-apps.presentation": "text/plain",
 };