Add select domain for move tool (#7057)

Signed-off-by: Denis Bykhov <[email protected]>
hcengineering · Oct 28, 2024 · 3a37a16 · 3a37a16
1 parent 7a8ee3c
commit 3a37a16
Show file tree

Hide file tree

Showing 2 changed files with 46 additions and 21 deletions.
diff --git 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<string>
 ): Promise<void> {
   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<string>
 ): Promise<void> {
   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
@@ -1666,21 +1666,39 @@ export function devTool (
     })
   })
 
-  program.command('move-workspace-to-pg <workspace> <region>').action(async (workspace: string, region: string) => {
-    const { dbUrl } = prepareTools()
-    const mongodbUri = getMongoDBUrl()
+  program
+    .command('move-workspace-to-pg <workspace> <region>')
+    .option('-i, --include <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()