Skip to content

Commit

Permalink
[parents_migration] Add and backfill Google Drive "Shared with me" fo…
Browse files Browse the repository at this point in the history
…lders (#9507)

* create and delete SharedWithMe folder

* add backfill script

* add concurrency on the script

* fix the internalId of the sharedWithMe folder

* use the dataSourceId in the SharedWithMe folder id instead of the connector ID

* move the deletion of the shared with me folder to the clean method

* use dataSourceId in the folderId for the sharedWithMe folder

* delete from core first

* remove the dataSourceId from the GOOGLE_DRIVE_SHARED_WITH_ME_VIRTUAL_ID folder

* cleanup

* remove the delete, leaving it for another task

* inline upsertSharedWithMeFolder

* move the upsertion to an activity
  • Loading branch information
aubin-tchoi authored Dec 19, 2024
1 parent ded1e01 commit 5176705
Show file tree
Hide file tree
Showing 5 changed files with 70 additions and 6 deletions.
37 changes: 37 additions & 0 deletions connectors/migrations/20241218_backfill_gdrive_shared_with_me.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,37 @@
import { makeScript } from "scripts/helpers";

import { GOOGLE_DRIVE_SHARED_WITH_ME_VIRTUAL_ID } from "@connectors/connectors/google_drive/lib/consts";
import { getInternalId } from "@connectors/connectors/google_drive/temporal/utils";
import { dataSourceConfigFromConnector } from "@connectors/lib/api/data_source_config";
import { concurrentExecutor } from "@connectors/lib/async_utils";
import { upsertDataSourceFolder } from "@connectors/lib/data_sources";
import { ConnectorResource } from "@connectors/resources/connector_resource";

makeScript({}, async ({ execute }, logger) => {
const connectors = await ConnectorResource.listByType("google_drive", {});

await concurrentExecutor(
connectors,
async (connector) => {
const folderId = getInternalId(GOOGLE_DRIVE_SHARED_WITH_ME_VIRTUAL_ID);
if (execute) {
await upsertDataSourceFolder({
dataSourceConfig: dataSourceConfigFromConnector(connector),
folderId,
parents: [folderId],
parentId: null,
title: "Shared with me",
mimeType: "application/vnd.dust.googledrive.folder",
});
logger.info(
`Upserted folder ${folderId} for connector ${connector.id}`
);
} else {
logger.info(
`Would upsert folder ${folderId} for connector ${connector.id}`
);
}
},
{ concurrency: 10 }
);
});
7 changes: 5 additions & 2 deletions connectors/src/connectors/google_drive/index.ts
Original file line number Diff line number Diff line change
Expand Up @@ -409,7 +409,7 @@ export class GoogleDriveConnectorManager extends BaseConnectorManager<null> {
// that are not living in a shared drive.
nodes.push({
provider: c.type,
internalId: GOOGLE_DRIVE_SHARED_WITH_ME_VIRTUAL_ID,
internalId: getInternalId(GOOGLE_DRIVE_SHARED_WITH_ME_VIRTUAL_ID),
parentInternalId: null,
type: "folder" as const,
preventSelection: true,
Expand All @@ -435,7 +435,10 @@ export class GoogleDriveConnectorManager extends BaseConnectorManager<null> {
// The "Shared with me" view requires to look for folders
// with the flag `sharedWithMe=true`, but there is no need to check for the parents.
let gdriveQuery = `mimeType='application/vnd.google-apps.folder'`;
if (parentInternalId === GOOGLE_DRIVE_SHARED_WITH_ME_VIRTUAL_ID) {
if (
parentInternalId ===
getInternalId(GOOGLE_DRIVE_SHARED_WITH_ME_VIRTUAL_ID)
) {
gdriveQuery += ` and sharedWithMe=true`;
} else {
gdriveQuery += ` and '${parentDriveId}' in parents`;
Expand Down
25 changes: 24 additions & 1 deletion connectors/src/connectors/google_drive/temporal/activities.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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";
Expand Down Expand Up @@ -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[]> {
Expand Down
3 changes: 3 additions & 0 deletions connectors/src/connectors/google_drive/temporal/workflows.ts
Original file line number Diff line number Diff line change
Expand Up @@ -25,6 +25,7 @@ const {
garbageCollectorFinished,
markFolderAsVisited,
shouldGarbageCollect,
upsertSharedWithMeFolder,
} = proxyActivities<typeof activities>({
startToCloseTimeout: "20 minutes",
});
Expand Down Expand Up @@ -104,6 +105,8 @@ export async function googleDriveFullSync({
}
});

await upsertSharedWithMeFolder(connectorId);

// Temp to clean up the running workflows state
foldersToBrowse = uniq(foldersToBrowse);

Expand Down
4 changes: 1 addition & 3 deletions connectors/src/resources/connector/google_drive.ts
Original file line number Diff line number Diff line change
Expand Up @@ -3,11 +3,9 @@ import type { Transaction } from "sequelize";

import {
GoogleDriveConfig,
GoogleDriveSheet,
} from "@connectors/lib/models/google_drive";
import {
GoogleDriveFiles,
GoogleDriveFolders,
GoogleDriveSheet,
GoogleDriveSyncToken,
} from "@connectors/lib/models/google_drive";
import type {
Expand Down

0 comments on commit 5176705

Please sign in to comment.