118 lines
9.3 KiB
TypeScript
118 lines
9.3 KiB
TypeScript
import assert from "node:assert/strict"
|
|
import type { PoolClient } from "pg"
|
|
import { sql } from "drizzle-orm"
|
|
import { admitPartitionWriter, withPartitionWrite, PartitionWriteError, isPartitionWriteOutcomeUnknown, isPartitionWriteClient, PARTITION_WRITER_ADVISORY_KEY, type PartitionWriteOwner } from "./partition-write-transaction.js"
|
|
|
|
function fixture() {
|
|
const statements: Array<{ text: string; values: unknown[] }> = [], releases: boolean[] = []
|
|
let fail: string | undefined, committed = 0, connects = 0
|
|
let commitCommand: string | undefined = "COMMIT"
|
|
const client = {
|
|
query: async (raw: string | { text: string; values?: unknown[] }, values: unknown[] = []) => {
|
|
const text = typeof raw === "string" ? raw : raw.text
|
|
statements.push({ text, values: typeof raw === "string" ? values : raw.values ?? values })
|
|
if (text === "COMMIT") committed++
|
|
if (text === fail) throw Error("Injected uncertain SQL ACK")
|
|
return { rows: [{ value: 1 }], rowCount: 1, command: text === "COMMIT" ? commitCommand : "SELECT" }
|
|
},
|
|
release: (destroy: boolean) => { releases.push(destroy) },
|
|
} as unknown as PoolClient
|
|
return { client, statements, releases, connect: async () => { connects++; return client }, fail: (text: string) => { fail = text }, committed: () => committed, connects: () => connects, command: (value: string | undefined) => { commitCommand = value } }
|
|
}
|
|
let cases = 0
|
|
async function test(name: string, work: () => Promise<void>) { await work(); cases++; console.log("ok", name) }
|
|
await test("raw and Drizzle use same backend after bounded shared admission", async () => {
|
|
const f = fixture()
|
|
const result = await withPartitionWrite(async owner => { await owner.client.query("INSERT INTO traffic_samples VALUES ($1)", [7]); await owner.db.execute(sql`SELECT ${9} AS value`); return "done" }, { connect: f.connect })
|
|
assert.equal(result, "done"); assert.equal(f.connects(), 1); assert.deepEqual(f.releases, [false])
|
|
assert.equal(f.statements[0].text, "BEGIN")
|
|
assert.match(f.statements[1].text, /^SET LOCAL lock_timeout='\d+ms'$/)
|
|
assert.match(f.statements[2].text, /^SET LOCAL statement_timeout='\d+ms'$/)
|
|
assert.equal(f.statements[3].text, "SELECT pg_advisory_xact_lock_shared($1::bigint)"); assert.deepEqual(f.statements[3].values, [PARTITION_WRITER_ADVISORY_KEY])
|
|
assert.equal(f.statements[4].text, "INSERT INTO traffic_samples VALUES ($1)"); assert.equal(f.statements.at(-1)?.text, "COMMIT")
|
|
})
|
|
await test("caller-owned transaction admission never starts or commits lifecycle", async () => {
|
|
const f = fixture(); await admitPartitionWriter(f.client)
|
|
assert.equal(f.statements.length, 3); assert(f.statements[2].text.includes("pg_advisory_xact_lock_shared")); assert.equal(f.releases.length, 0)
|
|
})
|
|
for (const boundary of ["BEGIN", "COMMIT"] as const) await test("lost " + boundary + " ACK destroys without rollback or retry", async () => {
|
|
const f = fixture(); f.fail(boundary)
|
|
await assert.rejects(withPartitionWrite(async owner => { await owner.client.query("INSERT INTO traffic_samples VALUES (1)") }, { connect: f.connect }), error => isPartitionWriteOutcomeUnknown(error) && error.automaticReplay === false && error.code === "PARTITION_WRITE_UNKNOWN")
|
|
assert.deepEqual(f.releases, [true]); assert.equal(f.statements.filter(row => row.text === boundary).length, 1); assert(!f.statements.some(row => row.text === "ROLLBACK"))
|
|
assert.equal(f.committed(), boundary === "COMMIT" ? 1 : 0)
|
|
})
|
|
await test("detached raw and Drizzle handles expire before backend returns to pool", async () => {
|
|
const f = fixture(); let captured!: PartitionWriteOwner
|
|
await withPartitionWrite(async owner => { captured = owner }, { connect: f.connect })
|
|
const before = f.statements.length
|
|
await assert.rejects(captured.client.query("SELECT 1"), PartitionWriteError)
|
|
await assert.rejects(captured.db.execute(sql`SELECT 1`))
|
|
assert.equal(f.statements.length, before); assert.deepEqual(f.releases, [false])
|
|
})
|
|
for (const text of ["COMMIT", "/* hidden */ COMMIT", "-- leading\nROLLBACK", "SELECT 1; COMMIT", "SAVEPOINT bypass"]) await test("scope refuses transaction escape " + text, async () => {
|
|
const f = fixture()
|
|
await assert.rejects(withPartitionWrite(async owner => { await owner.client.query(text) }, { connect: f.connect }), PartitionWriteError)
|
|
assert(!f.statements.some(row => row.text === text)); assert.deepEqual(f.releases, [true])
|
|
})
|
|
await test("caught query failure still prevents commit", async () => {
|
|
const f = fixture(); f.fail("SELECT fail")
|
|
await assert.rejects(withPartitionWrite(async owner => { await owner.client.query("SELECT fail").catch(() => {}) }, { connect: f.connect }), error => error instanceof PartitionWriteError && !error.outcomeUnknown)
|
|
assert(!f.statements.some(row => row.text === "COMMIT")); assert.deepEqual(f.releases, [true])
|
|
})
|
|
await test("callback raw query remains guarded and completes before commit", async () => {
|
|
const f = fixture()
|
|
await withPartitionWrite(async owner => { await new Promise<void>((resolve, reject) => { owner.client.query("SELECT 1", error => error ? reject(error) : resolve()) }) }, { connect: f.connect })
|
|
assert.equal(f.statements.at(-2)?.text, "SELECT 1"); assert.equal(f.statements.at(-1)?.text, "COMMIT")
|
|
})
|
|
await test("abort destroys owner and settles hanging work without waiting for callback", async () => {
|
|
const f = fixture(), controller = new AbortController(); let entered!: () => void
|
|
const started = new Promise<void>(resolve => { entered = resolve })
|
|
const pending = withPartitionWrite(async () => { entered(); await new Promise<never>(() => {}) }, { connect: f.connect, signal: controller.signal })
|
|
await started; controller.abort()
|
|
await assert.rejects(pending, PartitionWriteError); assert.deepEqual(f.releases, [true]); assert(!f.statements.some(row => row.text === "COMMIT"))
|
|
})
|
|
await test("deadline destroys owner and settles hanging work", async () => {
|
|
const f = fixture()
|
|
await assert.rejects(withPartitionWrite(async () => new Promise<never>(() => {}), { connect: f.connect, deadlineAt: Date.now() + 30 }), PartitionWriteError)
|
|
assert.deepEqual(f.releases, [true]); assert(!f.statements.some(row => row.text === "COMMIT"))
|
|
})
|
|
await test("abort settles an authority callback that never returns", async () => {
|
|
const f = fixture(), controller = new AbortController(); let entered!: () => void
|
|
const started = new Promise<void>(resolve => { entered = resolve })
|
|
const pending = withPartitionWrite(async () => {}, { connect: f.connect, signal: controller.signal, authorize: async () => { entered(); await new Promise<never>(() => {}) } })
|
|
await started; controller.abort(); await assert.rejects(pending, PartitionWriteError)
|
|
assert.deepEqual(f.releases, [true]); assert.equal(f.statements.length, 0)
|
|
})
|
|
await test("expired, overly long and pre-aborted admission never connects", async () => {
|
|
const f = fixture(), controller = new AbortController(); controller.abort()
|
|
for (const options of [{ deadlineAt: Date.now() - 1 }, { deadlineAt: Date.now() + 130000 }, { signal: controller.signal }]) await assert.rejects(withPartitionWrite(async () => {}, { ...options, connect: f.connect }), PartitionWriteError)
|
|
assert.equal(f.connects(), 0)
|
|
})
|
|
await test("revoked authority before commit cannot publish a success", async () => {
|
|
const f = fixture(); let active = true
|
|
await assert.rejects(withPartitionWrite(async owner => { await owner.client.query("SELECT 1"); active = false }, { connect: f.connect, authorize: async () => { assert(active) } }), PartitionWriteError)
|
|
assert(!f.statements.some(row => row.text === "COMMIT")); assert.deepEqual(f.releases, [true])
|
|
})
|
|
await test("late connected backend after abort is destroyed rather than leaked", async () => {
|
|
const f = fixture(), controller = new AbortController(); let connected!: (client: PoolClient) => void
|
|
const pending = withPartitionWrite(async () => {}, { signal: controller.signal, connect: () => new Promise(resolve => { connected = resolve }) })
|
|
controller.abort(); await assert.rejects(pending, PartitionWriteError)
|
|
connected(f.client); await new Promise(resolve => setImmediate(resolve))
|
|
assert.deepEqual(f.releases, [true]); assert.equal(f.statements.length, 0)
|
|
})
|
|
await test("owned client identity cannot be forged and expires before pooling", async () => {
|
|
const f=fixture();let client:PoolClient|undefined
|
|
assert.equal(isPartitionWriteClient(f.client),false)
|
|
assert.equal(isPartitionWriteClient({query:f.client.query}),false)
|
|
await withPartitionWrite(async owner=>{client=owner.client;assert.equal(isPartitionWriteClient(client),true)}, {connect:f.connect})
|
|
assert.equal(isPartitionWriteClient(client),false)
|
|
})
|
|
for (const command of ["ROLLBACK", "INSERT", undefined]) await test("resolved non-COMMIT acknowledgement refuses " + String(command), async () => {
|
|
const f = fixture(); f.command(command)
|
|
await assert.rejects(withPartitionWrite(async owner => { await owner.client.query("INSERT INTO traffic_samples VALUES (1)") }, { connect: f.connect }), error => isPartitionWriteOutcomeUnknown(error) && error.automaticReplay === false)
|
|
assert.equal(f.statements.filter(row => row.text === "COMMIT").length, 1)
|
|
assert.equal(f.statements.filter(row => row.text.startsWith("INSERT")).length, 1)
|
|
assert(!f.statements.some(row => row.text === "ROLLBACK")); assert.deepEqual(f.releases, [true])
|
|
})
|
|
console.log(`Partition writer transaction: ${cases} tests passed`)
|