diff --git a/dev/tool/src/db.ts b/dev/tool/src/db.ts index a5c5885d3a..0e7372c4dc 100644 --- a/dev/tool/src/db.ts +++ b/dev/tool/src/db.ts @@ -67,24 +67,30 @@ async function moveWorkspace ( mongo: MongoClient, pgClient: postgres.Sql, ws: Workspace, - region: string + region: string, + include?: Set ): Promise { try { const wsId = getWorkspaceId(ws.workspace) const mongoDB = getWorkspaceMongoDB(mongo, wsId) const collections = await mongoDB.collections() - await createTable( - pgClient, - collections.map((c) => c.collectionName) - ) + let tables = collections.map((c) => c.collectionName) + if (include !== undefined) { + tables = tables.filter((t) => include.has(t)) + } + + await createTable(pgClient, tables) const token = generateToken(systemAccountEmail, wsId) const endpoint = await getTransactorEndpoint(token, 'external') const connection = (await connect(endpoint, wsId, undefined, { model: 'upgrade' })) as unknown as Client & BackupClient for (const collection of collections) { - const cursor = collection.find() const domain = translateDomain(collection.collectionName) + if (include !== undefined && !include.has(domain)) { + continue + } + const cursor = collection.find() const current = await pgClient`SELECT _id FROM ${pgClient(domain)} WHERE "workspaceId" = ${ws.workspace}` const currentIds = new Set(current.map((r) => r._id)) console.log('move domain', domain) @@ -131,7 +137,8 @@ export async function moveWorkspaceFromMongoToPG ( mongoUrl: string, dbUrl: string | undefined, ws: Workspace, - region: string + region: string, + include?: Set ): Promise { if (dbUrl === undefined) { throw new Error('dbUrl is required') @@ -141,7 +148,7 @@ export async function moveWorkspaceFromMongoToPG ( const pg = getDBClient(dbUrl) const pgClient = await pg.getClient() - await moveWorkspace(accountDb, mongo, pgClient, ws, region) + await moveWorkspace(accountDb, mongo, pgClient, ws, region, include) pg.close() client.close() } diff --git a/dev/tool/src/index.ts b/dev/tool/src/index.ts index ff43160457..5d2e7a55f2 100644 --- a/dev/tool/src/index.ts +++ b/dev/tool/src/index.ts @@ -1666,21 +1666,39 @@ export function devTool ( }) }) - program.command('move-workspace-to-pg ').action(async (workspace: string, region: string) => { - const { dbUrl } = prepareTools() - const mongodbUri = getMongoDBUrl() + program + .command('move-workspace-to-pg ') + .option('-i, --include ', 'A list of ; separated domain names to include during backup', '*') + .action( + async ( + workspace: string, + region: string, + cmd: { + include: string + } + ) => { + const { dbUrl } = prepareTools() + const mongodbUri = getMongoDBUrl() - await withDatabase(mongodbUri, async (db) => { - const workspaceInfo = await getWorkspaceById(db, workspace) - if (workspaceInfo === null) { - throw new Error(`workspace ${workspace} not found`) - } - if (workspaceInfo.region === region) { - throw new Error(`workspace ${workspace} is already migrated`) + await withDatabase(mongodbUri, async (db) => { + const workspaceInfo = await getWorkspaceById(db, workspace) + if (workspaceInfo === null) { + throw new Error(`workspace ${workspace} not found`) + } + if (workspaceInfo.region === region) { + throw new Error(`workspace ${workspace} is already migrated`) + } + await moveWorkspaceFromMongoToPG( + db, + mongodbUri, + dbUrl, + workspaceInfo, + region, + cmd.include === '*' ? undefined : new Set(cmd.include.split(';').map((it) => it.trim())) + ) + }) } - await moveWorkspaceFromMongoToPG(db, mongodbUri, dbUrl, workspaceInfo, region) - }) - }) + ) program.command('move-account-db-to-pg').action(async () => { const { dbUrl } = prepareTools()