From 06472069384734cbbf42eb206c3d26f38b1b244d Mon Sep 17 00:00:00 2001
From: Denis Bykhov <bykhov.denis@gmail.com>
Date: Mon, 4 Nov 2024 17:33:44 +0500
Subject: [PATCH] Force move tool (#7095)

---
 dev/tool/src/db.ts    | 22 ++++++++++++++++++----
 dev/tool/src/index.ts |  5 ++++-
 2 files changed, 22 insertions(+), 5 deletions(-)

diff --git a/dev/tool/src/db.ts b/dev/tool/src/db.ts
index 0e7372c4dc..603d18eb93 100644
--- a/dev/tool/src/db.ts
+++ b/dev/tool/src/db.ts
@@ -68,7 +68,8 @@ async function moveWorkspace (
   pgClient: postgres.Sql,
   ws: Workspace,
   region: string,
-  include?: Set<string>
+  include?: Set<string>,
+  force = false
 ): Promise<void> {
   try {
     const wsId = getWorkspaceId(ws.workspace)
@@ -102,12 +103,24 @@ async function moveWorkspace (
         insertFields.push(field)
       }
       while (true) {
+        const toRemove: string[] = []
         while (docs.length < 50000) {
           const doc = (await cursor.next()) as Doc | null
           if (doc === null) break
-          if (currentIds.has(doc._id)) continue
+          if (currentIds.has(doc._id)) {
+            if (force) {
+              toRemove.push(doc._id)
+            } else {
+              continue
+            }
+          }
           docs.push(doc)
         }
+        if (toRemove.length > 0) {
+          await retryTxn(pgClient, async (client) => {
+            await client`DELETE FROM ${client(translateDomain(domain))} WHERE "workspaceId" = ${ws.workspace} AND _id IN (${client(toRemove)})`
+          })
+        }
         if (docs.length === 0) break
         while (docs.length > 0) {
           const part = docs.splice(0, 500)
@@ -138,7 +151,8 @@ export async function moveWorkspaceFromMongoToPG (
   dbUrl: string | undefined,
   ws: Workspace,
   region: string,
-  include?: Set<string>
+  include?: Set<string>,
+  force?: boolean
 ): Promise<void> {
   if (dbUrl === undefined) {
     throw new Error('dbUrl is required')
@@ -148,7 +162,7 @@ export async function moveWorkspaceFromMongoToPG (
   const pg = getDBClient(dbUrl)
   const pgClient = await pg.getClient()
 
-  await moveWorkspace(accountDb, mongo, pgClient, ws, region, include)
+  await moveWorkspace(accountDb, mongo, pgClient, ws, region, include, force)
   pg.close()
   client.close()
 }
diff --git a/dev/tool/src/index.ts b/dev/tool/src/index.ts
index 74472c3248..3c7d991517 100644
--- a/dev/tool/src/index.ts
+++ b/dev/tool/src/index.ts
@@ -1669,12 +1669,14 @@ export function devTool (
   program
     .command('move-workspace-to-pg <workspace> <region>')
     .option('-i, --include <include>', 'A list of ; separated domain names to include during backup', '*')
+    .option('-f|--force [force]', 'Force update', false)
     .action(
       async (
         workspace: string,
         region: string,
         cmd: {
           include: string
+          force: boolean
         }
       ) => {
         const { dbUrl } = prepareTools()
@@ -1694,7 +1696,8 @@ export function devTool (
             dbUrl,
             workspaceInfo,
             region,
-            cmd.include === '*' ? undefined : new Set(cmd.include.split(';').map((it) => it.trim()))
+            cmd.include === '*' ? undefined : new Set(cmd.include.split(';').map((it) => it.trim())),
+            cmd.force
           )
         })
       }