feature: Add maintain command for automated runway provisioning

src/advisory-lock.test.ts

@@ -1,5 +1,5 @@
 import { describe, expect } from "vitest";
-import { createPool } from "slonik";
+import { createPool, sql } from "slonik";
 
 import { pgsliceTest as test } from "./testing/index.js";
 import { AdvisoryLock, AdvisoryLockError } from "./advisory-lock.js";
@@ -20,7 +20,9 @@ describe("AdvisoryLock.withLock", () => {
     expect(result).toBe("success");
   });
 
-  test("releases lock even if handler throws", async ({ transaction }) => {
+  test("propagates the handler's error and stays re-acquirable in the transaction", async ({
+    transaction,
+  }) => {
     const table = Table.parse("test_table");
 
     await expect(
@@ -29,7 +31,7 @@ describe("AdvisoryLock.withLock", () => {
       }),
     ).rejects.toThrow("handler error");
 
-    // Should be able to acquire the lock again since it was released
+    // Re-acquirable within the same transaction (xact advisory locks are re-entrant)
     const result = await AdvisoryLock.withLock(
       transaction,
       table,
@@ -75,6 +77,49 @@ describe("AdvisoryLock.withLock", () => {
       await pool2.end();
     }
   });
+
+  test("does not leak the lock when the handler aborts the transaction", async ({
+    databaseUrl,
+  }) => {
+    const table = Table.parse("test_table");
+    const operation = "test_op";
+
+    const pool1 = await createPool(databaseUrl.toString(), {
+      maximumPoolSize: 1,
+      queryRetryLimit: 0,
+    });
+    const pool2 = await createPool(databaseUrl.toString(), {
+      maximumPoolSize: 1,
+      queryRetryLimit: 0,
+    });
+
+    try {
+      // The handler aborts its transaction with a failing statement. A
+      // session-scoped lock survived the rollback (leaked); a transaction-scoped
+      // lock is released when the transaction rolls back.
+      await expect(
+        pool1.transaction((tx1) =>
+          AdvisoryLock.withLock(tx1, table, operation, async () => {
+            await tx1.query(sql.unsafe`SELECT 1 FROM pgslice_missing_relation`);
+          }),
+        ),
+      ).rejects.toThrow();
+
+      // A different session can now take the lock — proof it was released.
+      await pool2.transaction(async (tx2) => {
+        const result = await AdvisoryLock.withLock(
+          tx2,
+          table,
+          operation,
+          async () => "acquired after rollback",
+        );
+        expect(result).toBe("acquired after rollback");
+      });
+    } finally {
+      await pool1.end();
+      await pool2.end();
+    }
+  });
 });
 
 describe("AdvisoryLock.acquire", () => {

src/advisory-lock.ts

@@ -16,21 +16,27 @@ export class AdvisoryLockError extends Error {
 
 export abstract class AdvisoryLock {
   /**
-   * Executes a handler while holding an advisory lock.
-   * The lock is automatically released when the handler completes or throws.
+   * Runs a handler while holding a transaction-scoped advisory lock. The lock is
+   * acquired on the given transaction and released automatically when that
+   * transaction commits or rolls back — so a handler that aborts the transaction
+   * cannot leak the lock, and there is no unlock query that could fail on an
+   * aborted transaction and mask the handler's error.
+   *
+   * Must be called inside a transaction (every pgslice command wraps it in
+   * `connection.transaction(...)`). For a lock that must span multiple
+   * transactions — e.g. a batched generator — use {@link acquire} instead.
    */
   static async withLock<T>(
     connection: CommonQueryMethods,
     table: Table,
     operation: string,
     handler: () => Promise<T>,
   ): Promise<T> {
-    const release = await this.acquire(connection, table, operation);
-    try {
-      return await handler();
-    } finally {
-      await release();
+    const key = await this.#getKey(connection, table, operation);
+    if (!(await this.#tryAcquireXact(connection, key))) {
+      throw new AdvisoryLockError(table, operation);
     }
+    return handler();
   }
 
   /**
@@ -80,6 +86,18 @@ export abstract class AdvisoryLock {
     return result.acquired;
   }
 
+  static async #tryAcquireXact(
+    connection: CommonQueryMethods,
+    key: bigint,
+  ): Promise<boolean> {
+    const result = await connection.one(
+      sql.type(z.object({ acquired: z.boolean() }))`
+        SELECT pg_try_advisory_xact_lock(${key}) AS acquired
+      `,
+    );
+    return result.acquired;
+  }
+
   static async #release(
     connection: CommonQueryMethods,
     key: bigint,

src/cli.ts

@@ -7,6 +7,7 @@ import { DisableRetiredMirroringCommand } from "./commands/disable-retired-mirro
 import { EnableMirroringCommand } from "./commands/enable-mirroring.js";
 import { EnableRetiredMirroringCommand } from "./commands/enable-retired-mirroring.js";
 import { FillCommand } from "./commands/fill.js";
+import { MaintainCommand } from "./commands/maintain.js";
 import { PrepCommand } from "./commands/prep.js";
 import { StatusCommand } from "./commands/status.js";
 import { SwapCommand } from "./commands/swap.js";
@@ -30,6 +31,7 @@ export function createCli(): Cli {
   cli.register(EnableMirroringCommand);
   cli.register(EnableRetiredMirroringCommand);
   cli.register(FillCommand);
+  cli.register(MaintainCommand);
   cli.register(PrepCommand);
   cli.register(StatusCommand);
   cli.register(SwapCommand);

src/commands/base.ts

@@ -25,7 +25,7 @@ export abstract class BaseCommand extends Command<Context> {
         new URL(this.getDatabaseUrl()),
       );
 
-      await this.perform(this.context.pgslice);
+      return await this.perform(this.context.pgslice);
     } catch (error) {
       if (error instanceof Error) {
         this.context.stderr.write(`${error.message}\n`);

src/commands/maintain.test.ts

@@ -0,0 +1,233 @@
+import { describe, expect } from "vitest";
+import { sql, type DatabaseTransactionConnection } from "slonik";
+
+import { commandTest as test } from "../testing/index.js";
+import { MaintainCommand } from "./maintain.js";
+
+interface LogEntry {
+  jobId: string;
+  msg: string;
+  level: string;
+  target: { db: string; host: string; schema?: string; table?: string };
+  future?: { daily: number; weekly: number; monthly: number; yearly: number };
+  partitions?: { new: number; total: number };
+  success?: number;
+  succeeded?: { count: number; tables: string[] };
+  failed?: { count: number; tables: string[] };
+}
+
+/** Parse the command's stdout as JSONL (one log record per line). */
+function jsonLines(output: string | undefined): LogEntry[] {
+  return (output ?? "")
+    .split("\n")
+    .filter((line) => line.trim().length > 0)
+    .map((line): LogEntry => JSON.parse(line));
+}
+
+describe("MaintainCommand", () => {
+  test.scoped({ commandClass: ({}, use) => use(MaintainCommand) });
+
+  async function createPosts(transaction: DatabaseTransactionConnection) {
+    await transaction.query(sql.unsafe`
+      CREATE TABLE posts (
+        id bigint NOT NULL,
+        created_at timestamp without time zone NOT NULL,
+        PRIMARY KEY (id, created_at)
+      ) PARTITION BY RANGE (created_at)
+    `);
+    await transaction.query(sql.unsafe`
+      COMMENT ON TABLE posts IS 'column:created_at,period:month,cast:date,version:3'
+    `);
+    await transaction.query(sql.unsafe`
+      CREATE TABLE posts_y2026m01 PARTITION OF posts
+      FOR VALUES FROM ('2026-01-01') TO ('2026-02-01')
+    `);
+  }
+
+  test("emits JSONL and exits 0 on a healthy fleet", async ({
+    cli,
+    commandContext,
+    transaction,
+  }) => {
+    await createPosts(transaction);
+
+    const exitCode = await cli.run(["maintain"], commandContext);
+    expect(exitCode).toBe(0);
+
+    const logs = jsonLines(commandContext.stdout.read()?.toString());
+
+    // Every record shares the shape, uses only info/error, and never leaks the
+    // partitioning model.
+    for (const entry of logs) {
+      expect(typeof entry.msg).toBe("string");
+      expect(["info", "error"]).toContain(entry.level);
+      expect(entry.target.db).toBeTypeOf("string");
+      expect(entry.target.host).toBe("localhost");
+      expect(JSON.stringify(entry)).not.toContain("model");
+    }
+
+    // Every record from one run shares a single generated jobId (a UUID).
+    const jobIds = new Set(logs.map((entry) => entry.jobId));
+    expect(jobIds.size).toBe(1);
+    expect([...jobIds][0]).toMatch(
+      /^[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}$/i,
+    );
+
+    const start = logs[0];
+    expect(start.msg).toBe("Running pgslice maintain");
+    expect(start.level).toBe("info");
+    expect(start.future).toEqual({
+      daily: 90,
+      weekly: 26,
+      monthly: 6,
+      yearly: 1,
+    });
+
+    const table = logs.find((entry) => entry.target.table === "posts");
+    expect(table?.msg).toBe("Extended table successfully");
+    expect(table?.level).toBe("info");
+    expect(table?.success).toBe(1);
+    expect(table?.target.schema).toBe("public");
+    expect(table?.partitions).toEqual({
+      new: expect.any(Number),
+      total: expect.any(Number),
+    });
+
+    const final = logs.at(-1);
+    expect(final?.msg).toBe("Finished pgslice maintain successfully");
+    expect(final?.level).toBe("info");
+    expect(final?.succeeded).toEqual({ count: 1, tables: ["public.posts"] });
+    expect(final?.failed).toEqual({ count: 0, tables: [] });
+  });
+
+  test("exits 1 and logs an error record when a leaf has no usable replica identity", async ({
+    cli,
+    commandContext,
+    transaction,
+  }) => {
+    await createPosts(transaction);
+    await transaction.query(sql.unsafe`
+      ALTER TABLE posts_y2026m01 REPLICA IDENTITY NOTHING
+    `);
+
+    const exitCode = await cli.run(["maintain"], commandContext);
+    expect(exitCode).toBe(1);
+
+    const logs = jsonLines(commandContext.stdout.read()?.toString());
+
+    const table = logs.find((entry) => entry.target.table === "posts");
+    expect(table?.level).toBe("error");
+    expect(table?.success).toBe(0);
+    expect(table?.msg).toContain("replica identity");
+    expect(table?.target.host).toBe("localhost");
+
+    const final = logs.at(-1);
+    expect(final?.msg).toBe("Finished pgslice maintain with errors");
+    expect(final?.level).toBe("error");
+    expect(final?.failed).toEqual({ count: 1, tables: ["public.posts"] });
+  });
+
+  test("exits 0 with an empty summary when no managed tables exist", async ({
+    cli,
+    commandContext,
+    transaction,
+  }) => {
+    await transaction.query(sql.unsafe`
+      CREATE TABLE plain (id bigint NOT NULL, created_at timestamptz NOT NULL)
+      PARTITION BY RANGE (created_at)
+    `);
+
+    const exitCode = await cli.run(["maintain"], commandContext);
+    expect(exitCode).toBe(0);
+
+    const logs = jsonLines(commandContext.stdout.read()?.toString());
+    expect(logs[0].msg).toBe("Running pgslice maintain");
+    expect(logs.some((entry) => entry.target.table !== undefined)).toBe(false);
+
+    const final = logs.at(-1);
+    expect(final?.msg).toBe("Finished pgslice maintain successfully");
+    expect(final?.succeeded).toEqual({ count: 0, tables: [] });
+    expect(final?.failed).toEqual({ count: 0, tables: [] });
+  });
+
+  test("logs a no-op record when a table already has runway beyond the horizon", async ({
+    cli,
+    commandContext,
+    transaction,
+  }) => {
+    // A managed parent whose only partition is far in the future: the forward
+    // horizon is already covered, so maintain creates nothing.
+    await transaction.query(sql.unsafe`
+      CREATE TABLE posts (
+        id bigint NOT NULL,
+        created_at timestamp without time zone NOT NULL,
+        PRIMARY KEY (id, created_at)
+      ) PARTITION BY RANGE (created_at)
+    `);
+    await transaction.query(sql.unsafe`
+      COMMENT ON TABLE posts IS 'column:created_at,period:month,cast:date,version:3'
+    `);
+    await transaction.query(sql.unsafe`
+      CREATE TABLE posts_y2099m01 PARTITION OF posts
+      FOR VALUES FROM ('2099-01-01') TO ('2099-02-01')
+    `);
+
+    const exitCode = await cli.run(["maintain"], commandContext);
+    expect(exitCode).toBe(0);
+
+    const logs = jsonLines(commandContext.stdout.read()?.toString());
+    const table = logs.find((entry) => entry.target.table === "posts");
+    expect(table?.msg).toBe("Table already up to date; no extension needed");
+    expect(table?.level).toBe("info");
+    expect(table?.success).toBe(1);
+    expect(table?.partitions).toEqual({ new: 0, total: 1 });
+    expect(table?.target.host).toBe("localhost");
+  });
+
+  test("reads per-period horizons from PGSLICE_FUTURE_* env vars", async ({
+    cli,
+    commandContext,
+    transaction,
+  }) => {
+    await createPosts(transaction);
+    process.env.PGSLICE_FUTURE_MONTHLY = "3";
+    process.env.PGSLICE_FUTURE_WEEKLY = "9";
+    try {
+      const exitCode = await cli.run(["maintain"], commandContext);
+      expect(exitCode).toBe(0);
+
+      const logs = jsonLines(commandContext.stdout.read()?.toString());
+      // env overrides the baked default for its period; unset periods keep theirs.
+      expect(logs[0].future).toEqual({
+        daily: 90,
+        weekly: 9,
+        monthly: 3,
+        yearly: 1,
+      });
+    } finally {
+      delete process.env.PGSLICE_FUTURE_MONTHLY;
+      delete process.env.PGSLICE_FUTURE_WEEKLY;
+    }
+  });
+
+  test("an explicit --future-monthly flag overrides the env var", async ({
+    cli,
+    commandContext,
+    transaction,
+  }) => {
+    await createPosts(transaction);
+    process.env.PGSLICE_FUTURE_MONTHLY = "3";
+    try {
+      const exitCode = await cli.run(
+        ["maintain", "--future-monthly", "7"],
+        commandContext,
+      );
+      expect(exitCode).toBe(0);
+
+      const logs = jsonLines(commandContext.stdout.read()?.toString());
+      expect(logs[0].future?.monthly).toBe(7);
+    } finally {
+      delete process.env.PGSLICE_FUTURE_MONTHLY;
+    }
+  });
+});