move the upsertion to an activity

dust-tt · Dec 19, 2024 · ca11df9 · ca11df9
1 parent a41b0a9
commit ca11df9
Show file tree

Hide file tree

Showing 3 changed files with 27 additions and 14 deletions.
diff --git a/connectors/src/connectors/google_drive/index.ts b/connectors/src/connectors/google_drive/index.ts
@@ -48,9 +48,7 @@ import {
   BaseConnectorManager,
   ConnectorManagerError,
 } from "@connectors/connectors/interface";
-import { dataSourceConfigFromConnector } from "@connectors/lib/api/data_source_config";
 import { concurrentExecutor } from "@connectors/lib/async_utils";
-import { upsertDataSourceFolder } from "@connectors/lib/data_sources";
 import {
   GoogleDriveConfig,
   GoogleDriveFiles,
@@ -131,17 +129,6 @@ export class GoogleDriveConnectorManager extends BaseConnectorManager<null> {
       googleDriveConfigurationBlob
     );
 
-    // Upserts to data_sources_folders (core) a top-level folder "Shared with me".
-    const folderId = getInternalId(GOOGLE_DRIVE_SHARED_WITH_ME_VIRTUAL_ID);
-    await upsertDataSourceFolder({
-      dataSourceConfig: dataSourceConfigFromConnector(connector),
-      folderId,
-      parents: [folderId],
-      parentId: null,
-      title: "Shared with me",
-      mimeType: "application/vnd.dust.googledrive.folder",
-    });
-
     // We mark it artificially as sync succeeded as google drive is created empty.
     await syncSucceeded(connector.id);
 

diff --git a/connectors/src/connectors/google_drive/temporal/activities.ts b/connectors/src/connectors/google_drive/temporal/activities.ts
@@ -7,7 +7,10 @@ import StatsD from "hot-shots";
 import PQueue from "p-queue";
 import { Op } from "sequelize";
 
-import { GOOGLE_DRIVE_USER_SPACE_VIRTUAL_DRIVE_ID } from "@connectors/connectors/google_drive/lib/consts";
+import {
+  GOOGLE_DRIVE_SHARED_WITH_ME_VIRTUAL_ID,
+  GOOGLE_DRIVE_USER_SPACE_VIRTUAL_DRIVE_ID,
+} from "@connectors/connectors/google_drive/lib/consts";
 import { getGoogleDriveObject } from "@connectors/connectors/google_drive/lib/google_drive_api";
 import { getFileParentsMemoized } from "@connectors/connectors/google_drive/lib/hierarchy";
 import { syncOneFile } from "@connectors/connectors/google_drive/temporal/file";
@@ -54,6 +57,26 @@ type LightGoogleDrive = {
 
 export const statsDClient = new StatsD();
 
+/**
+ *  Upserts to data_sources_folders (core) a top-level folder "Shared with me".
+ */
+export async function upsertSharedWithMeFolder(connectorId: ModelId) {
+  const connector = await ConnectorResource.fetchById(connectorId);
+  if (!connector) {
+    throw new Error(`Connector ${connectorId} not found`);
+  }
+
+  const folderId = getInternalId(GOOGLE_DRIVE_SHARED_WITH_ME_VIRTUAL_ID);
+  await upsertDataSourceFolder({
+    dataSourceConfig: dataSourceConfigFromConnector(connector),
+    folderId,
+    parents: [folderId],
+    parentId: null,
+    title: "Shared with me",
+    mimeType: "application/vnd.dust.googledrive.folder",
+  });
+}
+
 export async function getDrives(
   connectorId: ModelId
 ): Promise<LightGoogleDrive[]> {

diff --git a/connectors/src/connectors/google_drive/temporal/workflows.ts b/connectors/src/connectors/google_drive/temporal/workflows.ts
@@ -25,6 +25,7 @@ const {
   garbageCollectorFinished,
   markFolderAsVisited,
   shouldGarbageCollect,
+  upsertSharedWithMeFolder,
 } = proxyActivities<typeof activities>({
   startToCloseTimeout: "20 minutes",
 });
@@ -104,6 +105,8 @@ export async function googleDriveFullSync({
     }
   });
 
+  await upsertSharedWithMeFolder(connectorId);
+
   // Temp to clean up the running workflows state
   foldersToBrowse = uniq(foldersToBrowse);