Allow terminating idle migration blockers

backend/scripts/apply-tenant-export-lock-migration.ts

@@ -12,31 +12,40 @@ const logStep = (message: string) => {
     console.log(`[tenant-lock-migration] ${message}`)
 }
 
-const printRelationLocks = async (db: ReturnType<typeof drizzle>, relationName: string) => {
-    const result = await db.execute(sql`
-        SELECT
-            activity.pid,
-            activity.state,
-            activity.wait_event_type,
-            activity.wait_event,
-            now() - activity.query_start AS query_age,
-            locks.mode,
-            locks.granted,
-            left(activity.query, 500) AS query
-        FROM pg_locks locks
-        JOIN pg_class relation ON relation.oid = locks.relation
-        LEFT JOIN pg_stat_activity activity ON activity.pid = locks.pid
-        WHERE relation.relname = ${relationName}
-        ORDER BY locks.granted ASC, activity.query_start ASC NULLS LAST;
-    `)
+const terminateIdleBlockers = process.argv.includes("--terminate-idle-blockers")
 
-    if (!result.rows.length) {
+const relationLocksQuery = (relationName: string) => sql`
+    SELECT
+        activity.pid,
+        activity.state,
+        activity.wait_event_type,
+        activity.wait_event,
+        now() - activity.query_start AS query_age,
+        locks.mode,
+        locks.granted,
+        left(activity.query, 500) AS query
+    FROM pg_locks locks
+    JOIN pg_class relation ON relation.oid = locks.relation
+    LEFT JOIN pg_stat_activity activity ON activity.pid = locks.pid
+    WHERE relation.relname = ${relationName}
+    ORDER BY locks.granted ASC, activity.query_start ASC NULLS LAST;
+`
+
+const getRelationLocks = async (db: ReturnType<typeof drizzle>, relationName: string) => {
+    const result = await db.execute(relationLocksQuery(relationName))
+    return result.rows as Record<string, any>[]
+}
+
+const printRelationLocks = async (db: ReturnType<typeof drizzle>, relationName: string) => {
+    const rows = await getRelationLocks(db, relationName)
+
+    if (!rows.length) {
         logStep(`Keine aktiven Locks fuer ${relationName} gefunden`)
         return
     }
 
     logStep(`Aktive Locks fuer ${relationName}:`)
-    for (const row of result.rows as Record<string, any>[]) {
+    for (const row of rows) {
         console.log(JSON.stringify({
             pid: row.pid,
             state: row.state,
@@ -50,6 +59,54 @@ const printRelationLocks = async (db: ReturnType<typeof drizzle>, relationName:
     }
 }
 
+const terminateIdleRelationBlockers = async (db: ReturnType<typeof drizzle>, relationName: string) => {
+    const rows = await getRelationLocks(db, relationName)
+    const blockerPids = rows
+        .filter((row) =>
+            row.granted === true &&
+            row.state === "idle in transaction" &&
+            Number(row.pid) !== process.pid
+        )
+        .map((row) => Number(row.pid))
+        .filter((pid) => Number.isFinite(pid))
+
+    if (!blockerPids.length) {
+        logStep(`Keine idle-in-transaction Blocker fuer ${relationName} gefunden`)
+        return false
+    }
+
+    for (const pid of blockerPids) {
+        logStep(`Beende idle-in-transaction DB-Session ${pid}`)
+        await db.execute(sql`SELECT pg_terminate_backend(${pid})`)
+    }
+
+    return true
+}
+
+const executeStatement = async (
+    db: ReturnType<typeof drizzle>,
+    statement: ReturnType<typeof sql>,
+) => {
+    await db.execute(statement)
+}
+
+const handleLockError = async (
+    db: ReturnType<typeof drizzle>,
+    relationName: string,
+) => {
+    logStep(`Konnte Lock fuer ${relationName} nicht rechtzeitig bekommen`)
+    await printRelationLocks(db, relationName)
+
+    if (!terminateIdleBlockers) return false
+
+    const terminated = await terminateIdleRelationBlockers(db, relationName)
+    if (terminated) {
+        logStep(`Retry nach beendeten idle-in-transaction Sessions fuer ${relationName}`)
+    }
+
+    return terminated
+}
+
 const executeStep = async (
     db: ReturnType<typeof drizzle>,
     message: string,
@@ -59,19 +116,25 @@ const executeStep = async (
     logStep(message)
 
     try {
-        await db.execute(statement)
+        await executeStatement(db, statement)
     } catch (err: any) {
         const code = err?.cause?.code || err?.code
-        if (lockRelationName && ["55P03", "57014"].includes(code)) {
-            logStep(`Konnte Lock fuer ${lockRelationName} nicht rechtzeitig bekommen`)
-            await printRelationLocks(db, lockRelationName)
+        if (!lockRelationName || !["55P03", "57014"].includes(code)) {
+            throw err
         }
 
-        throw err
+        const shouldRetry = await handleLockError(db, lockRelationName)
+        if (!shouldRetry) throw err
+
+        await executeStatement(db, statement)
     }
 }
 
 async function run() {
+    if (terminateIdleBlockers) {
+        logStep("Option --terminate-idle-blockers aktiv")
+    }
+
     let connectionString = process.env.DATABASE_URL
 
     if (!connectionString) {