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

import "dotenv/config"
import { drizzle } from "drizzle-orm/node-postgres"
import { sql } from "drizzle-orm"
import { Pool } from "pg"

import { loadSecrets, secrets } from "../src/utils/secrets"

const maskConnectionString = (value: string) =>
    value.replace(/:\/\/([^:@]+):([^@]+)@/, "://$1:***@")

const logStep = (message: string) => {
    console.log(`[tenant-lock-migration] ${message}`)
}

const terminateIdleBlockers = process.argv.includes("--terminate-idle-blockers")

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 rows) {
        console.log(JSON.stringify({
            pid: row.pid,
            state: row.state,
            waitEventType: row.wait_event_type,
            waitEvent: row.wait_event,
            queryAge: row.query_age,
            mode: row.mode,
            granted: row.granted,
            query: row.query,
        }, null, 2))
    }
}

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,
    statement: ReturnType<typeof sql>,
    lockRelationName?: string,
) => {
    logStep(message)

    try {
        await executeStatement(db, statement)
    } catch (err: any) {
        const code = err?.cause?.code || err?.code
        if (!lockRelationName || !["55P03", "57014"].includes(code)) {
            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) {
        logStep("DATABASE_URL nicht in process.env gefunden, lade Secrets")
        await loadSecrets()
        connectionString = secrets.DATABASE_URL
    } else {
        logStep("DATABASE_URL aus process.env gefunden")
    }

    if (!connectionString) {
        throw new Error("DATABASE_URL not configured")
    }

    logStep(`Verbinde mit ${maskConnectionString(connectionString)}`)

    const pool = new Pool({
        connectionString,
        max: 1,
        connectionTimeoutMillis: 10000,
        idleTimeoutMillis: 10000,
    })

    try {
        const db = drizzle(pool)

        logStep("Teste DB-Verbindung")
        await db.execute(sql`SELECT 1`)

        logStep("Setze Lock- und Statement-Timeout")
        await db.execute(sql`SET lock_timeout = '10s'`)
        await db.execute(sql`SET statement_timeout = '2min'`)

        await executeStep(db, "Erweitere tenant_export_jobs", sql`
            ALTER TABLE "tenant_export_jobs"
                ADD COLUMN IF NOT EXISTS "operation" text DEFAULT 'export' NOT NULL,
                ADD COLUMN IF NOT EXISTS "previous_tenant_locked" "locked_tenant";
        `, "tenant_export_jobs")

        await executeStep(db, "Erweitere tenants", sql`
            ALTER TABLE "tenants"
                ADD COLUMN IF NOT EXISTS "locked_by_export_job_id" uuid;
        `, "tenants")

        await executeStep(db, "Pruefe Foreign Key", sql`
            DO $$
            BEGIN
                IF NOT EXISTS (
                    SELECT 1
                    FROM information_schema.table_constraints
                    WHERE constraint_schema = 'public'
                      AND constraint_name = 'tenants_locked_by_export_job_id_tenant_export_jobs_id_fk'
                ) THEN
                    ALTER TABLE "tenants"
                        ADD CONSTRAINT "tenants_locked_by_export_job_id_tenant_export_jobs_id_fk"
                        FOREIGN KEY ("locked_by_export_job_id")
                        REFERENCES "tenant_export_jobs"("id")
                        ON DELETE SET NULL;
                END IF;
            END $$;
        `, "tenants")

        await executeStep(db, "Pruefe Drizzle-Schema", sql`
            CREATE SCHEMA IF NOT EXISTS drizzle;
        `)

        await executeStep(db, "Pruefe Drizzle-Migrationstabelle", sql`
            CREATE TABLE IF NOT EXISTS drizzle.__drizzle_migrations (
                id SERIAL PRIMARY KEY,
                hash text NOT NULL,
                created_at bigint
            );
        `)

        await executeStep(db, "Setze Drizzle-Journal auf 0055", sql`
            INSERT INTO drizzle.__drizzle_migrations (hash, created_at)
            SELECT 'manual-0055_tenant_export_maintenance_lock', 1784021669805
            WHERE NOT EXISTS (
                SELECT 1
                FROM drizzle.__drizzle_migrations
                WHERE created_at >= 1784021669805
            );
        `)

        logStep("Fertig")
        console.log("✅ Tenant export lock migration applied")
    } finally {
        logStep("Schließe DB-Verbindung")
        await pool.end()
    }
}

run().catch((err) => {
    console.error("❌ Tenant export lock migration failed")
    console.error(err)
    process.exit(1)
})