fix: Lock the table before dropping mirroring triggers so disable cannot deadlock on a partitioned parent

CHANGELOG.md

@@ -1,5 +1,7 @@
 ## Unreleased
 
+- Fixed `disable_mirroring` and `disable_retired_mirroring` deadlocking against live traffic on a partitioned table. `DROP TRIGGER` on a partitioned parent locks every partition before the parent while statements lock the parent first; the source table is now locked explicitly before the trigger is dropped.
+- Added `--lock-timeout` (default `5s`) to the enable and disable mirroring commands, matching `swap`.
 - Added support for managing existing (already-partitioned) tables: bounds-anchored extension that recognizes legacy-named partitions and extends contiguously without renaming, gaps, or overlaps
 - Added a weekly (ISO-week, Monday-aligned) period
 - Added composite / parent-owned primary key support to `add_partitions` (skips the redundant per-partition key when the parent owns one)

README.md

@@ -144,6 +144,8 @@ pgslice disable_mirroring <table>
 
 After the swap, the original mirroring triggers are no longer needed since the tables have been swapped.
 
+The table is locked before the trigger is dropped, parent first and then its partitions, which is the order application statements lock in, so the drop cannot deadlock against live traffic. While the command waits for the lock, new statements on the table queue behind it; the wait is bounded by `--lock-timeout` (default: `5s`), and a timeout aborts the command cleanly so it can be retried.
+
 11. Enable Reverse Mirroring (now-partitioned table to retired table)
 
 This will make unswapping later less problematic as the two tables are kept in sync. Note that
@@ -168,7 +170,7 @@ pgslice fill <table> --swapped
 pgslice disable_retired_mirroring <table>
 ```
 
-Once you're confident the retired table is no longer needed and you're ready to drop it, disable the retired mirroring triggers.
+Once you're confident the retired table is no longer needed and you're ready to drop it, disable the retired mirroring triggers. The table is locked before the trigger is dropped, parent first and then its partitions, which is the order application statements lock in, so the drop cannot deadlock against live traffic. While the command waits for the lock, new statements on the table queue behind it; the wait is bounded by `--lock-timeout` (default: `5s`), and a timeout aborts the command cleanly so it can be retried.
 
 14. Back up the retired table with a tool like [pg_dump](https://www.postgresql.org/docs/current/static/app-pgdump.html) and drop it
 

src/commands/disable-mirroring.ts

@@ -11,21 +11,37 @@ export class DisableMirroringCommand extends BaseCommand {
     details: `
       This command removes the triggers on the source table that mirror operations
       to the intermediate table. Use this after partitioning is complete.
+
+      The table is locked before the trigger is dropped, parent first and
+      then its partitions, the same order application statements lock in,
+      so the drop cannot deadlock against live traffic. Statements on the
+      table queue behind the lock for the duration; the wait is bounded by
+      --lock-timeout, and a timeout aborts cleanly and can be retried.
     `,
     examples: [
       ["Disable mirroring for a table", "$0 mirroring disable posts"],
       [
         "Disable mirroring with explicit schema",
         "$0 mirroring disable myschema.posts",
       ],
+      [
+        "With custom lock timeout",
+        "$0 mirroring disable posts --lock-timeout 10s",
+      ],
     ],
   });
 
   table = Option.String({ required: true, name: "table" });
+  lockTimeout = Option.String("--lock-timeout", "5s", {
+    description: "Bound on how long to wait for the table lock",
+  });
 
   override async perform(pgslice: Pgslice): Promise<void> {
     await pgslice.start(async (tx) => {
-      await pgslice.disableMirroring(tx, { table: this.table });
+      await pgslice.disableMirroring(tx, {
+        table: this.table,
+        lockTimeout: this.lockTimeout,
+      });
       this.context.stdout.write(
         `Mirroring triggers disabled for ${this.table}\n`,
       );

src/commands/disable-retired-mirroring.ts

@@ -15,6 +15,12 @@ export class DisableRetiredMirroringCommand extends BaseCommand {
       This command removes the triggers on the source table that mirror operations
       to the retired table. Use this when you no longer need to keep the retired
       table in sync with the main table.
+
+      The table is locked before the trigger is dropped, parent first and
+      then its partitions, the same order application statements lock in,
+      so the drop cannot deadlock against live traffic. Statements on the
+      table queue behind the lock for the duration; the wait is bounded by
+      --lock-timeout, and a timeout aborts cleanly and can be retried.
     `,
     examples: [
       [
@@ -25,16 +31,24 @@ export class DisableRetiredMirroringCommand extends BaseCommand {
         "Disable retired mirroring with explicit schema",
         "$0 mirroring disable-retired myschema.posts",
       ],
+      [
+        "With custom lock timeout",
+        "$0 mirroring disable-retired posts --lock-timeout 10s",
+      ],
     ],
   });
 
   table = Option.String({ required: true, name: "table" });
+  lockTimeout = Option.String("--lock-timeout", "5s", {
+    description: "Bound on how long to wait for the table lock",
+  });
 
   override async perform(pgslice: Pgslice): Promise<void> {
     await pgslice.start(async (tx) => {
       await pgslice.disableMirroring(tx, {
         table: this.table,
         targetType: "retired",
+        lockTimeout: this.lockTimeout,
       });
       this.context.stdout.write(
         `Retired mirroring triggers disabled for ${this.table}\n`,

src/commands/enable-mirroring.ts

@@ -24,10 +24,16 @@ export class EnableMirroringCommand extends BaseCommand {
   });
 
   table = Option.String({ required: true, name: "table" });
+  lockTimeout = Option.String("--lock-timeout", "5s", {
+    description: "Bound on how long to wait for the table lock",
+  });
 
   override async perform(pgslice: Pgslice): Promise<void> {
     await pgslice.start(async (tx) =>
-      pgslice.enableMirroring(tx, { table: this.table }),
+      pgslice.enableMirroring(tx, {
+        table: this.table,
+        lockTimeout: this.lockTimeout,
+      }),
     );
 
     this.context.stdout.write(`Mirroring triggers enabled for ${this.table}\n`);

src/commands/enable-retired-mirroring.ts

@@ -34,10 +34,17 @@ export class EnableRetiredMirroringCommand extends BaseCommand {
   });
 
   table = Option.String({ required: true, name: "table" });
+  lockTimeout = Option.String("--lock-timeout", "5s", {
+    description: "Bound on how long to wait for the table lock",
+  });
 
   override async perform(pgslice: Pgslice): Promise<void> {
     await pgslice.start(async (tx) =>
-      pgslice.enableMirroring(tx, { table: this.table, targetType: "retired" }),
+      pgslice.enableMirroring(tx, {
+        table: this.table,
+        targetType: "retired",
+        lockTimeout: this.lockTimeout,
+      }),
     );
 
     this.context.stdout.write(

src/mirroring.test.ts

@@ -1,8 +1,10 @@
 import { describe, expect } from "vitest";
-import { sql } from "slonik";
+import { DatabaseTransactionConnection, sql } from "slonik";
 import { z } from "zod";
 
 import { pgsliceTest as test } from "./testing/index.js";
+import { Mirroring } from "./mirroring.js";
+import { Table } from "./table.js";
 
 describe("Pgslice.enableMirroring", () => {
   test.beforeEach(async ({ transaction }) => {
@@ -358,3 +360,169 @@ describe("Pgslice.disableMirroring", () => {
     });
   });
 });
+
+describe("Pgslice.disableMirroring (partitioned table)", () => {
+  const triggerCount = async (
+    transaction: DatabaseTransactionConnection,
+    name: string,
+  ) =>
+    transaction.oneFirst(
+      sql.type(z.object({ count: z.coerce.number() }))`
+        SELECT COUNT(*)::int FROM pg_trigger WHERE tgname = ${name}
+      `,
+    );
+
+  test.beforeEach(async ({ pgslice, transaction }) => {
+    await transaction.query(sql.unsafe`
+      CREATE TABLE posts (
+        id SERIAL,
+        title TEXT NOT NULL,
+        created_at DATE NOT NULL
+      ) PARTITION BY RANGE (created_at)
+    `);
+    await transaction.query(sql.unsafe`
+      CREATE UNIQUE INDEX posts_replica_identity_uidx ON posts (id, created_at)
+    `);
+    await transaction.query(sql.unsafe`
+      CREATE TABLE posts_202401 PARTITION OF posts
+        FOR VALUES FROM ('2024-01-01') TO ('2024-02-01')
+    `);
+    await transaction.query(sql.unsafe`
+      CREATE TABLE posts_202402 PARTITION OF posts
+        FOR VALUES FROM ('2024-02-01') TO ('2024-03-01')
+    `);
+    await transaction.query(sql.unsafe`
+      ALTER TABLE posts_202401 ADD PRIMARY KEY (id)
+    `);
+    await transaction.query(sql.unsafe`
+      ALTER TABLE posts_202402 ADD PRIMARY KEY (id)
+    `);
+    await transaction.query(sql.unsafe`
+      CREATE TABLE posts_retired (
+        id INTEGER PRIMARY KEY,
+        title TEXT NOT NULL,
+        created_at DATE NOT NULL
+      )
+    `);
+    await pgslice.enableMirroring(transaction, {
+      table: "posts",
+      targetType: "retired",
+    });
+  });
+
+  test("clones the trigger onto every partition", async ({ transaction }) => {
+    expect(
+      await triggerCount(transaction, "posts_retired_mirror_trigger"),
+    ).toBe(3);
+  });
+
+  test("locks the parent before it drops anything", async ({ transaction }) => {
+    const issued: string[] = [];
+    const recording: DatabaseTransactionConnection = {
+      ...transaction,
+      async query(query) {
+        issued.push(query.sql.replace(/\s+/g, " ").trim());
+        return transaction.query(query);
+      },
+    };
+
+    await new Mirroring({
+      source: Table.parse("posts"),
+      targetType: "retired",
+    }).disable(recording);
+
+    expect(issued[0]).toMatch(
+      /^LOCK TABLE "public"\."posts" IN ACCESS EXCLUSIVE MODE$/,
+    );
+    expect(issued[1]).toMatch(/^DROP TRIGGER IF EXISTS/);
+    expect(issued[2]).toMatch(/^DROP FUNCTION IF EXISTS/);
+
+    const held = await transaction.oneFirst(
+      sql.type(z.object({ count: z.coerce.number() }))`
+        SELECT COUNT(*)::int FROM pg_locks
+        WHERE relation = 'posts'::regclass
+          AND pid = pg_backend_pid()
+          AND mode = 'AccessExclusiveLock'
+          AND granted
+      `,
+    );
+    expect(held).toBe(1);
+  });
+
+  test("enable locks the parent before it drops and recreates the trigger", async ({
+    transaction,
+  }) => {
+    const issued: string[] = [];
+    const recording: DatabaseTransactionConnection = {
+      ...transaction,
+      async query(query) {
+        issued.push(query.sql.replace(/\s+/g, " ").trim());
+        return transaction.query(query);
+      },
+    };
+
+    await new Mirroring({
+      source: Table.parse("posts"),
+      targetType: "retired",
+    }).enable(recording, Table.parse("posts_retired"));
+
+    const lockIndex = issued.findIndex((statement) =>
+      /^LOCK TABLE "public"\."posts" IN ACCESS EXCLUSIVE MODE$/.test(statement),
+    );
+    const dropIndex = issued.findIndex((statement) =>
+      /^DROP TRIGGER IF EXISTS/.test(statement),
+    );
+    expect(lockIndex).toBeGreaterThanOrEqual(0);
+    expect(dropIndex).toBeGreaterThan(lockIndex);
+  });
+
+  test("re-running enable leaves one trigger on the parent and each partition", async ({
+    pgslice,
+    transaction,
+  }) => {
+    await pgslice.enableMirroring(transaction, {
+      table: "posts",
+      targetType: "retired",
+    });
+
+    expect(
+      await triggerCount(transaction, "posts_retired_mirror_trigger"),
+    ).toBe(3);
+  });
+
+  test("removes the trigger from the parent and every partition", async ({
+    pgslice,
+    transaction,
+  }) => {
+    await pgslice.disableMirroring(transaction, {
+      table: "posts",
+      targetType: "retired",
+    });
+
+    expect(
+      await triggerCount(transaction, "posts_retired_mirror_trigger"),
+    ).toBe(0);
+    const functions = await transaction.oneFirst(
+      sql.type(z.object({ count: z.coerce.number() }))`
+        SELECT COUNT(*)::int FROM pg_proc WHERE proname = 'posts_mirror_to_retired'
+      `,
+    );
+    expect(functions).toBe(0);
+  });
+
+  test("bounds the lock wait with the requested lock_timeout", async ({
+    pgslice,
+    transaction,
+  }) => {
+    await pgslice.disableMirroring(transaction, {
+      table: "posts",
+      targetType: "retired",
+      lockTimeout: "7s",
+    });
+
+    const lockTimeout = await transaction.oneFirst(
+      sql.type(z.object({ lock_timeout: z.string() }))`SHOW lock_timeout`,
+    );
+    expect(lockTimeout).toBe("7s");
+  });
+});

src/mirroring.ts

@@ -14,6 +14,14 @@ interface MirroringOptions {
  * Manages database mirroring triggers between a source table and a target table.
  * These triggers ensure that INSERT, UPDATE, and DELETE operations on the source
  * table are automatically replicated to the target table.
+ *
+ * When the source is partitioned, Postgres clones the trigger onto every
+ * partition, and `DROP TRIGGER` removes the clones before the parent trigger,
+ * taking an ACCESS EXCLUSIVE lock on each partition before the parent.
+ * Statements enter through the parent and lock it first, so the two walk the
+ * same relations in opposite orders and deadlock under live traffic. Dropping
+ * therefore locks the parent explicitly before anything else, and so does
+ * enabling, which drops any existing trigger before recreating it.
  */
 export class Mirroring {
   readonly #source: Table;
@@ -40,6 +48,7 @@ export class Mirroring {
       target,
     );
 
+    await this.#lockSource(tx);
     await tx.query(sql.typeAlias("void")`${functionSql}`);
     await tx.query(sql.typeAlias("void")`${this.#dropTriggerSql}`);
     await tx.query(sql.typeAlias("void")`${this.#createTriggerSql}`);
@@ -49,10 +58,19 @@ export class Mirroring {
    * Disables mirroring by dropping the trigger and function.
    */
   async disable(tx: DatabaseTransactionConnection): Promise<void> {
+    await this.#lockSource(tx);
     await tx.query(sql.typeAlias("void")`${this.#dropTriggerSql}`);
     await tx.query(sql.typeAlias("void")`${this.#dropFunctionSql}`);
   }
 
+  async #lockSource(tx: DatabaseTransactionConnection): Promise<void> {
+    await tx.query(
+      sql.typeAlias("void")`
+        LOCK TABLE ${this.#source.sqlIdentifier} IN ACCESS EXCLUSIVE MODE
+      `,
+    );
+  }
+
   get #functionName() {
     const suffix =
       this.#targetType === "intermediate"

src/pgslice.ts

@@ -803,6 +803,11 @@ export class Pgslice {
           throw new Error(`Table not found: ${target.toString()}`);
         }
 
+        await tx.query(
+          sql.typeAlias(
+            "void",
+          )`SET LOCAL lock_timeout = ${sql.literalValue(options.lockTimeout ?? "5s")}`,
+        );
         await new Mirroring({ source: table, targetType }).enable(tx, target);
       }),
     );
@@ -826,6 +831,11 @@ export class Pgslice {
           throw new Error(`Table not found: ${table.toString()}`);
         }
 
+        await tx.query(
+          sql.typeAlias(
+            "void",
+          )`SET LOCAL lock_timeout = ${sql.literalValue(options.lockTimeout ?? "5s")}`,
+        );
         await new Mirroring({ source: table, targetType }).disable(tx);
       }),
     );