diff --git a/cloud/apps/push/src/push-database.ts b/cloud/apps/push/src/push-database.ts index a2ed2bce9d2..019a76a7cc4 100644 --- a/cloud/apps/push/src/push-database.ts +++ b/cloud/apps/push/src/push-database.ts @@ -213,7 +213,10 @@ async function applySchemaOnUntimedPool( const database = new PostgresDatabase(pool) try { await applyPostgresSchema(pushSchemaStatements(), (statement) => database.query(statement), { - eventPrefix: 'orca_push_postgres_schema' + eventPrefix: 'orca_push_postgres_schema', + // Push has no catalog pre-check, so a lock timeout here says nothing about whether the + // object already exists and the old bounded retry is still the right answer. + retryLockTimeout: true }) } finally { await database.close().catch(() => undefined) diff --git a/cloud/apps/relay/src/database-postgres-timeout.test.ts b/cloud/apps/relay/src/database-postgres-timeout.test.ts index df300383c1b..aeb730496c9 100644 --- a/cloud/apps/relay/src/database-postgres-timeout.test.ts +++ b/cloud/apps/relay/src/database-postgres-timeout.test.ts @@ -119,15 +119,21 @@ describe('PostgreSQL relay deadlines', () => { dataDir: './unused' }) - expect(ddl.length).toBeGreaterThan(0) - expect(ddl).toContain(POSTGRES_STATEMENT_STATS_MIGRATION) + // The catalog pre-check reads pg_catalog on the same untimed connection before each + // lock-taking statement, so the schema pool now carries reads as well as DDL. + const probes = ddl.filter((statement) => /^SELECT\b/i.test(statement)) + const statements = ddl.filter((statement) => !/^SELECT\b/i.test(statement)) + expect(probes.length).toBeGreaterThan(0) + expect(probes.every((statement) => statement.includes('pg_catalog'))).toBe(true) + expect(statements.length).toBeGreaterThan(0) + expect(statements).toContain(POSTGRES_STATEMENT_STATS_MIGRATION.trim()) // Statements can open with a leading `--` rationale comment. const body = (statement: string): string => statement.replace(/^(?:\s*--[^\n]*\n)*\s*/, '') expect( - ddl.every( + statements.every( (statement) => - statement === POSTGRES_STATEMENT_STATS_MIGRATION || + statement === POSTGRES_STATEMENT_STATS_MIGRATION.trim() || /^(?:CREATE|ALTER TABLE)\b/i.test(body(statement)) ) ).toBe(true) @@ -201,11 +207,11 @@ describe('PostgreSQL relay deadlines', () => { }) describe('PostgreSQL schema startup', () => { - it('retries lock and statement timeouts with bounded backoff', async () => { + it('retries statement timeouts with bounded backoff', async () => { vi.spyOn(console, 'warn').mockImplementation(() => undefined) const query = vi .fn<(statement: string) => Promise>() - .mockRejectedValueOnce(Object.assign(new Error('lock timeout'), { code: '55P03' })) + .mockRejectedValueOnce(Object.assign(new Error('statement timeout'), { code: '57014' })) .mockRejectedValueOnce(Object.assign(new Error('statement timeout'), { code: '57014' })) .mockResolvedValue(undefined) const delays: number[] = [] @@ -221,6 +227,31 @@ describe('PostgreSQL schema startup', () => { expect(delays).toEqual([125, 250]) }) + it('fails the boot on a lock timeout instead of re-entering the lock queue', async () => { + // The catalog pre-check already answered that the object is missing, so a lock timeout means + // this boot lost the queue. Relation locks are granted in queue order, so each retry parks + // every writer behind it for another timeout. + const errors: string[] = [] + vi.spyOn(console, 'error').mockImplementation((line: string) => { + errors.push(line) + }) + const error = Object.assign(new Error('lock timeout'), { code: '55P03' }) + const query = vi.fn<(statement: string) => Promise>().mockRejectedValue(error) + const pause = vi.fn(async () => undefined) + + await expect( + applyPostgresSchema(['CREATE INDEX IF NOT EXISTS i ON t(c)'], query, { wait: pause }) + ).rejects.toBe(error) + + expect(query).toHaveBeenCalledTimes(1) + expect(pause).not.toHaveBeenCalled() + expect(JSON.parse(errors[0] ?? '{}')).toMatchObject({ + event: 'orca_relay_postgres_schema_lock_timeout', + code: '55P03', + statement: 'CREATE INDEX IF NOT EXISTS i ON t(c)' + }) + }) + it('retries only the PostgreSQL concurrent type-creation collision', async () => { vi.spyOn(console, 'warn').mockImplementation(() => undefined) const collision = Object.assign(new Error('duplicate type'), { @@ -386,7 +417,7 @@ describe('PostgreSQL schema startup', () => { it('stops retrying at the shared startup deadline', async () => { vi.spyOn(console, 'warn').mockImplementation(() => undefined) - const error = Object.assign(new Error('lock timeout'), { code: '55P03' }) + const error = Object.assign(new Error('statement timeout'), { code: '57014' }) const delays: number[] = [] let now = 0 const query = vi @@ -413,7 +444,7 @@ describe('PostgreSQL schema startup', () => { expect(console.warn).toHaveBeenLastCalledWith( JSON.stringify({ event: 'orca_relay_postgres_schema_retry_exhausted', - code: '55P03', + code: '57014', attempts: 2 }) ) diff --git a/cloud/apps/relay/src/database-statement-timeout-postgres.test.ts b/cloud/apps/relay/src/database-statement-timeout-postgres.test.ts index 6b7ebb0334e..a3d69d013d2 100644 --- a/cloud/apps/relay/src/database-statement-timeout-postgres.test.ts +++ b/cloud/apps/relay/src/database-statement-timeout-postgres.test.ts @@ -42,10 +42,12 @@ describePostgres('PostgreSQL statement deadline', () => { expect(result).toBe(2) }, 15_000) - // Why: DDL runs on its own untimed connection. relay_invites carries a - // CREATE INDEX IF NOT EXISTS, which (unlike CREATE TABLE IF NOT EXISTS) - // really does queue behind an ACCESS EXCLUSIVE lock on the table. - it('applies the schema behind a held ACCESS EXCLUSIVE lock', async () => { + // Why: relay_invites carries a CREATE INDEX IF NOT EXISTS, which (unlike CREATE TABLE IF NOT + // EXISTS) really does queue behind an ACCESS EXCLUSIVE lock on the table. The catalog pre-check + // asks pg_class whether that index is already there, and the read takes no lock on relay_invites, + // so a boot on a migrated database no longer joins the queue at all. It used to, and every writer + // queued behind it in lock order. + it('boots without queueing behind a held ACCESS EXCLUSIVE lock', async () => { let releaseTable!: () => void const tableReleased = new Promise((resolve) => { releaseTable = resolve @@ -61,7 +63,9 @@ describePostgres('PostgreSQL statement deadline', () => { }) await tableHeldPromise - const opening = openRelayDatabase({ + // Resolving while the lock is still held is the whole proof: a statement that queued would hit + // the schema connection's 1s lock_timeout and fail the boot, which is no longer retried. + const database = await openRelayDatabase({ databaseUrl, dataDir: '', applicationName, @@ -69,27 +73,18 @@ describePostgres('PostgreSQL statement deadline', () => { // connection must not. statementTimeoutMs: 200 }) - const blockedOnSchemaConnection = async (): Promise => { - const deadline = Date.now() + 4_000 - while (Date.now() < deadline) { - const rows = await databases[0]!.query( - `SELECT count(*) AS waiting FROM pg_stat_activity - WHERE datname = current_database() AND wait_event_type = 'Lock' - AND application_name = ?`, - [`${applicationName}/schema`] - ) - if (Number(rows[0]!.waiting) > 0) return true - await new Promise((resolve) => setTimeout(resolve, 10)) - } - return false - } - const blocked = await blockedOnSchemaConnection() + databases.push(database) + const waiting = await databases[0]!.query( + `SELECT count(*) AS waiting FROM pg_stat_activity + WHERE datname = current_database() AND wait_event_type = 'Lock' + AND application_name = ?`, + [`${applicationName}/schema`] + ) + expect(Number(waiting[0]!.waiting)).toBe(0) + releaseTable() await holder - const database = await opening - databases.push(database) - expect(blocked).toBe(true) // The serving pool still carries the short deadline it was opened with. expect(await database.query(`SELECT current_setting('statement_timeout') AS statement_timeout`)).toEqual([ { statement_timeout: '200ms' } diff --git a/cloud/apps/relay/src/database.ts b/cloud/apps/relay/src/database.ts index 53c178215df..4a8fd3ee975 100644 --- a/cloud/apps/relay/src/database.ts +++ b/cloud/apps/relay/src/database.ts @@ -68,6 +68,17 @@ export interface RelayDatabase { close(): Promise } +// RULE - no new index and no new column on `relay_control_connection_reservations`, +// `relay_confirm_results`, `relay_audit_events`, `relay_connection_bases`, or any other large +// table may be added to SCHEMA or to POSTGRES_SCHEMA_MIGRATIONS. The catalog pre-check skips a +// lock-taking statement only once the object exists, so a brand-new one reports missing on every +// director at once and each runs a non-concurrent build over the whole table. POSTGRES_LOCK_TIMEOUT_MS +// bounds how long that build waits for its lock, not how long it holds it. Build the index out of +// band with CREATE INDEX CONCURRENTLY first, then add it here, where the pre-check skips it forever +// after. relay-schema-lock-targets.test.ts pins the current list, so an addition fails CI. +// Constraint swaps are matched by NAME in pg_constraint, never by body, because the CHECK list is +// generated from REGION_LIST. Changing a constraint's definition under the same name therefore does +// nothing on boot: an operator drops it, and the next boot adds the current definition back. const SCHEMA = ` CREATE TABLE IF NOT EXISTS relay_invites ( user_id TEXT NOT NULL, @@ -636,6 +647,14 @@ export const POSTGRES_SCHEMA_MIGRATIONS = [ `ALTER TABLE relay_region_rehome_attempts ADD COLUMN IF NOT EXISTS source_generation BIGINT NOT NULL DEFAULT 0` ] +// The exact statement list a Postgres boot applies, in order, so the lock-target census can read +// what production runs rather than a copy of it. The SQLite path keeps SCHEMA on its own. +export function relayPostgresSchemaStatements(): string[] { + return [...SCHEMA.split(';'), ...POSTGRES_SCHEMA_MIGRATIONS] + .map((statement) => statement.trim()) + .filter((statement) => statement.length > 0) +} + function postgresSql(sql: string): string { let index = 0 return sql.replace(/\?/g, () => `$${++index}`) @@ -1114,11 +1133,14 @@ async function applySchemaOnUntimedPool( const database = new PostgresDatabase(pool) try { await applyPostgresSchema( - [ - ...SCHEMA.split(';').filter((statement) => statement.trim()), - ...POSTGRES_SCHEMA_MIGRATIONS - ], - async (statement) => await database.query(statement) + relayPostgresSchemaStatements(), + async (statement) => await database.query(statement), + // Asks the catalog whether each index or column is already there. CREATE INDEX IF NOT EXISTS + // and ALTER TABLE ADD COLUMN IF NOT EXISTS take their relation lock before the server + // evaluates the existence test, so on an already-migrated database the boot still joins the + // lock queue - and relation locks are granted in queue order, so every writer queues behind + // it. The catalog read takes no lock on the table. + { catalogQuery: async (sql, params) => await database.query(sql, params) } ) } finally { await database.close().catch(() => undefined) diff --git a/cloud/apps/relay/src/relay-schema-catalog-precheck-postgres.test.ts b/cloud/apps/relay/src/relay-schema-catalog-precheck-postgres.test.ts new file mode 100644 index 00000000000..2ce9ad9dcb7 --- /dev/null +++ b/cloud/apps/relay/src/relay-schema-catalog-precheck-postgres.test.ts @@ -0,0 +1,216 @@ +import pg from 'pg' +import { afterAll, beforeAll, beforeEach, describe, expect, it } from 'vitest' +import { + applyPostgresSchema, + catalogObjectPresence, + schemaLockTarget, + takesRelationLock +} from '@orca-cloud/postgres-schema' +import { + openRelayDatabase, + POSTGRES_LOCK_TIMEOUT_MS, + relayPostgresSchemaStatements, + type RelayDatabase +} from './database.js' + +// The outage this guards against: CREATE INDEX IF NOT EXISTS takes its relation lock before the +// server evaluates the existence test, so on a database that already has the index the boot still +// joins the lock queue, and relation locks are granted in queue order, so every writer queues +// behind it. Only a real server can show that the catalog pre-check removes those statements. +const databaseUrl = process.env.ORCA_RELAY_TEST_POSTGRES_URL +const describePostgres = databaseUrl ? describe : describe.skip +const schema = 'relay_schema_precheck_test' + +// A table the boot would otherwise touch with two CREATE INDEX statements, and the largest table +// in production. +const LOCKED_TABLE = 'relay_control_connection_reservations' + +function scopedUrl(): string { + const url = new URL(databaseUrl!) + url.searchParams.set('options', `-c search_path=${schema}`) + return url.toString() +} + +async function onAdmin(operation: (client: pg.Client) => Promise): Promise { + const client = new pg.Client({ connectionString: databaseUrl }) + await client.connect() + try { + return await operation(client) + } finally { + await client.end() + } +} + +describePostgres('relay boot-time schema against PostgreSQL', () => { + let url = '' + let pool: pg.Pool + let sent: string[] + const opened: RelayDatabase[] = [] + + beforeAll(() => { + url = scopedUrl() + }) + + beforeEach(async () => { + await onAdmin(async (client) => { + await client.query(`DROP SCHEMA IF EXISTS ${schema} CASCADE`) + await client.query(`CREATE SCHEMA ${schema}`) + }) + // Same lock bound the boot pool uses, so a statement that still queues fails instead of + // hanging the test. + pool = new pg.Pool({ connectionString: url, max: 1, lock_timeout: POSTGRES_LOCK_TIMEOUT_MS }) + sent = [] + }) + + afterAll(async () => { + await Promise.all(opened.map((database) => database.close().catch(() => undefined))) + await onAdmin((client) => client.query(`DROP SCHEMA IF EXISTS ${schema} CASCADE`)) + }) + + async function applyRecording(): Promise<{ ran: number; skipped: number }> { + const record = async (statement: string, params: unknown[] = []): Promise => { + sent.push(statement) + return await pool.query(statement, params) + } + return await applyPostgresSchema( + relayPostgresSchemaStatements(), + (statement) => record(statement), + { catalogQuery: async (sql, params) => (await record(sql, params)).rows } + ) + } + + function lockTaking(): string[] { + return sent.filter(takesRelationLock) + } + + it('creates the schema cold, then issues no lock-taking statement on the next boot', async () => { + const cold = await applyRecording() + expect(lockTaking().length).toBeGreaterThan(0) + expect(cold.skipped).toBeGreaterThan(0) + + sent = [] + const warm = await applyRecording() + // Zero, with no exceptions: every statement that takes a relation lock has a pre-check. + expect(lockTaking()).toEqual([]) + const preCheckedCount = relayPostgresSchemaStatements().filter( + (statement) => schemaLockTarget(statement) !== undefined + ).length + expect(warm.skipped).toBe(preCheckedCount) + expect(warm.ran).toBe(relayPostgresSchemaStatements().length - preCheckedCount) + // The CREATE TABLEs still run: they resolve a name and take no lock on an existing table. + expect(warm.ran).toBeGreaterThan(0) + await pool.end() + }) + + it('skips an index a failed concurrent build left invalid instead of rebuilding it', async () => { + await applyRecording() + // A cancelled CREATE INDEX CONCURRENTLY leaves exactly this state, and IF NOT EXISTS skips it + // too, so reading indisvalid as a condition would newly take the lock it used to avoid. + await pool.query( + `UPDATE pg_catalog.pg_index SET indisvalid = false + WHERE indexrelid = to_regclass('${schema}.relay_audit_events_at')` + ) + sent = [] + await applyRecording() + expect(lockTaking()).toEqual([]) + await pool.end() + }) + + it('re-adds a constraint an operator dropped, then skips it again', async () => { + // The pre-check matches the constraint by name only, so this is the one shape where a changed + // body needs an operator: drop it, and the next boot puts the current definition back. + await applyRecording() + const named = async (): Promise => + Number( + ( + await pool.query( + `SELECT count(*) AS present FROM pg_catalog.pg_constraint + WHERE conrelid = to_regclass('relay_region_rehome_attempts') + AND conname = 'relay_region_rehome_attempts_preferred_region_valid'` + ) + ).rows[0]?.present + ) + expect(await named()).toBe(1) + + await pool.query( + `ALTER TABLE relay_region_rehome_attempts + DROP CONSTRAINT relay_region_rehome_attempts_preferred_region_valid` + ) + sent = [] + await applyRecording() + expect(lockTaking()).toHaveLength(1) + expect(await named()).toBe(1) + + sent = [] + await applyRecording() + expect(lockTaking()).toEqual([]) + await pool.end() + }) + + it('does not let a same-named index on a sibling table answer for this one', async () => { + // Index names are unique per schema, not per table, so a name freed on one table and taken on + // another is reachable. Without tying the index to the table, the pre-check reads that sibling + // as this table's index and skips the real CREATE INDEX for good. + await applyRecording() + const ask = async (table: string): Promise => + ( + await catalogObjectPresence( + async (sql, params) => (await pool.query(sql, params)).rows, + { kind: 'index', table, name: 'relay_audit_events_at', skipWhen: 'present' } + ) + ).present + + expect(await ask('relay_audit_events')).toBe(true) + await pool.query(`DROP INDEX ${schema}.relay_audit_events_at`) + await pool.query(`CREATE TABLE ${schema}.precheck_sibling (at BIGINT)`) + await pool.query(`CREATE INDEX relay_audit_events_at ON ${schema}.precheck_sibling(at)`) + + expect(await ask('relay_audit_events')).toBe(false) + expect(await ask('precheck_sibling')).toBe(true) + await pool.end() + }) + + it('boots while another session holds ACCESS EXCLUSIVE on the largest table', async () => { + // The end-to-end proof through openRelayDatabase: with the lock held, any DDL the boot still + // sent against this table would hit lock_timeout, and 55P03 is no longer retried. + const cold = await openRelayDatabase({ databaseUrl: url, dataDir: '' }) + opened.push(cold) + await pool.end() + + const holder = new pg.Client({ connectionString: url }) + await holder.connect() + await holder.query('BEGIN') + await holder.query(`LOCK TABLE ${LOCKED_TABLE} IN ACCESS EXCLUSIVE MODE`) + try { + const warm = await openRelayDatabase({ databaseUrl: url, dataDir: '' }) + opened.push(warm) + } finally { + await holder.query('ROLLBACK') + await holder.end() + } + }) + + it('fails that same boot with 55P03 when the pre-check is not wired in', async () => { + // Keeps the test above from passing vacuously: the lock really does block relay's DDL. + const cold = await openRelayDatabase({ databaseUrl: url, dataDir: '' }) + opened.push(cold) + + const holder = new pg.Client({ connectionString: url }) + await holder.connect() + await holder.query('BEGIN') + await holder.query(`LOCK TABLE ${LOCKED_TABLE} IN ACCESS EXCLUSIVE MODE`) + try { + await expect( + applyPostgresSchema( + relayPostgresSchemaStatements(), + (statement) => pool.query(statement), + { retryDeadlineMs: 0 } + ) + ).rejects.toMatchObject({ code: '55P03' }) + } finally { + await holder.query('ROLLBACK') + await holder.end() + await pool.end() + } + }) +}) diff --git a/cloud/apps/relay/src/relay-schema-lock-targets.test.ts b/cloud/apps/relay/src/relay-schema-lock-targets.test.ts new file mode 100644 index 00000000000..a490a394378 --- /dev/null +++ b/cloud/apps/relay/src/relay-schema-lock-targets.test.ts @@ -0,0 +1,200 @@ +import { describe, expect, it } from 'vitest' +import { + requireSchemaLockTarget, + schemaLockTarget, + sqlWithoutComments, + takesRelationLock, + type SchemaLockTarget +} from '@orca-cloud/postgres-schema' +import { relayPostgresSchemaStatements } from './database.js' + +// Golden pin of every boot-time statement that takes a relation lock on Postgres. Each entry with +// a kind is gated by the catalog pre-check, so it costs a catalog read on a migrated database and +// nothing more. An addition to this list is the case the RULE comment beside SCHEMA forbids: a +// brand-new index reports missing on every director at once and each one runs a non-concurrent +// build over the whole table, which is how a boot takes the site down. Build it out of band with +// CREATE INDEX CONCURRENTLY first, then add it to SCHEMA and update this list. +const GOLDEN_LOCK_TAKING: SchemaLockTarget[] = [ + { kind: 'index', table: 'relay_invites', name: 'relay_invites_device', skipWhen: 'present' }, + { kind: 'index', table: 'relay_devices', name: 'relay_devices_current_hash', skipWhen: 'present' }, + { kind: 'index', table: 'relay_devices', name: 'relay_devices_grace_hash', skipWhen: 'present' }, + { kind: 'index', table: 'relay_connection_bases', name: 'relay_connection_bases_active_deadline', skipWhen: 'present' }, + { + kind: 'index', + table: 'relay_assignment_region_preferences', + name: 'relay_assignment_region_preferences_observed', + skipWhen: 'present' + }, + { + kind: 'index', + table: 'relay_region_rehome_attempts', + name: 'relay_region_rehome_attempts_pending', + skipWhen: 'present' + }, + { + kind: 'index', + table: 'relay_region_rehome_attempts', + name: 'relay_region_rehome_attempts_host_recency', + skipWhen: 'present' + }, + { kind: 'index', table: 'relay_cell_runtime', name: 'relay_cell_runtime_heartbeat', skipWhen: 'present' }, + { + kind: 'index', + table: 'relay_cell_connection_runtime', + name: 'relay_cell_connection_runtime_heartbeat', + skipWhen: 'present' + }, + { + kind: 'index', + table: 'relay_cell_connection_snapshots', + name: 'relay_cell_connection_snapshot_freshness', + skipWhen: 'present' + }, + { kind: 'index', table: 'relay_cell_fences', name: 'relay_cell_fences_expiry', skipWhen: 'present' }, + { kind: 'index', table: 'relay_cell_committed_fences', name: 'relay_cell_committed_fences_expiry', skipWhen: 'present' }, + { + kind: 'index', + table: 'relay_cell_legacy_fence_adoptions', + name: 'relay_cell_legacy_fence_adoptions_expiry', + skipWhen: 'present' + }, + { kind: 'index', table: 'relay_cell_fence_attempts', name: 'relay_cell_fence_attempts_expiry', skipWhen: 'present' }, + { kind: 'index', table: 'relay_cell_fence_attempts', name: 'relay_cell_fence_attempts_cell', skipWhen: 'present' }, + { + kind: 'index', + table: 'relay_cell_fence_apply_invocations', + name: 'relay_cell_fence_apply_invocations_attempt', + skipWhen: 'present' + }, + { + kind: 'index', + table: 'relay_cell_drain_attempt_states', + name: 'relay_cell_drain_attempt_states_cell', + skipWhen: 'present' + }, + { + kind: 'index', + table: 'relay_assignment_activity_leases', + name: 'relay_assignment_activity_expiry', + skipWhen: 'present' + }, + { + kind: 'index', + table: 'relay_control_connection_reservations', + name: 'relay_control_connection_reservation_headroom', + skipWhen: 'present' + }, + { + kind: 'index', + table: 'relay_control_connection_reservations', + name: 'relay_control_connection_reservation_assignment', + skipWhen: 'present' + }, + { kind: 'index', table: 'relay_assignment_migrations', name: 'relay_assignment_migrations_active', skipWhen: 'present' }, + { + kind: 'index', + table: 'relay_post_drain_migration_pins', + name: 'relay_post_drain_migration_pins_attempt', + skipWhen: 'present' + }, + { kind: 'index', table: 'relay_audit_events', name: 'relay_audit_events_at', skipWhen: 'present' }, + { kind: 'column', table: 'relay_region_decisions', name: 'last_considered_at', skipWhen: 'present' }, + { kind: 'column', table: 'relay_region_decisions', name: 'cohort_bucket', skipWhen: 'present' }, + // Constraint swaps are matched by name in pg_constraint, with opposite polarities: nothing to + // drop is nothing to do, and a name already there is nothing to add. + { + kind: 'constraint', + table: 'relay_region_rehome_attempts', + name: 'relay_region_rehome_attempts_preferred_region_check', + skipWhen: 'absent' + }, + { + kind: 'constraint', + table: 'relay_region_rehome_attempts', + name: 'relay_region_rehome_attempts_preferred_region_valid', + skipWhen: 'present' + }, + { kind: 'column', table: 'relay_region_rehome_control', name: 'host_cooldown_ms', skipWhen: 'present' }, + { kind: 'column', table: 'relay_control_capabilities', name: 'idle_regional_rehome', skipWhen: 'present' }, + { kind: 'column', table: 'relay_region_rehome_attempts', name: 'source_generation', skipWhen: 'present' } +] + +const INDEX_OR_ADD_COLUMN = /^(?:CREATE\s+(?:UNIQUE\s+)?INDEX|ALTER\s+TABLE\s+[^\s]+\s+ADD\s+COLUMN)/i + +function lockTakingStatements(): string[] { + return relayPostgresSchemaStatements().filter(takesRelationLock) +} + +describe('relay boot-time lock targets', () => { + it('matches the pinned list of lock-taking statements', () => { + expect(lockTakingStatements().map(schemaLockTarget)).toEqual(GOLDEN_LOCK_TAKING) + }) + + it('derives a target for every CREATE INDEX and every ALTER TABLE ADD COLUMN', () => { + // A census over the real schema, not two hand-picked cases: a statement that lands here + // without a target is sent on every boot and takes the lock the pre-check exists to avoid. + // requireSchemaLockTarget is what boot calls, so this fails the same way boot would. + for (const statement of relayPostgresSchemaStatements()) { + expect(() => requireSchemaLockTarget(statement)).not.toThrow() + } + const unparsed = relayPostgresSchemaStatements().filter( + (statement) => + INDEX_OR_ADD_COLUMN.test(sqlWithoutComments(statement)) && + schemaLockTarget(statement) === undefined + ) + expect(unparsed).toEqual([]) + }) + + it('reads every derived name as a bare identifier, never a keyword or a qualified name', () => { + for (const statement of relayPostgresSchemaStatements()) { + const target = schemaLockTarget(statement) + if (!target) continue + expect(target.name).toMatch(/^[a-z_][a-z0-9_]*$/) + expect(target.table).toMatch(/^[a-z_][a-z0-9_]*$/) + } + }) + + it('pre-checks every lock-taking statement, with no exceptions', () => { + // The invariant the rule comment beside SCHEMA depends on: nothing that takes a relation lock + // reaches the server on a warm boot. A statement with no target breaks it. + const unchecked = lockTakingStatements().filter( + (statement) => schemaLockTarget(statement) === undefined + ) + expect(unchecked).toEqual([]) + expect(lockTakingStatements()).toHaveLength(GOLDEN_LOCK_TAKING.length) + }) + + it('derives a target through the comment block a split schema glues on', () => { + // Not vacuous: SCHEMA really does carry a comment-prefixed statement, and it is a CREATE INDEX + // on relay_connection_bases. Classifying the raw text would give it no target at all. + const commented = relayPostgresSchemaStatements().filter((statement) => + statement.startsWith('--') + ) + expect(commented.length).toBeGreaterThan(0) + for (const statement of commented) { + if (!takesRelationLock(statement)) continue + expect(schemaLockTarget(statement)).toBeDefined() + } + expect(commented.map(schemaLockTarget)).toContainEqual({ + kind: 'index', + table: 'relay_connection_bases', + name: 'relay_connection_bases_active_deadline', + skipWhen: 'present' + }) + }) + + it('leaves the dollar-quoted statement-stats migration byte-identical', () => { + // Its body is a PL/pgSQL block full of commas and parentheses. Reading the tag as anything but + // opaque would change the text classification sees, and it is the only such statement relay has. + const doBlock = relayPostgresSchemaStatements().find((statement) => statement.startsWith('DO ')) + expect(doBlock).toBeDefined() + expect(sqlWithoutComments(doBlock!)).toBe(doBlock) + expect(takesRelationLock(doBlock!)).toBe(false) + }) + + it('leaves every statement classifiable once its leading comments are stripped', () => { + for (const statement of relayPostgresSchemaStatements()) { + expect(sqlWithoutComments(statement)).toMatch(/^(?:CREATE|ALTER|DO)\s/i) + } + }) +}) diff --git a/cloud/packages/postgres-schema/package.json b/cloud/packages/postgres-schema/package.json index 86dea887cac..05cd4bd7221 100644 --- a/cloud/packages/postgres-schema/package.json +++ b/cloud/packages/postgres-schema/package.json @@ -9,7 +9,7 @@ "build": "tsc -p tsconfig.build.json", "clean": "node -e \"require('fs').rmSync('dist', { recursive: true, force: true })\"", "lint": "tsc -p tsconfig.json --noEmit", - "test": "pnpm build", + "test": "pnpm build && vitest run", "typecheck": "tsc -p tsconfig.json --noEmit" }, "devDependencies": { diff --git a/cloud/packages/postgres-schema/src/apply-postgres-schema.test.ts b/cloud/packages/postgres-schema/src/apply-postgres-schema.test.ts new file mode 100644 index 00000000000..b576e3a1129 --- /dev/null +++ b/cloud/packages/postgres-schema/src/apply-postgres-schema.test.ts @@ -0,0 +1,390 @@ +import { afterEach, describe, expect, it, vi } from 'vitest' +import { applyPostgresSchema } from './apply-postgres-schema.js' +import type { SchemaCatalogRow } from './catalog-object-precheck.js' + +function postgresError(code: string, constraint?: string): Error { + return Object.assign(new Error(code), constraint === undefined ? { code } : { code, constraint }) +} + +const COMMENTED_INDEX = `-- Why the sweep needs this index +CREATE INDEX IF NOT EXISTS relay_bases_active ON relay_connection_bases(active, deadline)` + +const COMMENTED_TABLE = `-- Two comment lines, the other shape a split schema carries +-- above a statement +CREATE TABLE IF NOT EXISTS relay_cells ( + cell_id TEXT PRIMARY KEY +)` + +// Answers absent first and present afterwards, the state a concurrent create leaves behind. +function catalogAnswersInSequence(answers: SchemaCatalogRow[][]): { + catalogQuery: (sql: string, params: unknown[]) => Promise + asked: unknown[][] +} { + const asked: unknown[][] = [] + return { + asked, + catalogQuery: async (sql, params) => { + asked.push([sql, ...params]) + return answers[asked.length - 1] ?? [] + } + } +} + +function catalogAnswers(rows: SchemaCatalogRow[]): { + catalogQuery: (sql: string, params: unknown[]) => Promise + asked: unknown[][] +} { + const asked: unknown[][] = [] + return { + asked, + catalogQuery: async (sql, params) => { + asked.push([sql, ...params]) + return rows + } + } +} + +afterEach(() => { + vi.restoreAllMocks() +}) + +describe('applyPostgresSchema classification', () => { + it('classifies a comment-prefixed CREATE INDEX by its first SQL keyword', async () => { + let calls = 0 + const query = vi.fn(async () => { + calls += 1 + if (calls === 1) throw postgresError('42P07') + return undefined + }) + await applyPostgresSchema([COMMENTED_INDEX], query, { wait: async () => undefined }) + expect(query).toHaveBeenCalledTimes(2) + }) + + it('classifies a comment-prefixed CREATE TABLE by its own collision codes', async () => { + // pg_type_typname_nsp_index is reached only through the CREATE TABLE branch, so a statement + // misread as unknown would fail the boot on a benign concurrent create instead of retrying. + let calls = 0 + const query = vi.fn(async () => { + calls += 1 + if (calls === 1) throw postgresError('23505', 'pg_type_typname_nsp_index') + return undefined + }) + await applyPostgresSchema([COMMENTED_TABLE], query, { wait: async () => undefined }) + expect(query).toHaveBeenCalledTimes(2) + }) + + it('retries a concurrent index collision until it succeeds', async () => { + let calls = 0 + const query = vi.fn(async () => { + calls += 1 + if (calls < 3) throw postgresError('23505', 'pg_class_relname_nsp_index') + return undefined + }) + const summary = await applyPostgresSchema(['CREATE INDEX IF NOT EXISTS i ON t(c)'], query, { + wait: async () => undefined + }) + expect(query).toHaveBeenCalledTimes(3) + expect(summary).toEqual({ ran: 1, skipped: 0 }) + }) + + it('treats an already-applied constraint as skipped rather than an error', async () => { + // Still the answer for a caller with no pre-check, and for a constraint another director + // committed between this boot's pre-check and its ALTER TABLE. + const query = vi.fn(async () => { + throw postgresError('42710') + }) + const summary = await applyPostgresSchema(['ALTER TABLE t ADD CONSTRAINT c CHECK (x > 0)'], query) + expect(summary).toEqual({ ran: 0, skipped: 1 }) + }) + + it('propagates an unrelated error without retrying', async () => { + const query = vi.fn(async () => { + throw postgresError('42501') + }) + await expect( + applyPostgresSchema(['CREATE TABLE IF NOT EXISTS t (id TEXT)'], query) + ).rejects.toThrow(/42501/) + expect(query).toHaveBeenCalledTimes(1) + }) +}) + +describe('applyPostgresSchema lock timeouts', () => { + it('does not retry a lock timeout', async () => { + vi.spyOn(console, 'error').mockImplementation(() => undefined) + const query = vi.fn(async () => { + throw postgresError('55P03') + }) + await expect( + applyPostgresSchema(['CREATE INDEX IF NOT EXISTS i ON t(c)'], query, { + wait: async () => undefined + }) + ).rejects.toThrow(/55P03/) + expect(query).toHaveBeenCalledTimes(1) + }) + + it('names the statement that could not take its lock', async () => { + const lines: string[] = [] + vi.spyOn(console, 'error').mockImplementation((line: string) => { + lines.push(line) + }) + const query = vi.fn(async () => { + throw postgresError('55P03') + }) + await expect(applyPostgresSchema([COMMENTED_INDEX], query)).rejects.toThrow(/55P03/) + expect(JSON.parse(lines[0] ?? '{}')).toMatchObject({ + event: 'orca_relay_postgres_schema_lock_timeout', + code: '55P03', + statement: 'CREATE INDEX IF NOT EXISTS relay_bases_active ON relay_connection_bases(active, deadline)' + }) + }) + + it('still retries a lock timeout for a caller that opts in', async () => { + // A caller with no catalog pre-check learns nothing from a lock timeout about whether the + // object exists, so its old bounded retry is the correct behaviour there. + vi.spyOn(console, 'warn').mockImplementation(() => undefined) + let calls = 0 + const query = vi.fn(async () => { + calls += 1 + if (calls < 3) throw postgresError('55P03') + return undefined + }) + await applyPostgresSchema(['CREATE INDEX IF NOT EXISTS i ON t(c)'], query, { + retryLockTimeout: true, + wait: async () => undefined + }) + expect(query).toHaveBeenCalledTimes(3) + }) +}) + +describe('applyPostgresSchema catalog pre-check', () => { + it('sends no lock-taking statement when the catalog has the object', async () => { + vi.spyOn(console, 'log').mockImplementation(() => undefined) + const query = vi.fn(async (_statement: string) => undefined) + const { catalogQuery, asked } = catalogAnswers([{ indisvalid: true }]) + const summary = await applyPostgresSchema([COMMENTED_TABLE, COMMENTED_INDEX], query, { + catalogQuery + }) + expect(query.mock.calls.map(([sql]) => sql)).toEqual([COMMENTED_TABLE]) + expect(asked).toEqual([ + [expect.stringContaining('pg_catalog.pg_index'), 'relay_connection_bases', 'relay_bases_active'] + ]) + expect(summary).toEqual({ ran: 1, skipped: 1 }) + }) + + it('skips an index the catalog reports as invalid rather than rebuilding it', async () => { + // A cancelled CREATE INDEX CONCURRENTLY leaves exactly this state, and IF NOT EXISTS skips it + // too, so reading indisvalid as a condition would newly take the lock it used to avoid. + const logged: { event?: string }[] = [] + vi.spyOn(console, 'log').mockImplementation((line: string) => { + logged.push(JSON.parse(line)) + }) + const query = vi.fn(async (_statement: string) => undefined) + const { catalogQuery } = catalogAnswers([{ indisvalid: false }]) + await applyPostgresSchema([COMMENTED_INDEX], query, { catalogQuery }) + expect(query).not.toHaveBeenCalled() + expect(logged.filter((entry) => entry.event?.endsWith('_object_present'))).toEqual([ + { + event: 'orca_relay_postgres_schema_object_present', + kind: 'index', + table: 'relay_connection_bases', + name: 'relay_bases_active', + indisvalid: false + } + ]) + }) + + it('reports how many statements ran and how many were skipped', async () => { + const logged: string[] = [] + vi.spyOn(console, 'log').mockImplementation((line: string) => { + logged.push(line) + }) + const { catalogQuery } = catalogAnswers([{ indisvalid: true }]) + await applyPostgresSchema([COMMENTED_TABLE, COMMENTED_INDEX], vi.fn(async () => undefined), { + catalogQuery, + eventPrefix: 'orca_push_postgres_schema' + }) + expect(JSON.parse(logged[logged.length - 1] ?? '{}')).toEqual({ + event: 'orca_push_postgres_schema_applied', + ran: 1, + skipped: 1 + }) + }) + + it('asks pg_attribute for a column and sends the ALTER TABLE when no row comes back', async () => { + const query = vi.fn(async (_statement: string) => undefined) + const { catalogQuery, asked } = catalogAnswers([]) + const statement = 'ALTER TABLE relay_control_capabilities ADD COLUMN IF NOT EXISTS idle BIGINT' + const summary = await applyPostgresSchema([statement], query, { catalogQuery }) + expect(asked).toEqual([ + [expect.stringContaining('pg_catalog.pg_attribute'), 'relay_control_capabilities', 'idle'] + ]) + expect(query.mock.calls.map(([sql]) => sql)).toEqual([statement]) + expect(summary).toEqual({ ran: 1, skipped: 0 }) + }) + + it('never probes the catalog for a statement that takes no relation lock', async () => { + const query = vi.fn(async (_statement: string) => undefined) + const { catalogQuery, asked } = catalogAnswers([{ indisvalid: true }]) + await applyPostgresSchema([COMMENTED_TABLE], query, { catalogQuery }) + expect(asked).toEqual([]) + expect(query).toHaveBeenCalledTimes(1) + }) + + it('skips an ADD CONSTRAINT the catalog already names', async () => { + vi.spyOn(console, 'log').mockImplementation(() => undefined) + const query = vi.fn(async (_statement: string) => undefined) + const { catalogQuery, asked } = catalogAnswers([{}]) + const summary = await applyPostgresSchema( + ['ALTER TABLE relay_region_rehome_attempts ADD CONSTRAINT region_valid CHECK (r IN (1))'], + query, + { catalogQuery } + ) + expect(asked).toEqual([ + [ + expect.stringContaining('pg_catalog.pg_constraint'), + 'relay_region_rehome_attempts', + 'region_valid' + ] + ]) + expect(query).not.toHaveBeenCalled() + expect(summary).toEqual({ ran: 0, skipped: 1 }) + }) + + it('skips a DROP CONSTRAINT IF EXISTS when the constraint is already gone', async () => { + // Inverse polarity: an absent constraint is what means there is nothing to drop. Sending it + // anyway takes ACCESS EXCLUSIVE to discover the same thing. + const logged: { event?: string }[] = [] + vi.spyOn(console, 'log').mockImplementation((line: string) => { + logged.push(JSON.parse(line)) + }) + const query = vi.fn(async (_statement: string) => undefined) + const { catalogQuery } = catalogAnswers([]) + const summary = await applyPostgresSchema( + ['ALTER TABLE relay_region_rehome_attempts DROP CONSTRAINT IF EXISTS region_check'], + query, + { catalogQuery } + ) + expect(query).not.toHaveBeenCalled() + expect(summary).toEqual({ ran: 0, skipped: 1 }) + expect(logged).toContainEqual({ + event: 'orca_relay_postgres_schema_object_absent', + kind: 'constraint', + table: 'relay_region_rehome_attempts', + name: 'region_check', + indisvalid: undefined + }) + }) + + it('sends a DROP CONSTRAINT IF EXISTS when the constraint is still there', async () => { + const query = vi.fn(async (_statement: string) => undefined) + const { catalogQuery } = catalogAnswers([{}]) + const statement = 'ALTER TABLE t DROP CONSTRAINT IF EXISTS region_check' + const summary = await applyPostgresSchema([statement], query, { catalogQuery }) + expect(query.mock.calls.map(([sql]) => sql)).toEqual([statement]) + expect(summary).toEqual({ ran: 1, skipped: 0 }) + }) + + it('sends an ADD CONSTRAINT the catalog does not name yet', async () => { + const query = vi.fn(async (_statement: string) => undefined) + const { catalogQuery } = catalogAnswers([]) + const statement = 'ALTER TABLE t ADD CONSTRAINT region_valid CHECK (r IN (1))' + const summary = await applyPostgresSchema([statement], query, { catalogQuery }) + expect(query.mock.calls.map(([sql]) => sql)).toEqual([statement]) + expect(summary).toEqual({ ran: 1, skipped: 0 }) + }) + + it('sends every statement when no catalog query is supplied', async () => { + const query = vi.fn(async (_statement: string) => undefined) + const summary = await applyPostgresSchema([COMMENTED_TABLE, COMMENTED_INDEX], query) + expect(query).toHaveBeenCalledTimes(2) + expect(summary).toEqual({ ran: 2, skipped: 0 }) + }) +}) + +describe('applyPostgresSchema concurrent creates', () => { + it('re-asks the catalog on a collision instead of retrying the CREATE INDEX', async () => { + // Another director created the index between the pre-check and this statement. Retrying would + // take SHARE on the table again for an object that is already there. + vi.spyOn(console, 'log').mockImplementation(() => undefined) + const query = vi.fn(async (_statement: string) => { + throw postgresError('42P07') + }) + const { catalogQuery, asked } = catalogAnswersInSequence([[], [{ indisvalid: true }]]) + const summary = await applyPostgresSchema([COMMENTED_INDEX], query, { + catalogQuery, + wait: async () => undefined + }) + expect(query).toHaveBeenCalledTimes(1) + expect(asked).toHaveLength(2) + expect(summary).toEqual({ ran: 0, skipped: 1 }) + }) + + it('still retries when the catalog says the object is not there after all', async () => { + let calls = 0 + const query = vi.fn(async (_statement: string) => { + calls += 1 + if (calls === 1) throw postgresError('23505', 'pg_class_relname_nsp_index') + return undefined + }) + const { catalogQuery } = catalogAnswersInSequence([[], []]) + const summary = await applyPostgresSchema([COMMENTED_INDEX], query, { + catalogQuery, + wait: async () => undefined + }) + expect(query).toHaveBeenCalledTimes(2) + expect(summary).toEqual({ ran: 1, skipped: 0 }) + }) + + it('retries a CREATE TABLE collision without a catalog re-ask, having no target to ask about', async () => { + let calls = 0 + const query = vi.fn(async (_statement: string) => { + calls += 1 + if (calls === 1) throw postgresError('42710') + return undefined + }) + const { catalogQuery, asked } = catalogAnswersInSequence([[], []]) + await applyPostgresSchema([COMMENTED_TABLE], query, { + catalogQuery, + wait: async () => undefined + }) + expect(asked).toEqual([]) + expect(query).toHaveBeenCalledTimes(2) + }) +}) + +describe('applyPostgresSchema unparseable statements', () => { + it('fails the boot rather than sending an index whose target cannot be read', async () => { + const query = vi.fn(async (_statement: string) => undefined) + await expect(applyPostgresSchema(['CREATE INDEX ON t(c)'], query)).rejects.toThrow( + /unparsed_schema_lock_target/ + ) + expect(query).not.toHaveBeenCalled() + }) + + it('fails even with no catalog query, because the statement would take the lock either way', async () => { + const query = vi.fn(async (_statement: string) => undefined) + await expect( + applyPostgresSchema(['ALTER TABLE t ADD COLUMN IF NOT EXISTS'], query) + ).rejects.toThrow(/unparsed_schema_lock_target/) + expect(query).not.toHaveBeenCalled() + }) +}) + +describe('applyPostgresSchema statement text', () => { + it('sends the original statement, comments included, not the classified form', async () => { + // Classification reads a comment-free copy. Rewriting what the server runs would change the + // DDL itself, and a comment inside a string literal or a quoted name is part of the statement. + const statement = `ALTER TABLE t ADD /* note */ COLUMN c TEXT DEFAULT '-- keep'` + const query = vi.fn(async (_sql: string) => undefined) + const { catalogQuery, asked } = catalogAnswers([]) + await applyPostgresSchema([statement], query, { catalogQuery }) + expect(query.mock.calls.map(([sql]) => sql)).toEqual([statement]) + expect(asked).toEqual([[expect.stringContaining('pg_catalog.pg_attribute'), 't', 'c']]) + }) + + it('sends a comment-prefixed statement unchanged too', async () => { + const query = vi.fn(async (_sql: string) => undefined) + await applyPostgresSchema([COMMENTED_TABLE], query) + expect(query.mock.calls.map(([sql]) => sql)).toEqual([COMMENTED_TABLE]) + }) +}) diff --git a/cloud/packages/postgres-schema/src/apply-postgres-schema.ts b/cloud/packages/postgres-schema/src/apply-postgres-schema.ts new file mode 100644 index 00000000000..3417d9ccd70 --- /dev/null +++ b/cloud/packages/postgres-schema/src/apply-postgres-schema.ts @@ -0,0 +1,179 @@ +import { catalogObjectPresence, type SchemaCatalogQuery } from './catalog-object-precheck.js' +import { + requireSchemaLockTarget, + sqlWithoutComments, + type SchemaLockTarget +} from './schema-lock-target.js' + +const RETRYABLE_SCHEMA_CODES = new Set(['57014']) +const LOCK_NOT_AVAILABLE = '55P03' +const DEFAULT_RETRY_DEADLINE_MS = 30_000 +const RETRY_BASE_DELAY_MS = 250 +const RETRY_MAX_DELAY_MS = 2_000 +const DEFAULT_EVENT_PREFIX = 'orca_relay_postgres_schema' + +export type SchemaStartupOptions = { + // Enables the catalog pre-check. Without it every lock-taking statement is sent as before. + catalogQuery?: SchemaCatalogQuery + eventPrefix?: string + now?: () => number + random?: () => number + retryDeadlineMs?: number + // Only for a caller with no catalog pre-check, where a lock timeout still says nothing about + // whether the object exists. + retryLockTimeout?: boolean + wait?: (delayMs: number) => Promise +} + +export type SchemaApplySummary = { ran: number; skipped: number } + +function retryDelayMs(attempt: number, random: () => number): number { + const ceiling = Math.min(RETRY_BASE_DELAY_MS * 2 ** (attempt - 1), RETRY_MAX_DELAY_MS) + return Math.ceil(ceiling * (0.5 + random() * 0.5)) +} + +function wait(delayMs: number): Promise { + return new Promise((resolve) => setTimeout(resolve, delayMs)) +} + +const CREATE_TABLE_IF_NOT_EXISTS = /^CREATE\s+TABLE\s+IF\s+NOT\s+EXISTS\b/i +const CREATE_INDEX_IF_NOT_EXISTS = /^CREATE\s+(?:UNIQUE\s+)?INDEX\s+IF\s+NOT\s+EXISTS\b/i +const ALTER_TABLE_ADD_CONSTRAINT = /^ALTER\s+TABLE\s+\S+\s+ADD\s+CONSTRAINT\b/i + +// `IF NOT EXISTS` only checks the name before the catalog inserts, so the loser of a concurrent +// CREATE can fail on the catalog unique index (23505) or, when the winner has already committed by +// the time the loser reaches TypeCreate/heap_create_with_catalog, on the name check those routines +// repeat (42710 duplicate type, 42P07 duplicate relation). Each is a no-op on the next attempt. +function concurrentCreateCollision( + value: { code?: unknown; constraint?: unknown }, + sql: string +): boolean { + if (CREATE_TABLE_IF_NOT_EXISTS.test(sql)) { + return ( + (value.code === '23505' && value.constraint === 'pg_type_typname_nsp_index') || + value.code === '42710' || + value.code === '42P07' + ) + } + if (CREATE_INDEX_IF_NOT_EXISTS.test(sql)) { + return ( + (value.code === '23505' && value.constraint === 'pg_class_relname_nsp_index') || + value.code === '42P07' + ) + } + return false +} + +function constraintAlreadyApplied(error: unknown, sql: string): boolean { + return ( + ALTER_TABLE_ADD_CONSTRAINT.test(sql) && (error as { code?: unknown } | null)?.code === '42710' + ) +} + +function retryableSchemaError(error: unknown, sql: string): boolean { + const value = (error as { code?: unknown; constraint?: unknown } | null) ?? {} + return RETRYABLE_SCHEMA_CODES.has(String(value.code)) || concurrentCreateCollision(value, sql) +} + +// Evaluated immediately before each statement, so a pre-check still sees the objects the statements +// ahead of it created in this same boot. +async function nothingToDo( + target: SchemaLockTarget | undefined, + options: SchemaStartupOptions, + eventPrefix: string +): Promise { + const catalogQuery = options.catalogQuery + if (!catalogQuery || !target) return false + const presence = await catalogObjectPresence(catalogQuery, target) + if (presence.present !== (target.skipWhen === 'present')) return false + console.log( + JSON.stringify({ + event: `${eventPrefix}_object_${target.skipWhen}`, + kind: target.kind, + table: target.table, + name: target.name, + indisvalid: presence.indisvalid + }) + ) + return true +} + +export async function applyPostgresSchema( + statements: string[], + query: (statement: string) => Promise, + options: SchemaStartupOptions = {} +): Promise { + const eventPrefix = options.eventPrefix ?? DEFAULT_EVENT_PREFIX + const now = options.now ?? Date.now + const random = options.random ?? Math.random + const pause = options.wait ?? wait + const deadlineAt = now() + (options.retryDeadlineMs ?? DEFAULT_RETRY_DEADLINE_MS) + const summary: SchemaApplySummary = { ran: 0, skipped: 0 } + + for (const statement of statements) { + // Throws when an index or column statement's target cannot be read, rather than sending it + // unchecked into the lock queue. + const target = requireSchemaLockTarget(statement) + if (await nothingToDo(target, options, eventPrefix)) { + summary.skipped += 1 + continue + } + const sql = sqlWithoutComments(statement) + let attempt = 1 + while (true) { + try { + await query(statement) + summary.ran += 1 + break + } catch (error) { + if (constraintAlreadyApplied(error, sql)) { + summary.skipped += 1 + break + } + const code = String((error as { code?: unknown } | null)?.code) + // With the pre-check ahead of it a lock timeout means the object is genuinely missing and + // this boot lost the queue. Relation locks are granted in queue order, so each retry parks + // every writer behind it again for another timeout. Fail once, loudly. + if (code === LOCK_NOT_AVAILABLE && !options.retryLockTimeout) { + console.error( + JSON.stringify({ + event: `${eventPrefix}_lock_timeout`, + code, + statement: sql.split('\n')[0], + detail: 'boot-time DDL could not take its lock; retrying would requeue every writer' + }) + ) + throw error + } + // The object was created between the pre-check and this statement. Re-asking the catalog + // is the cheap answer; retrying the CREATE INDEX would take SHARE on the table again for + // an object that is already there. + if ( + concurrentCreateCollision((error as { code?: unknown; constraint?: unknown }) ?? {}, sql) && + (await nothingToDo(target, options, eventPrefix)) + ) { + summary.skipped += 1 + break + } + const remainingMs = deadlineAt - now() + const retryable = + retryableSchemaError(error, sql) || + (code === LOCK_NOT_AVAILABLE && options.retryLockTimeout === true) + if (!retryable || remainingMs <= 0) { + if (retryable) { + console.warn( + JSON.stringify({ event: `${eventPrefix}_retry_exhausted`, code, attempts: attempt }) + ) + } + throw error + } + const delayMs = Math.min(remainingMs, retryDelayMs(attempt, random)) + console.warn(JSON.stringify({ event: `${eventPrefix}_retry`, code, attempt, delayMs })) + await pause(delayMs) + attempt += 1 + } + } + } + console.log(JSON.stringify({ event: `${eventPrefix}_applied`, ...summary })) + return summary +} diff --git a/cloud/packages/postgres-schema/src/catalog-object-precheck.ts b/cloud/packages/postgres-schema/src/catalog-object-precheck.ts new file mode 100644 index 00000000000..66bc5848608 --- /dev/null +++ b/cloud/packages/postgres-schema/src/catalog-object-precheck.ts @@ -0,0 +1,48 @@ +import type { SchemaLockTarget } from './schema-lock-target.js' + +export type SchemaCatalogRow = Record + +// Runs with `$n` placeholders bound to the lock target, on the same connection the DDL would use. +export type SchemaCatalogQuery = ( + sql: string, + params: unknown[] +) => Promise + +// The index name is resolved inside the table's own namespace, and `i.indrelid = t.oid` ties it to +// this table: index names are unique per schema, not per table, so without that condition a +// same-named index on a sibling table answers yes and the real index is skipped forever. +// `to_regclass` returns NULL rather than erroring when the table does not exist yet, which is the +// whole of a cold start. +const INDEX_PRESENT = `SELECT i.indisvalid FROM pg_catalog.pg_class t +JOIN pg_catalog.pg_class c ON c.relnamespace = t.relnamespace AND c.relname = $2 +JOIN pg_catalog.pg_index i ON i.indexrelid = c.oid AND i.indrelid = t.oid +WHERE t.oid = to_regclass($1)` + +const COLUMN_PRESENT = `SELECT 1 FROM pg_catalog.pg_attribute +WHERE attrelid = to_regclass($1) AND attname = $2 AND attnum > 0 AND NOT attisdropped` + +// Name only. The CHECK body is generated from RELAY_REGIONS, so comparing it would re-run the swap +// on every region change, and an ADD CONSTRAINT is the one statement here that scans the table. +const CONSTRAINT_PRESENT = `SELECT 1 FROM pg_catalog.pg_constraint +WHERE conrelid = to_regclass($1) AND conname = $2` + +const PRESENCE_SQL = { + index: INDEX_PRESENT, + column: COLUMN_PRESENT, + constraint: CONSTRAINT_PRESENT +} as const + +export type SchemaCatalogPresence = { present: boolean; indisvalid: unknown } + +// Row presence is the answer, whatever the row says. An index left invalid by a cancelled +// concurrent build is skipped by `IF NOT EXISTS` today as well, so reading `indisvalid` as a +// condition would newly take the lock for exactly the indexes a failed build left behind. +export async function catalogObjectPresence( + query: SchemaCatalogQuery, + target: SchemaLockTarget +): Promise { + const sql = PRESENCE_SQL[target.kind] + const rows = await query(sql, [target.table, target.name]) + const row = rows[0] + return row ? { present: true, indisvalid: row.indisvalid } : { present: false, indisvalid: undefined } +} diff --git a/cloud/packages/postgres-schema/src/index.ts b/cloud/packages/postgres-schema/src/index.ts index 9981a03ead5..af101ade9b7 100644 --- a/cloud/packages/postgres-schema/src/index.ts +++ b/cloud/packages/postgres-schema/src/index.ts @@ -1,113 +1,18 @@ -const RETRYABLE_SCHEMA_CODES = new Set(['55P03', '57014']) -const DEFAULT_RETRY_DEADLINE_MS = 30_000 -const RETRY_BASE_DELAY_MS = 250 -const RETRY_MAX_DELAY_MS = 2_000 - -type SchemaStartupOptions = { - eventPrefix?: string - now?: () => number - random?: () => number - retryDeadlineMs?: number - wait?: (delayMs: number) => Promise -} - -function retryDelayMs(attempt: number, random: () => number): number { - const ceiling = Math.min(RETRY_BASE_DELAY_MS * 2 ** (attempt - 1), RETRY_MAX_DELAY_MS) - return Math.ceil(ceiling * (0.5 + random() * 0.5)) -} - -function wait(delayMs: number): Promise { - return new Promise((resolve) => setTimeout(resolve, delayMs)) -} - -const CREATE_TABLE_IF_NOT_EXISTS = /^\s*CREATE\s+TABLE\s+IF\s+NOT\s+EXISTS\b/i -const CREATE_INDEX_IF_NOT_EXISTS = /^\s*CREATE\s+(?:UNIQUE\s+)?INDEX\s+IF\s+NOT\s+EXISTS\b/i - -// `IF NOT EXISTS` only checks the name before the catalog inserts, so the loser of a concurrent -// CREATE can fail on the catalog unique index (23505) or, when the winner has already committed by -// the time the loser reaches TypeCreate/heap_create_with_catalog, on the name check those routines -// repeat (42710 duplicate type, 42P07 duplicate relation). Each is a no-op on the next attempt. -function concurrentCreateCollision( - value: { code?: unknown; constraint?: unknown }, - statement: string -): boolean { - if (CREATE_TABLE_IF_NOT_EXISTS.test(statement)) { - return ( - (value.code === '23505' && value.constraint === 'pg_type_typname_nsp_index') || - value.code === '42710' || - value.code === '42P07' - ) - } - if (CREATE_INDEX_IF_NOT_EXISTS.test(statement)) { - return ( - (value.code === '23505' && value.constraint === 'pg_class_relname_nsp_index') || - value.code === '42P07' - ) - } - return false -} - -const ALTER_TABLE_ADD_CONSTRAINT = /^\s*ALTER\s+TABLE\s+\S+\s+ADD\s+CONSTRAINT\b/i - -function constraintAlreadyApplied(error: unknown, statement: string): boolean { - return ( - ALTER_TABLE_ADD_CONSTRAINT.test(statement) && - (error as { code?: unknown }).code === '42710' - ) -} - -function retryableSchemaError(error: unknown, statement: string): boolean { - const value = error as { code?: unknown; constraint?: unknown } - return ( - RETRYABLE_SCHEMA_CODES.has(String(value.code)) || concurrentCreateCollision(value, statement) - ) -} - -export async function applyPostgresSchema( - statements: string[], - query: (statement: string) => Promise, - options: SchemaStartupOptions = {} -): Promise { - const now = options.now ?? Date.now - const random = options.random ?? Math.random - const pause = options.wait ?? wait - const deadlineAt = now() + (options.retryDeadlineMs ?? DEFAULT_RETRY_DEADLINE_MS) - - for (const statement of statements) { - let attempt = 1 - while (true) { - try { - await query(statement) - break - } catch (error) { - if (constraintAlreadyApplied(error, statement)) break - const code = String((error as { code?: unknown }).code) - const remainingMs = deadlineAt - now() - const retryable = retryableSchemaError(error, statement) - if (!retryable || remainingMs <= 0) { - if (retryable) { - console.warn( - JSON.stringify({ - event: `${options.eventPrefix ?? 'orca_relay_postgres_schema'}_retry_exhausted`, - code, - attempts: attempt - }) - ) - } - throw error - } - const delayMs = Math.min(remainingMs, retryDelayMs(attempt, random)) - console.warn( - JSON.stringify({ - event: `${options.eventPrefix ?? 'orca_relay_postgres_schema'}_retry`, - code, - attempt, - delayMs - }) - ) - await pause(delayMs) - attempt += 1 - } - } - } -} +export { + applyPostgresSchema, + type SchemaApplySummary, + type SchemaStartupOptions +} from './apply-postgres-schema.js' +export { + catalogObjectPresence, + type SchemaCatalogPresence, + type SchemaCatalogQuery, + type SchemaCatalogRow +} from './catalog-object-precheck.js' +export { + requireSchemaLockTarget, + schemaLockTarget, + sqlWithoutComments, + takesRelationLock, + type SchemaLockTarget +} from './schema-lock-target.js' diff --git a/cloud/packages/postgres-schema/src/schema-lock-target.test.ts b/cloud/packages/postgres-schema/src/schema-lock-target.test.ts new file mode 100644 index 00000000000..ffd04dd81cd --- /dev/null +++ b/cloud/packages/postgres-schema/src/schema-lock-target.test.ts @@ -0,0 +1,419 @@ +import { describe, expect, it } from 'vitest' +import { + requireSchemaLockTarget, + schemaLockTarget, + sqlWithoutComments, + takesRelationLock +} from './schema-lock-target.js' + +// The shape a schema string split on ';' actually produces: the comment written above a statement +// arrives glued to the front of it. +const COMMENTED_INDEX = `-- Why: the maintenance sweep matches (active, deadline) while inactive +-- bases accumulate unboundedly. +CREATE INDEX IF NOT EXISTS relay_connection_bases_active_deadline + ON relay_connection_bases(active, deadline)` + +const COMMENTED_TABLE = `-- Rehoming is bidirectional, but tables created before that carry the +-- original single-region column check. +CREATE TABLE IF NOT EXISTS relay_cells ( + cell_id TEXT PRIMARY KEY +)` + +describe('sqlWithoutComments', () => { + it('strips the line comments a split schema glues above a statement', () => { + expect(sqlWithoutComments(COMMENTED_INDEX)).toMatch(/^CREATE INDEX IF NOT EXISTS/) + }) + + it('strips a leading block comment', () => { + expect(sqlWithoutComments('/* note */\n ALTER TABLE t ADD COLUMN c TEXT')).toBe( + 'ALTER TABLE t ADD COLUMN c TEXT' + ) + }) + + it('strips a comment sitting between two keywords', () => { + expect(sqlWithoutComments('ALTER TABLE t ADD /* note */ COLUMN c TEXT')).toBe( + 'ALTER TABLE t ADD COLUMN c TEXT' + ) + }) + + it('strips a trailing comment', () => { + expect(sqlWithoutComments('SELECT 1 -- note')).toBe('SELECT 1') + }) + + it('closes the inner block comment first when they nest', () => { + expect(sqlWithoutComments('ALTER TABLE t /* a /* b */ c */ ADD COLUMN d TEXT')).toBe( + 'ALTER TABLE t ADD COLUMN d TEXT' + ) + }) + + it('leaves a comment marker inside a string literal alone', () => { + expect(sqlWithoutComments(`ALTER TABLE t ADD COLUMN c TEXT DEFAULT '-- not a comment'`)).toBe( + `ALTER TABLE t ADD COLUMN c TEXT DEFAULT '-- not a comment'` + ) + }) + + it('leaves a comment marker inside a quoted identifier alone', () => { + expect(sqlWithoutComments('ALTER TABLE t ADD COLUMN "a/* b */c" TEXT')).toBe( + 'ALTER TABLE t ADD COLUMN "a/* b */c" TEXT' + ) + }) +}) + +describe('comments between keywords', () => { + // Before this, the classification regexes and the must-parse shapes both needed `ADD COLUMN` + // contiguous, so this statement derived NO target and threw NOTHING: the DDL ran with no + // pre-check, taking ACCESS EXCLUSIVE on every boot. + it('derives a column target through a comment between ADD and COLUMN', () => { + expect(requireSchemaLockTarget('ALTER TABLE t ADD /* note */ COLUMN c TEXT')).toEqual({ + kind: 'column', + table: 't', + name: 'c', + skipWhen: 'present' + }) + }) + + it('derives an index target through a line comment before ON', () => { + expect( + requireSchemaLockTarget('CREATE INDEX IF NOT EXISTS i\n-- why this index exists\nON t(c)') + ).toEqual({ kind: 'index', table: 't', name: 'i', skipWhen: 'present' }) + }) + + it('still counts a commented statement as taking a relation lock', () => { + expect(takesRelationLock('/* note */ ALTER TABLE t ADD COLUMN c TEXT')).toBe(true) + }) + + it('does not read a comment marker inside a quoted name as a comment', () => { + expect(schemaLockTarget('ALTER TABLE t ADD COLUMN "a--b" TEXT')).toEqual({ + kind: 'column', + table: 't', + name: 'a--b', + skipWhen: 'present' + }) + }) +}) + +describe('catalog name folding', () => { + it('reads a quoted identifier containing a dot as one name', () => { + expect(schemaLockTarget('CREATE INDEX IF NOT EXISTS "a.b" ON t(c)')).toEqual({ + kind: 'index', + table: 't', + name: 'a.b', + skipWhen: 'present' + }) + }) + + it('folds an unquoted name to lower case, the form the catalog stores', () => { + expect(schemaLockTarget('CREATE INDEX IF NOT EXISTS Foo ON Bar(c)')).toEqual({ + kind: 'index', + table: 'Bar', + name: 'foo', + skipWhen: 'present' + }) + }) + + it('keeps a quoted name in its written case', () => { + expect(schemaLockTarget('CREATE INDEX IF NOT EXISTS public."Mixed.Name" ON t(c)')).toEqual({ + kind: 'index', + table: 't', + name: 'Mixed.Name', + skipWhen: 'present' + }) + }) + + it('unescapes a doubled quote and leaves the qualified table text as written', () => { + expect(schemaLockTarget('ALTER TABLE App."My Table" ADD COLUMN "od""d" TEXT')).toEqual({ + kind: 'column', + table: 'App."My Table"', + name: 'od"d', + skipWhen: 'present' + }) + }) + + it('folds an unquoted column name too', () => { + expect(schemaLockTarget('ALTER TABLE t ADD COLUMN IF NOT EXISTS HostCooldownMs BIGINT')).toEqual( + { kind: 'column', table: 't', name: 'hostcooldownms', skipWhen: 'present' } + ) + }) +}) + +describe('square brackets in an ALTER TABLE', () => { + it.each([ + ['an array type', 'ALTER TABLE t ADD COLUMN c bigint[] DEFAULT ARRAY[1, 2]'], + ['a nested array default', "ALTER TABLE t ADD COLUMN c TEXT[] DEFAULT ARRAY['a', 'b']"] + ])('does not read a comma inside %s as a second subcommand', (_label, statement) => { + expect(() => requireSchemaLockTarget(statement)).not.toThrow() + }) + + it('still catches a second subcommand after an array default', () => { + expect(() => + requireSchemaLockTarget('ALTER TABLE t ADD COLUMN a bigint[] DEFAULT ARRAY[1, 2], ADD COLUMN b TEXT') + ).toThrow(/unparsed_schema_lock_target/) + }) +}) + + +describe('takesRelationLock', () => { + it('classifies a comment-prefixed CREATE INDEX by its first SQL keyword', () => { + expect(takesRelationLock(COMMENTED_INDEX)).toBe(true) + }) + + it('classifies a comment-prefixed CREATE TABLE as taking no relation lock', () => { + expect(takesRelationLock(COMMENTED_TABLE)).toBe(false) + }) + + it('counts every ALTER TABLE, including the constraint swaps', () => { + expect(takesRelationLock('ALTER TABLE t DROP CONSTRAINT IF EXISTS c')).toBe(true) + }) +}) + +describe('schemaLockTarget', () => { + it('derives an index target through the comments above it', () => { + expect(schemaLockTarget(COMMENTED_INDEX)).toEqual({ + kind: 'index', + table: 'relay_connection_bases', + name: 'relay_connection_bases_active_deadline', + skipWhen: 'present' + }) + }) + + it('derives an index target across the line break before ON', () => { + expect( + schemaLockTarget(`CREATE INDEX IF NOT EXISTS relay_reservation_assignment + ON relay_control_connection_reservations( + user_id, relay_host_id + )`) + ).toEqual({ + kind: 'index', + table: 'relay_control_connection_reservations', + name: 'relay_reservation_assignment', + skipWhen: 'present' + }) + }) + + it('derives a unique concurrent index target', () => { + expect(schemaLockTarget('CREATE UNIQUE INDEX CONCURRENTLY i ON t(c)')).toEqual({ + kind: 'index', + table: 't', + name: 'i', + skipWhen: 'present' + }) + }) + + it('keeps the schema qualification on the table and drops it from the object name', () => { + // `table` is fed to to_regclass, which needs the qualification; `name` is matched against + // relname, which stores the bare identifier. + expect(schemaLockTarget('CREATE INDEX IF NOT EXISTS app.i ON app.t(c)')).toEqual({ + kind: 'index', + table: 'app.t', + name: 'i', + skipWhen: 'present' + }) + }) + + it('unquotes a quoted identifier, doubled quote included', () => { + expect(schemaLockTarget('CREATE INDEX IF NOT EXISTS "od""d" ON "My Table"(c)')).toEqual({ + kind: 'index', + table: '"My Table"', + name: 'od"d', + skipWhen: 'present' + }) + }) + + it('derives a column target from a multi-line ADD COLUMN IF NOT EXISTS', () => { + expect( + schemaLockTarget(`ALTER TABLE relay_region_rehome_control + ADD COLUMN IF NOT EXISTS host_cooldown_ms BIGINT NOT NULL + DEFAULT 604800000`) + ).toEqual({ + kind: 'column', + table: 'relay_region_rehome_control', + name: 'host_cooldown_ms', + skipWhen: 'present' + }) + }) + + it('derives a column target without IF NOT EXISTS', () => { + expect(schemaLockTarget('ALTER TABLE ONLY t ADD COLUMN c TEXT')).toEqual({ + kind: 'column', + table: 't', + name: 'c', + skipWhen: 'present' + }) + }) + + it('skips an ADD CONSTRAINT once the constraint name is there', () => { + expect(schemaLockTarget('ALTER TABLE t ADD CONSTRAINT c CHECK (x > 0)')).toEqual({ + kind: 'constraint', + table: 't', + name: 'c', + skipWhen: 'present' + }) + }) + + it('skips a DROP CONSTRAINT IF EXISTS when the constraint is already gone', () => { + // The inverse polarity: nothing to drop is nothing to do. + expect(schemaLockTarget('ALTER TABLE t DROP CONSTRAINT IF EXISTS c')).toEqual({ + kind: 'constraint', + table: 't', + name: 'c', + skipWhen: 'absent' + }) + }) + + it('derives a constraint target across a line break', () => { + expect( + schemaLockTarget(`ALTER TABLE relay_region_rehome_attempts + ADD CONSTRAINT relay_region_rehome_attempts_preferred_region_valid + CHECK (preferred_region IN ('us-central1'))`) + ).toEqual({ + kind: 'constraint', + table: 'relay_region_rehome_attempts', + name: 'relay_region_rehome_attempts_preferred_region_valid', + skipWhen: 'present' + }) + }) + + it('gives CREATE TABLE IF NOT EXISTS no target', () => { + expect(schemaLockTarget(COMMENTED_TABLE)).toBeUndefined() + }) +}) + +// Each of these reads as an index or column statement and each fails to yield a target. Letting any +// of them through would send an unchecked lock-taking statement on every boot. +const MALFORMED = [ + ['an index with no ON clause', 'CREATE INDEX IF NOT EXISTS i'], + ['an auto-named index', 'CREATE INDEX ON t(c)'], + ['an auto-named unique concurrent index', 'CREATE UNIQUE INDEX CONCURRENTLY ON t(c)'], + ['an index whose name ran into a comment', '-- note\nCREATE INDEX IF NOT EXISTS\nON t(c)'], + ['an ALTER TABLE with no table', 'ALTER TABLE ADD COLUMN c TEXT'], + ['an ADD COLUMN with no column', 'ALTER TABLE t ADD COLUMN'], + ['an ADD COLUMN IF NOT EXISTS with no column', 'ALTER TABLE t ADD COLUMN IF NOT EXISTS'] +] as const + +describe('requireSchemaLockTarget', () => { + it.each(MALFORMED)('throws with the statement text on %s', (_label, statement) => { + expect(() => requireSchemaLockTarget(statement)).toThrow(/unparsed_schema_lock_target/) + }) + + it('names the offending statement in the error', () => { + expect(() => requireSchemaLockTarget('CREATE INDEX ON t(c)')).toThrow( + 'unparsed_schema_lock_target: CREATE INDEX ON t(c)' + ) + }) + + it('returns the target for a statement that parses', () => { + expect(requireSchemaLockTarget(COMMENTED_INDEX)).toEqual({ + kind: 'index', + table: 'relay_connection_bases', + name: 'relay_connection_bases_active_deadline', + skipWhen: 'present' + }) + }) + + it.each([ + ['CREATE TABLE IF NOT EXISTS t (id TEXT)'], + ['ALTER TABLE t ALTER COLUMN c SET DEFAULT 0'] + ])('leaves %s alone, because no target is expected of it', (statement) => { + expect(requireSchemaLockTarget(statement)).toBeUndefined() + }) +}) + +describe('multi-action ALTER TABLE', () => { + it('throws rather than deriving only the first subcommand', () => { + // Deriving `a` and skipping on it would drop `b` for the life of the database, and the first + // subcommand parses fine, so nothing else here would catch it. + const statement = + 'ALTER TABLE t ADD COLUMN IF NOT EXISTS a TEXT, ADD COLUMN IF NOT EXISTS b TEXT' + expect(schemaLockTarget(statement)).toEqual({ + kind: 'column', + table: 't', + name: 'a', + skipWhen: 'present' + }) + expect(() => requireSchemaLockTarget(statement)).toThrow(/unparsed_schema_lock_target/) + }) + + it('throws on a constraint swap written as one statement', () => { + expect(() => + requireSchemaLockTarget( + 'ALTER TABLE t DROP CONSTRAINT IF EXISTS old, ADD CONSTRAINT new CHECK (x > 0)' + ) + ).toThrow(/unparsed_schema_lock_target/) + }) + + it.each([ + ['a parenthesised type', 'ALTER TABLE t ADD COLUMN IF NOT EXISTS a NUMERIC(10, 2)'], + ['a CHECK body', "ALTER TABLE t ADD CONSTRAINT c CHECK (r IN ('us-central1', 'asia-east2'))"], + ['a quoted comma', `ALTER TABLE t ADD COLUMN IF NOT EXISTS a TEXT DEFAULT 'x, y'`], + ['a doubled quote before a comma', `ALTER TABLE t ADD COLUMN a TEXT DEFAULT 'it''s, fine'`], + ['a trailing line comment', 'ALTER TABLE t ADD COLUMN a TEXT -- one, two'], + ['a trailing block comment', 'ALTER TABLE t ADD COLUMN a TEXT /* one, two */'] + ])('does not throw on %s', (_label, statement) => { + expect(() => requireSchemaLockTarget(statement)).not.toThrow() + }) + + it('derives through a block comment sitting where the column name belongs', () => { + expect(requireSchemaLockTarget('ALTER TABLE t ADD COLUMN /* note */ a TEXT')).toEqual({ + kind: 'column', + table: 't', + name: 'a', + skipWhen: 'present' + }) + }) + + it('leaves a multi-column CREATE INDEX alone', () => { + expect(() => requireSchemaLockTarget('CREATE INDEX IF NOT EXISTS i ON t(a, b)')).not.toThrow() + }) +}) + +describe('dollar-quoted bodies', () => { + it('does not read a comment marker inside a dollar-quoted default as a comment', () => { + expect(sqlWithoutComments('ALTER TABLE t ADD COLUMN c TEXT DEFAULT $$--$$')).toBe( + 'ALTER TABLE t ADD COLUMN c TEXT DEFAULT $$--$$' + ) + expect(requireSchemaLockTarget('ALTER TABLE t ADD COLUMN c TEXT DEFAULT $$--$$')).toEqual({ + kind: 'column', + table: 't', + name: 'c', + skipWhen: 'present' + }) + }) + + it('does not count a comma inside a dollar-quoted default as a second subcommand', () => { + expect(() => + requireSchemaLockTarget('ALTER TABLE t ADD COLUMN c TEXT DEFAULT $$a, b$$') + ).not.toThrow() + }) + + it('reads an inner $$ inside a tagged body as text, not as the close', () => { + // The closing delimiter has to match the opening tag, so the comma and the comment marker + // between the inner $$ pair are still inside the body. + const statement = 'ALTER TABLE t ADD COLUMN c TEXT DEFAULT $tag$ a $$ -- b, c $$ d $tag$' + expect(sqlWithoutComments(statement)).toBe(statement) + expect(() => requireSchemaLockTarget(statement)).not.toThrow() + expect(schemaLockTarget(statement)).toEqual({ + kind: 'column', + table: 't', + name: 'c', + skipWhen: 'present' + }) + }) + + it('still catches a second subcommand after a dollar-quoted default', () => { + expect(() => + requireSchemaLockTarget('ALTER TABLE t ADD COLUMN a TEXT DEFAULT $$x, y$$, ADD COLUMN b TEXT') + ).toThrow(/unparsed_schema_lock_target/) + }) + + it('leaves a numbered placeholder alone, because a tag cannot start with a digit', () => { + expect(sqlWithoutComments('ALTER TABLE t ADD COLUMN c TEXT -- $1 and $2')).toBe( + 'ALTER TABLE t ADD COLUMN c TEXT' + ) + }) + + it('treats an unterminated dollar quote as opaque to the end', () => { + expect(sqlWithoutComments('ALTER TABLE t ADD COLUMN c TEXT DEFAULT $$ -- unterminated')).toBe( + 'ALTER TABLE t ADD COLUMN c TEXT DEFAULT $$ -- unterminated' + ) + }) +}) diff --git a/cloud/packages/postgres-schema/src/schema-lock-target.ts b/cloud/packages/postgres-schema/src/schema-lock-target.ts new file mode 100644 index 00000000000..b4b42fc3b5a --- /dev/null +++ b/cloud/packages/postgres-schema/src/schema-lock-target.ts @@ -0,0 +1,264 @@ +// The object a boot-time DDL statement locks on Postgres, so the catalog can be asked whether it +// already exists before the statement joins the lock queue. `table` is kept exactly as the +// statement wrote it, schema qualification and quoting included, because it is fed to +// `to_regclass`; `name` is the bare identifier the catalog stores in `relname`/`attname`. +export type SchemaLockTarget = { + kind: 'index' | 'column' | 'constraint' + table: string + name: string + // The catalog answer that means this statement has nothing left to do. Creating statements skip + // on present; `DROP CONSTRAINT IF EXISTS` is the inverse, because nothing to drop is done. + skipWhen: 'present' | 'absent' +} + +// Keywords that sit in an identifier position when the optional clause before them is absent. +// Without this, `CREATE UNIQUE INDEX CONCURRENTLY ON t(c)` reads CONCURRENTLY as the index name and +// `ADD COLUMN IF NOT EXISTS` with no column reads IF as the column: a silently wrong target, which +// is worse than no target. Excluding them makes both throw instead. A column genuinely named `if` +// has to be quoted to be derivable, which is the safe direction to fail in. +const NOT_KEYWORD = '(?!(?:CONCURRENTLY|IF|NOT|EXISTS|ON|ONLY)\\b)' +const IDENTIFIER = `"(?:[^"]|"")*"|${NOT_KEYWORD}[A-Za-z_][A-Za-z0-9_$]*` +const QUALIFIED = `((?:${IDENTIFIER})(?:\\.(?:${IDENTIFIER}))?)` + +// `$$...$$` and `$tag$...$tag$` are opaque: a comment marker, comma, parenthesis or bracket inside +// one is text. The closing delimiter must match the opening tag exactly, so an inner `$$` inside a +// `$tag$` body is more text rather than the end. The tag cannot start with a digit, which is what +// keeps a `$1` placeholder from reading as an opener. +const DOLLAR_QUOTE = /\$(?:[A-Za-z_][A-Za-z0-9_]*)?\$/y + +function dollarQuoteEnd(sql: string, index: number): number | undefined { + DOLLAR_QUOTE.lastIndex = index + const opener = DOLLAR_QUOTE.exec(sql)?.[0] + if (opener === undefined) return undefined + const close = sql.indexOf(opener, index + opener.length) + return close === -1 ? sql.length : close + opener.length +} + +// Every comment, not only the block a ';'-split schema glues above a statement. A comment between +// two keywords (`ADD /* note */ COLUMN`) is invisible to the classification regexes AND to the +// must-parse shapes, so it used to yield no target and no throw: the statement ran with no +// pre-check at all, which is the one direction this must never fail in. Postgres treats a comment +// as whitespace, so each becomes a single space. Only classification reads this; the server is +// always sent the original text. +export function sqlWithoutComments(statement: string): string { + let stripped = '' + let quote: string | undefined + for (let index = 0; index < statement.length; index += 1) { + const character = statement[index]! + if (quote !== undefined) { + stripped += character + if (character !== quote) continue + if (statement[index + 1] === quote) { + stripped += quote + index += 1 + } else quote = undefined + continue + } + if (character === "'" || character === '"') { + quote = character + stripped += character + continue + } + if (character === '$') { + const end = dollarQuoteEnd(statement, index) + if (end !== undefined) { + stripped += statement.slice(index, end) + index = end - 1 + continue + } + } + if (character === '-' && statement[index + 1] === '-') { + const newline = statement.indexOf('\n', index) + index = newline === -1 ? statement.length : newline + stripped += ' ' + continue + } + if (character === '/' && statement[index + 1] === '*') { + // Postgres nests block comments, so a depth counter is what closes the right one. + let depth = 1 + index += 2 + while (index < statement.length && depth > 0) { + if (statement[index] === '/' && statement[index + 1] === '*') { + depth += 1 + index += 2 + } else if (statement[index] === '*' && statement[index + 1] === '/') { + depth -= 1 + index += 2 + } else index += 1 + } + index -= 1 + stripped += ' ' + continue + } + stripped += character + } + return stripped.trim() +} + +const CREATE_INDEX = new RegExp( + `^CREATE\\s+(?:UNIQUE\\s+)?INDEX\\s+(?:CONCURRENTLY\\s+)?(?:IF\\s+NOT\\s+EXISTS\\s+)?` + + `${QUALIFIED}\\s+ON\\s+(?:ONLY\\s+)?${QUALIFIED}`, + 'i' +) +const ADD_COLUMN = new RegExp( + `^ALTER\\s+TABLE\\s+(?:IF\\s+EXISTS\\s+)?(?:ONLY\\s+)?${QUALIFIED}\\s+` + + `ADD\\s+COLUMN\\s+(?:IF\\s+NOT\\s+EXISTS\\s+)?${QUALIFIED}`, + 'i' +) +const ADD_CONSTRAINT = new RegExp( + `^ALTER\\s+TABLE\\s+(?:IF\\s+EXISTS\\s+)?(?:ONLY\\s+)?${QUALIFIED}\\s+` + + `ADD\\s+CONSTRAINT\\s+${QUALIFIED}`, + 'i' +) +// `IF EXISTS` is required, not optional. A bare `DROP CONSTRAINT` on a missing constraint is an +// error the server is supposed to raise, and skipping it would swallow that. Without a target the +// statement throws at boot instead, which tells the author to write `IF EXISTS`. +const DROP_CONSTRAINT = new RegExp( + `^ALTER\\s+TABLE\\s+(?:IF\\s+EXISTS\\s+)?(?:ONLY\\s+)?${QUALIFIED}\\s+` + + `DROP\\s+CONSTRAINT\\s+IF\\s+EXISTS\\s+${QUALIFIED}`, + 'i' +) + +// Every statement shape that takes a relation lock before Postgres evaluates its existence test. +// `CREATE TABLE IF NOT EXISTS` is absent on purpose: it resolves a name against the schema and +// takes no lock on an existing table. +const TAKES_RELATION_LOCK = /^(?:CREATE\s+(?:UNIQUE\s+)?INDEX|ALTER\s+TABLE)\b/i + +export function takesRelationLock(statement: string): boolean { + return TAKES_RELATION_LOCK.test(sqlWithoutComments(statement)) +} + +// Splitting on '.' is not enough: `"a.b"` is one identifier containing a dot, not two parts. Each +// part is read quote-aware, with a doubled quote unescaped to one. +function qualifiedParts(written: string): { text: string; quoted: boolean }[] { + const parts: { text: string; quoted: boolean }[] = [] + let text = '' + let quoted = false + let wasQuoted = false + for (let index = 0; index < written.length; index += 1) { + const character = written[index]! + if (quoted) { + if (character !== '"') { + text += character + continue + } + if (written[index + 1] === '"') { + text += '"' + index += 1 + } else quoted = false + continue + } + if (character === '"') { + quoted = true + wasQuoted = true + } else if (character === '.') { + parts.push({ text, quoted: wasQuoted }) + text = '' + wasQuoted = false + } else text += character + } + parts.push({ text, quoted: wasQuoted }) + return parts +} + +// Postgres folds an unquoted identifier to lower case before storing it, so `Foo` is `foo` in +// relname, attname and conname. Comparing the written case would miss the row and rebuild the +// object on every boot. +function catalogName(written: string): string { + const last = qualifiedParts(written).pop() + if (!last) return written + return last.quoted ? last.text : last.text.toLowerCase() +} + +// Shapes whose lock target the pre-check must be able to derive. Deliberately looser than the +// regexes that parse them, so a statement that reads as one of these but does not parse is caught +// rather than falling through to the lock path. +const MUST_PARSE = [ + /^CREATE\s+(?:UNIQUE\s+)?INDEX\b/i, + /^ALTER\s+TABLE\b[\s\S]*\bADD\s+COLUMN\b/i, + /^ALTER\s+TABLE\b[\s\S]*\bADD\s+CONSTRAINT\b/i, + /^ALTER\s+TABLE\b[\s\S]*\bDROP\s+CONSTRAINT\b/i +] + +// Derived from the statement itself so a renamed index cannot drift away from its pre-check. +export function schemaLockTarget(statement: string): SchemaLockTarget | undefined { + const sql = sqlWithoutComments(statement) + const index = CREATE_INDEX.exec(sql) + if (index?.[1] && index[2]) { + return { kind: 'index', table: index[2], name: catalogName(index[1]), skipWhen: 'present' } + } + const column = ADD_COLUMN.exec(sql) + if (column?.[1] && column[2]) { + return { kind: 'column', table: column[1], name: catalogName(column[2]), skipWhen: 'present' } + } + const added = ADD_CONSTRAINT.exec(sql) + if (added?.[1] && added[2]) { + return { + kind: 'constraint', + table: added[1], + name: catalogName(added[2]), + skipWhen: 'present' + } + } + const dropped = DROP_CONSTRAINT.exec(sql) + if (dropped?.[1] && dropped[2]) { + return { + kind: 'constraint', + table: dropped[1], + name: catalogName(dropped[2]), + skipWhen: 'absent' + } + } + return undefined +} + +const ALTER_TABLE = /^ALTER\s+TABLE\b/i + +// A comma that separates ALTER TABLE subcommands rather than sitting inside a type, a default, a +// CHECK body or a dollar-quoted body. Takes comment-free SQL. Square brackets count as depth too, +// or an array type or `DEFAULT ARRAY[1, 2]` reads as a second subcommand and fails the boot. +function hasTopLevelComma(sql: string): boolean { + let depth = 0 + let quote: string | undefined + for (let index = 0; index < sql.length; index += 1) { + const character = sql[index] + if (quote !== undefined) { + if (character !== quote) continue + if (sql[index + 1] === quote) index += 1 + else quote = undefined + continue + } + if (character === '$') { + const end = dollarQuoteEnd(sql, index) + if (end !== undefined) { + index = end - 1 + continue + } + } + if (character === "'" || character === '"') quote = character + else if (character === '(' || character === '[') depth += 1 + else if (character === ')' || character === ']') depth -= 1 + else if (character === ',' && depth === 0) return true + } + return false +} + +// An index or column statement whose target cannot be read is the dangerous case: it would be sent +// unchecked and take the lock the pre-check exists to avoid, silently and on every boot. An +// auto-named `CREATE INDEX ON t(c)` lands here too, because nothing in the text says what the +// catalog will call it. Fail the boot with the statement instead. +export function requireSchemaLockTarget(statement: string): SchemaLockTarget | undefined { + const sql = sqlWithoutComments(statement) + // A multi-action ALTER TABLE parses to its FIRST subcommand's target only, so skipping on that + // one object would silently drop every later action for the life of the database. One action per + // statement, or no pre-check is possible. + if (ALTER_TABLE.test(sql) && hasTopLevelComma(sql)) { + throw new Error(`unparsed_schema_lock_target: ${sql}`) + } + const target = schemaLockTarget(statement) + if (target) return target + if (MUST_PARSE.some((shape) => shape.test(sql))) { + throw new Error(`unparsed_schema_lock_target: ${sql}`) + } + return undefined +}