From c7269ed08b71b61c933a09f9d25e1de1b90bade1 Mon Sep 17 00:00:00 2001 From: rgarcia <72655+rgarcia@users.noreply.github.com> Date: Fri, 30 Oct 2026 14:43:26 +0000 Subject: [PATCH 1/2] feat(durable): Postgres storage on the grouped commit contract Port the adapter to Storage.commit(batches, receive, context). A group of batches goes to a new commit_batches function in one round trip: it locks the session row and checks the owner once, then runs commit_batch per batch in a subtransaction, so a refusal by one of the adapter's checks keeps the batches before it and ends the call with the refusal, while any other error rolls the call back and rejects it. A batch with a document.copy starts a new call once the batches before it are durable, since the copy's source is read while preparing; every other group is one round trip. Schema version 3 also admits the sleeping task status added in 1.4.0. A round-trip test pins one query for a group of plain batches and three for a group whose last batch copies a document the earlier ones wrote. --- .../durable/src/storage/postgres/database.ts | 69 + .../durable/src/storage/postgres/index.ts | 15 + .../src/storage/postgres/migrations.ts | 325 +++++ .../durable/src/storage/postgres/storage.ts | 1228 +++++++++++++++++ .../durable/test/postgres-storage.test.ts | 225 +++ 5 files changed, 1862 insertions(+) create mode 100644 packages/durable/src/storage/postgres/database.ts create mode 100644 packages/durable/src/storage/postgres/index.ts create mode 100644 packages/durable/src/storage/postgres/migrations.ts create mode 100644 packages/durable/src/storage/postgres/storage.ts create mode 100644 packages/durable/test/postgres-storage.test.ts diff --git a/packages/durable/src/storage/postgres/database.ts b/packages/durable/src/storage/postgres/database.ts new file mode 100644 index 000000000000..19e61c43ebb9 --- /dev/null +++ b/packages/durable/src/storage/postgres/database.ts @@ -0,0 +1,69 @@ +/** Values the Postgres storage binds as query parameters. */ +export type PostgresValue = null | boolean | number | string; + +export interface PostgresQueryResult { + readonly rows: R[]; + readonly rowCount: number | null; +} + +/** Anything that runs one parameterized statement, such as a pool or a checked-out client. */ +export interface PostgresQueryable { + query(text: string, values?: PostgresValue[]): Promise>; +} + +/** A connection checked out of the pool. `release(error)` discards it instead of returning it. */ +export interface PostgresPoolClient extends PostgresQueryable { + release(error?: Error | boolean): void; + on(event: "error", listener: (error: Error) => void): unknown; + off(event: "error", listener: (error: Error) => void): unknown; +} + +/** + * The connection pool `PostgresStorage` runs on. A `pg.Pool` satisfies it structurally, so this package + * never imports a driver. The pool belongs to the caller: closing the storage does not end it. + * + * The storage keeps one client checked out for its whole lifetime to hold the session lock, so the pool + * needs at least two connections. + */ +export interface PostgresPool extends PostgresQueryable { + connect(): Promise; +} + +/** + * Run `callback` in a transaction on its own pooled client. When BEGIN, COMMIT, or ROLLBACK fails, the client is + * discarded so the pool never reuses a connection in an unknown transaction state. + */ +export async function inTransaction( + pool: PostgresPool, + begin: string, + callback: (client: PostgresQueryable) => Promise, +): Promise { + const client = await pool.connect(); + let discard: Error | undefined; + const control = async (sql: string) => { + try { + await client.query(sql); + } catch (error) { + discard = error as Error; + throw error; + } + }; + try { + await control(begin); + let result: T; + try { + result = await callback(client); + } catch (error) { + try { + await control("ROLLBACK"); + } catch (rollbackError) { + throw new AggregateError([error, rollbackError], "Durable Postgres transaction failed to roll back"); + } + throw error; + } + await control("COMMIT"); + return result; + } finally { + client.release(discard); + } +} diff --git a/packages/durable/src/storage/postgres/index.ts b/packages/durable/src/storage/postgres/index.ts new file mode 100644 index 000000000000..92f9b9557a86 --- /dev/null +++ b/packages/durable/src/storage/postgres/index.ts @@ -0,0 +1,15 @@ +export type { + PostgresPool, + PostgresPoolClient, + PostgresQueryable, + PostgresQueryResult, + PostgresValue, +} from "./database.ts"; +export { + applyPostgresMigrations, + CURRENT_POSTGRES_SCHEMA_VERSION, + DEFAULT_POSTGRES_SCHEMA, + POSTGRES_MIGRATIONS, + type PostgresMigration, +} from "./migrations.ts"; +export { PostgresStorage, type PostgresStorageOptions } from "./storage.ts"; diff --git a/packages/durable/src/storage/postgres/migrations.ts b/packages/durable/src/storage/postgres/migrations.ts new file mode 100644 index 000000000000..546e51df4d62 --- /dev/null +++ b/packages/durable/src/storage/postgres/migrations.ts @@ -0,0 +1,325 @@ +import { inTransaction, type PostgresPool, type PostgresQueryable } from "./database.ts"; + +export type PostgresMigration = { + readonly version: number; + /** Statements with `{schema}` standing for the quoted schema name. */ + readonly statements: readonly string[]; +}; + +// Every table is keyed by session_id first, so one schema holds any number of independent Sessions. +// Records are `json`, not `jsonb`: `json` keeps the exact text, including key order and `\u0000` escapes that `jsonb` +// rejects. +const INITIAL_SCHEMA: readonly string[] = [ + `CREATE TABLE {schema}.sessions ( + session_id TEXT PRIMARY KEY, + next_id BIGINT NOT NULL, + next_seq BIGINT NOT NULL, + owner TEXT NOT NULL + )`, + `CREATE TABLE {schema}.record_ids ( + session_id TEXT NOT NULL, + id BIGINT NOT NULL, + record_type TEXT NOT NULL CHECK (record_type IN ('conversation', 'entry', 'task', 'submission', 'document')), + PRIMARY KEY (session_id, id) + )`, + `CREATE TABLE {schema}.conversations ( + session_id TEXT NOT NULL, + id BIGINT NOT NULL, + owner_conversation_id BIGINT, + owner_task_id BIGINT, + record JSON NOT NULL, + PRIMARY KEY (session_id, id) + )`, + "CREATE INDEX conversations_by_owner_conversation ON {schema}.conversations (session_id, owner_conversation_id, id)", + "CREATE INDEX conversations_by_owner_task ON {schema}.conversations (session_id, owner_task_id, id)", + `CREATE TABLE {schema}.entries ( + session_id TEXT NOT NULL, + id BIGINT NOT NULL, + conversation_id BIGINT NOT NULL, + head BIGINT, + commit_seq BIGINT NOT NULL, + record JSON NOT NULL, + PRIMARY KEY (session_id, id) + )`, + "CREATE INDEX entries_by_conversation ON {schema}.entries (session_id, conversation_id, id DESC)", + `CREATE INDEX entry_heads_by_conversation ON {schema}.entries (session_id, conversation_id, id DESC) + WHERE head IS NOT NULL`, + `CREATE TABLE {schema}.tasks ( + session_id TEXT NOT NULL, + id BIGINT NOT NULL, + conversation_id BIGINT NOT NULL, + kind TEXT NOT NULL, + status TEXT NOT NULL CHECK (status IN ('pending', 'running', 'waiting', 'completing', 'terminal')), + abort_requested BOOLEAN NOT NULL, + background BOOLEAN NOT NULL, + record JSON NOT NULL, + PRIMARY KEY (session_id, id) + )`, + "CREATE INDEX tasks_by_status ON {schema}.tasks (session_id, status, id)", + "CREATE INDEX tasks_by_conversation ON {schema}.tasks (session_id, conversation_id, id)", + "CREATE INDEX tasks_by_kind ON {schema}.tasks (session_id, kind, id)", + "CREATE INDEX tasks_by_abort_requested ON {schema}.tasks (session_id, abort_requested, id)", + "CREATE INDEX tasks_by_background ON {schema}.tasks (session_id, background, id)", + `CREATE TABLE {schema}.submissions ( + session_id TEXT NOT NULL, + id BIGINT NOT NULL, + conversation_id BIGINT NOT NULL, + request_id TEXT, + status TEXT NOT NULL CHECK (status IN ('queued', 'placed', 'done', 'unanswered')), + record JSON NOT NULL, + PRIMARY KEY (session_id, id) + )`, + "CREATE INDEX submissions_by_request ON {schema}.submissions (session_id, conversation_id, request_id)", + "CREATE INDEX submissions_by_conversation ON {schema}.submissions (session_id, conversation_id, id)", + "CREATE INDEX submissions_by_status ON {schema}.submissions (session_id, status, id)", + `CREATE TABLE {schema}.documents ( + session_id TEXT NOT NULL, + id BIGINT NOT NULL, + kind TEXT NOT NULL, + family SMALLINT NOT NULL CHECK (family IN (0, 1)), + key_value TEXT NOT NULL, + scope_kind TEXT NOT NULL CHECK (scope_kind IN ('session', 'conversation', 'task')), + owner_id BIGINT NOT NULL, + created_at BIGINT NOT NULL, + retired_at BIGINT, + record JSON NOT NULL, + PRIMARY KEY (session_id, id) + )`, + `CREATE INDEX documents_by_address ON {schema}.documents + (session_id, kind, scope_kind, owner_id, family, key_value, created_at DESC, retired_at)`, + "CREATE INDEX documents_by_scope ON {schema}.documents (session_id, scope_kind, owner_id, id)", + "CREATE INDEX documents_by_scope_kind ON {schema}.documents (session_id, scope_kind, owner_id, kind, id)", + `CREATE TABLE {schema}.document_revisions ( + session_id TEXT NOT NULL, + document_id BIGINT NOT NULL, + seq BIGINT NOT NULL, + kind TEXT NOT NULL CHECK (kind IN ('base', 'delta')), + version INTEGER NOT NULL, + content JSON NOT NULL, + PRIMARY KEY (session_id, document_id, seq) + )`, + "CREATE INDEX document_revisions_by_kind ON {schema}.document_revisions (session_id, document_id, kind, seq DESC)", +]; + +// One commit is one call. The function locks the session row, checks that the caller still owns the session at +// the sequence it expects, validates the batch against the stored state, writes every table, and advances the +// session. Each rule raises the adapter's message with its own SQLSTATE so the adapter can map it back; see +// `commitErrorCodes` in storage.ts. +// +// Records and content travel inside the batch as JSON text and are cast to `json` on insert. Postgres accepts a +// lone UTF-16 surrogate escape when it stores `json` but rejects it in any function that decodes string values, +// such as `json_to_recordset` and `->>`; so nothing here decodes a record, and the current-only test matches the +// stored record's text. +const COMMIT_FUNCTION = `CREATE OR REPLACE FUNCTION {schema}.commit_batch( + p_session TEXT, p_owner TEXT, p_seq BIGINT, p_next_id BIGINT, p_batch JSON +) RETURNS BIGINT LANGUAGE plpgsql AS $fn$ +DECLARE + v_owner TEXT; + v_seq BIGINT; + v_id BIGINT; + v_type TEXT; + v_exists BOOLEAN; + v_retired BIGINT; + v_version INTEGER; + v_key TEXT; + v_live JSONB := '{}'::jsonb; + v_count INTEGER; + doc RECORD; + stored RECORD; + live RECORD; +BEGIN + SELECT owner, next_seq INTO v_owner, v_seq FROM {schema}.sessions WHERE session_id = p_session FOR UPDATE; + IF NOT FOUND THEN + RAISE EXCEPTION 'Durable Postgres session metadata is missing' USING ERRCODE = 'DU001'; + END IF; + IF v_owner <> p_owner OR v_seq <> p_seq THEN + RAISE EXCEPTION 'Durable Postgres session % was opened by another process', p_session USING ERRCODE = 'DU002'; + END IF; + + SELECT c.id, r.record_type INTO v_id, v_type + FROM json_to_recordset(p_batch->'claims') AS c(id BIGINT, type TEXT) + JOIN {schema}.record_ids r ON r.session_id = p_session AND r.id = c.id + WHERE c.type IN ('conversation', 'entry', 'document') OR r.record_type <> c.type + LIMIT 1; + IF FOUND THEN + RAISE EXCEPTION 'ID % already belongs to %', v_id, v_type USING ERRCODE = 'DU003'; + END IF; + + FOR doc IN SELECT * FROM json_to_recordset(p_batch->'documents') AS x( + id BIGINT, is_create BOOLEAN, retire BOOLEAN, kind TEXT, family SMALLINT, key_value TEXT, + scope_kind TEXT, owner_id BIGINT, content_kind TEXT, version INTEGER) + LOOP + SELECT x.retired_at, x.kind, x.family, x.key_value, x.scope_kind, x.owner_id + INTO stored FROM {schema}.documents x WHERE x.session_id = p_session AND x.id = doc.id; + v_exists := FOUND; + IF NOT doc.is_create AND NOT v_exists THEN + RAISE EXCEPTION 'Unknown document: %', doc.id USING ERRCODE = 'DU004'; + END IF; + IF doc.is_create AND v_exists THEN + RAISE EXCEPTION 'Document % already exists', doc.id USING ERRCODE = 'DU005'; + END IF; + IF v_exists AND stored.retired_at IS NOT NULL THEN + RAISE EXCEPTION 'Document % is retired', doc.id USING ERRCODE = 'DU006'; + END IF; + IF doc.content_kind = 'delta' THEN + SELECT r.version INTO v_version FROM {schema}.document_revisions r + WHERE r.session_id = p_session AND r.document_id = doc.id ORDER BY r.seq DESC LIMIT 1; + IF NOT FOUND THEN + RAISE EXCEPTION 'Document % delta has no base', doc.id USING ERRCODE = 'DU007'; + END IF; + IF v_version <> doc.version THEN + RAISE EXCEPTION 'Document % version transition requires a base', doc.id USING ERRCODE = 'DU008'; + END IF; + END IF; + IF doc.is_create OR doc.retire THEN + IF doc.is_create THEN + v_key := concat_ws('|', doc.kind, doc.scope_kind, doc.owner_id, doc.family, doc.key_value); + ELSE + v_key := concat_ws('|', stored.kind, stored.scope_kind, stored.owner_id, stored.family, stored.key_value); + END IF; + IF NOT v_live ? v_key THEN + IF doc.is_create THEN + SELECT count(*) INTO v_count FROM {schema}.documents x WHERE x.session_id = p_session + AND x.kind = doc.kind AND x.scope_kind = doc.scope_kind AND x.owner_id = doc.owner_id + AND x.family = doc.family AND x.key_value = doc.key_value AND x.retired_at IS NULL; + ELSE + SELECT count(*) INTO v_count FROM {schema}.documents x WHERE x.session_id = p_session + AND x.kind = stored.kind AND x.scope_kind = stored.scope_kind AND x.owner_id = stored.owner_id + AND x.family = stored.family AND x.key_value = stored.key_value AND x.retired_at IS NULL; + END IF; + v_count := LEAST(v_count, 1); + ELSE + v_count := (v_live->>v_key)::integer; + END IF; + IF doc.retire AND v_exists THEN v_count := v_count - 1; END IF; + IF doc.is_create AND NOT doc.retire THEN v_count := v_count + 1; END IF; + v_live := v_live || jsonb_build_object(v_key, v_count); + END IF; + END LOOP; + FOR live IN SELECT value FROM jsonb_each_text(v_live) LOOP + IF live.value::integer > 1 THEN + RAISE EXCEPTION 'Document address already has a current incarnation' USING ERRCODE = 'DU009'; + END IF; + END LOOP; + + INSERT INTO {schema}.record_ids (session_id, id, record_type) + SELECT p_session, c.id, c.type FROM json_to_recordset(p_batch->'claims') AS c(id BIGINT, type TEXT) + ON CONFLICT (session_id, id) DO NOTHING; + INSERT INTO {schema}.conversations (session_id, id, owner_conversation_id, owner_task_id, record) + SELECT p_session, x.id, x.owner_conversation_id, x.owner_task_id, x.record::json + FROM json_to_recordset(p_batch->'conversations') + AS x(id BIGINT, owner_conversation_id BIGINT, owner_task_id BIGINT, record TEXT); + INSERT INTO {schema}.entries (session_id, id, conversation_id, head, commit_seq, record) + SELECT p_session, x.id, x.conversation_id, x.head, p_seq, x.record::json + FROM json_to_recordset(p_batch->'entries') AS x(id BIGINT, conversation_id BIGINT, head BIGINT, record TEXT); + INSERT INTO {schema}.tasks (session_id, id, conversation_id, kind, status, abort_requested, background, record) + SELECT p_session, x.id, x.conversation_id, x.kind, x.status, x.abort_requested, x.background, x.record::json + FROM json_to_recordset(p_batch->'tasks') AS x(id BIGINT, conversation_id BIGINT, kind TEXT, status TEXT, + abort_requested BOOLEAN, background BOOLEAN, record TEXT) + ON CONFLICT (session_id, id) DO UPDATE SET conversation_id = excluded.conversation_id, kind = excluded.kind, + status = excluded.status, abort_requested = excluded.abort_requested, background = excluded.background, + record = excluded.record; + INSERT INTO {schema}.submissions (session_id, id, conversation_id, request_id, status, record) + SELECT p_session, x.id, x.conversation_id, x.request_id, x.status, x.record::json + FROM json_to_recordset(p_batch->'submissions') + AS x(id BIGINT, conversation_id BIGINT, request_id TEXT, status TEXT, record TEXT) + ON CONFLICT (session_id, id) DO UPDATE SET conversation_id = excluded.conversation_id, + request_id = excluded.request_id, status = excluded.status, record = excluded.record; + + -- A retired record gains retiredAt by appending to its text, which keeps the rest of the record verbatim. + UPDATE {schema}.documents x SET retired_at = p_seq, + record = (left(x.record::text, length(x.record::text) - 1) || ',"retiredAt":' || p_seq || '}')::json + FROM json_to_recordset(p_batch->'documents') AS d(id BIGINT, is_create BOOLEAN, retire BOOLEAN) + WHERE x.session_id = p_session AND x.id = d.id AND d.retire AND NOT d.is_create; + INSERT INTO {schema}.documents + (session_id, id, kind, family, key_value, scope_kind, owner_id, created_at, retired_at, record) + SELECT p_session, d.id, d.kind, d.family, d.key_value, d.scope_kind, d.owner_id, p_seq, + CASE WHEN d.retire THEN p_seq END, d.record::json + FROM json_to_recordset(p_batch->'documents') AS d(id BIGINT, is_create BOOLEAN, retire BOOLEAN, kind TEXT, + family SMALLINT, key_value TEXT, scope_kind TEXT, owner_id BIGINT, record TEXT) + WHERE d.is_create; + -- A current-only document keeps no history: a base replaces its revisions, and a retirement clears them. The + -- record is matched as text; its only object values are the adapter's own, so the pair cannot appear elsewhere. + DELETE FROM {schema}.document_revisions r + USING json_to_recordset(p_batch->'documents') AS d(id BIGINT, retire BOOLEAN, content_kind TEXT) + JOIN {schema}.documents x ON x.session_id = p_session AND x.id = d.id + WHERE r.session_id = p_session AND r.document_id = d.id + AND (x.scope_kind <> 'conversation' OR x.record::text LIKE '%"history":"latest"%') + AND (d.retire OR d.content_kind = 'base'); + INSERT INTO {schema}.document_revisions (session_id, document_id, seq, kind, version, content) + SELECT p_session, d.id, p_seq, d.content_kind, d.version, d.content::json + FROM json_to_recordset(p_batch->'documents') + AS d(id BIGINT, retire BOOLEAN, content_kind TEXT, version INTEGER, content TEXT) + JOIN {schema}.documents x ON x.session_id = p_session AND x.id = d.id + WHERE d.content_kind IS NOT NULL + AND NOT (d.retire AND (x.scope_kind <> 'conversation' OR x.record::text LIKE '%"history":"latest"%')); + + UPDATE {schema}.sessions SET next_id = GREATEST(next_id, p_next_id), next_seq = p_seq + 1 + WHERE session_id = p_session; + RETURN p_seq; +END +$fn$`; + +/** Immutable, ordered schema history. Append new migrations after the initial schema ships. */ +export const POSTGRES_MIGRATIONS: readonly PostgresMigration[] = [ + { version: 1, statements: INITIAL_SCHEMA }, + { version: 2, statements: [COMMIT_FUNCTION] }, +]; + +export const CURRENT_POSTGRES_SCHEMA_VERSION = POSTGRES_MIGRATIONS.at(-1)?.version ?? 0; + +export const DEFAULT_POSTGRES_SCHEMA = "pi_durable"; + +/** Quote one identifier for interpolation into SQL text. */ +export const quoteIdentifier = (name: string): string => `"${name.replaceAll('"', '""')}"`; + +type SchemaRow = { readonly version: number }; + +/** + * Create the schema if needed and apply all pending migrations in one transaction. Concurrent callers on the + * same schema serialize on a transaction-scoped advisory lock. + */ +export async function applyPostgresMigrations( + pool: PostgresPool, + schema: string = DEFAULT_POSTGRES_SCHEMA, + migrations: readonly PostgresMigration[] = POSTGRES_MIGRATIONS, +): Promise { + for (let index = 0; index < migrations.length; index++) { + if (migrations[index]?.version !== index + 1) { + throw new Error("Durable Postgres migrations must have contiguous versions starting at 1"); + } + } + const quoted = quoteIdentifier(schema); + await inTransaction(pool, "BEGIN", (client) => migrate(client, schema, quoted, migrations)); +} + +async function migrate( + client: PostgresQueryable, + schema: string, + quoted: string, + migrations: readonly PostgresMigration[], +): Promise { + await client.query("SELECT pg_advisory_xact_lock(hashtextextended($1, 0))", [`pi-durable:migrate:${schema}`]); + await client.query(`CREATE SCHEMA IF NOT EXISTS ${quoted}`); + await client.query(`CREATE TABLE IF NOT EXISTS ${quoted}.durable_schema ( + singleton INTEGER PRIMARY KEY CHECK (singleton = 1), + version INTEGER NOT NULL CHECK (version >= 0) + )`); + await client.query( + `INSERT INTO ${quoted}.durable_schema (singleton, version) VALUES (1, 0) ON CONFLICT (singleton) DO NOTHING`, + ); + const row = (await client.query(`SELECT version FROM ${quoted}.durable_schema WHERE singleton = 1`)) + .rows[0]; + if (row === undefined) throw new Error("Durable Postgres schema metadata is missing"); + const currentVersion = migrations.at(-1)?.version ?? 0; + if (row.version > currentVersion) { + throw new Error( + `Durable Postgres schema version ${row.version} is newer than supported version ${currentVersion}`, + ); + } + for (const migration of migrations) { + if (migration.version <= row.version) continue; + for (const statement of migration.statements) await client.query(statement.replaceAll("{schema}", quoted)); + await client.query(`UPDATE ${quoted}.durable_schema SET version = $1 WHERE singleton = 1`, [migration.version]); + } +} diff --git a/packages/durable/src/storage/postgres/storage.ts b/packages/durable/src/storage/postgres/storage.ts new file mode 100644 index 000000000000..e95732a8b62e --- /dev/null +++ b/packages/durable/src/storage/postgres/storage.ts @@ -0,0 +1,1228 @@ +import type { Context, JsonValue } from "@earendil-works/chord"; +import { apply, type Op } from "@earendil-works/chord/delta"; +/** Thrown when the Session was opened by another process; any error from `commit` fails the Session. */ +export class StorageRejected extends Error { + constructor(message: string, options?: ErrorOptions) { + super(message, options); + this.name = "StorageRejected"; + } +} + +import { StorageRequestError } from "../../errors.ts"; +import { idFromNumber, seqFromNumber } from "../../ids.ts"; +import type { + ConversationId, + ConversationQuery, + ConversationRecord, + Cursor, + DocumentAddress, + DocumentContent, + DocumentCreate, + DocumentId, + DocumentPoint, + DocumentQuery, + DocumentRecord, + EntryId, + EntryQuery, + EntryRecord, + Id, + JsonObject, + Page, + ScanOrder, + Seq, + Storage, + StorageWrite, + StoredDocument, + SubmissionId, + SubmissionQuery, + SubmissionRecord, + TaskId, + TaskQuery, + TaskRecord, +} from "../../types.ts"; +import { nextCursor, type ScanStart, scanStart } from "../scan.ts"; +import type { PostgresPool, PostgresQueryable, PostgresValue } from "./database.ts"; +import { applyPostgresMigrations, DEFAULT_POSTGRES_SCHEMA, quoteIdentifier } from "./migrations.ts"; + +type StoredTask = TaskRecord; +type TableName = "conversation" | "entry" | "task" | "submission" | "document"; +type JsonRow = { readonly record: string }; +type EntryJsonRow = { readonly record: string; readonly commit_seq: string }; +type MetadataRow = { readonly next_id: string; readonly next_seq: string }; +/** A conversation's own entries over one contiguous id range ending at its newest entry, newest-last. */ +type EntryWindow = { ids: EntryId[]; complete: boolean }; +type DocumentState = { readonly record: string; readonly version: number; value: JsonObject; deltasSinceBase: number }; +type DocumentWrites = { + readonly claims: DocumentId[]; + readonly rows: unknown[]; + /** What each document action means for the caches once the commit succeeds. */ + readonly applied: Map< + DocumentId, + { + readonly created: DocumentRecord | undefined; + readonly retired: boolean; + readonly content: DocumentContent | undefined; + } + >; +}; + +/** SQLSTATEs the commit function raises, each for one adapter error. */ +const commitErrorCodes: Readonly> = { + DU001: true, + DU002: true, + DU003: true, + DU004: true, + DU005: true, + DU006: true, + DU007: true, + DU008: true, + DU009: true, +}; +type DocumentAction = { + create?: DocumentCreate; + copy?: Extract["source"]; + content?: DocumentContent; + retire: boolean; +}; +type ScopeColumns = { + readonly scopeKind: DocumentRecord["scope"]["kind"]; + readonly ownerId: number; +}; +type Tables = { + readonly sessions: string; + readonly recordIds: string; + readonly conversations: string; + readonly entries: string; + readonly tasks: string; + readonly submissions: string; + readonly documents: string; + readonly revisions: string; +}; + +export interface PostgresStorageOptions { + /** The Session stored by this storage. Sessions in one schema are independent. */ + readonly session: string; + /** Schema holding the tables, created on first open. Defaults to `pi_durable`. */ + readonly schema?: string; + /** + * Apply schema migrations on open. Defaults to true. A deployment that migrates in an explicit step before + * starting, or opens many Sessions per process, passes false so an open costs two statements and takes no + * schema-wide lock. + */ + readonly migrate?: boolean; +} + +const parseJson = (value: string): T => JSON.parse(value) as T; +const encodeJson = (value: unknown): string => JSON.stringify(value) as string; +// Drivers encode parameters as UTF-8, which replaces lone UTF-16 surrogates. JSON encoding keeps indexed identities +// lossless. +const encodeIndexedString = (value: string): string => JSON.stringify(value); + +const cursorId = >(cursor: Cursor | undefined): I | undefined => { + const after = cursor?.after; + if (after === undefined) return undefined; + if (typeof after !== "number" || !Number.isSafeInteger(after)) + throw new StorageRequestError("Invalid storage cursor"); + return idFromNumber(after); +}; + +const page = }>( + values: readonly T[], + limit: number, + order: ScanOrder = "ascending", +): Page => { + const items = values.slice(0, limit); + if (values.length <= limit) return { items }; + return { items, next: nextCursor(items.at(-1)!.id, order) }; +}; + +/** The ID condition, its parameter, and the `ORDER BY` direction of a table scan. */ +const scanSql = ( + start: ScanStart, +): { readonly comparator: string; readonly param: number; readonly direction: string } => + start.order === "ascending" + ? { comparator: ">", param: start.after ?? -1, direction: "ASC" } + : { comparator: "<", param: start.after ?? Number.MAX_SAFE_INTEGER, direction: "DESC" }; + +const scopeColumns = (scope: DocumentRecord["scope"]): ScopeColumns => { + switch (scope.kind) { + case "session": + return { scopeKind: "session", ownerId: 0 }; + case "conversation": + return { scopeKind: "conversation", ownerId: scope.conversationId }; + case "task": + return { scopeKind: "task", ownerId: scope.taskId }; + } +}; + +const addressParts = (address: DocumentAddress | DocumentCreate | DocumentRecord) => { + const scope = scopeColumns(address.scope); + return { + kind: encodeIndexedString(address.kind), + ...scope, + family: address.key === undefined ? 0 : 1, + keyValue: encodeIndexedString(address.key ?? ""), + }; +}; + +const addressKey = (address: DocumentAddress | DocumentCreate | DocumentRecord): string => { + const parts = addressParts(address); + return JSON.stringify([parts.kind, parts.scopeKind, parts.ownerId, parts.family, parts.keyValue]); +}; + +const isAliveAt = (record: DocumentRecord, at: DocumentPoint): boolean => { + if (at === "current") return record.retiredAt === undefined; + return record.createdAt <= at && (record.retiredAt === undefined || at < record.retiredAt); +}; + +const isCurrentOnly = (record: DocumentRecord): boolean => + record.scope.kind !== "conversation" || record.history === "latest"; + +/** Entries loaded per conversation on first use, and the most kept as commits extend the window. */ +const ENTRY_WINDOW_LOAD = 256; +const ENTRY_WINDOW_MAX = 1024; + +const requestKey = (conversationId: ConversationId, requestId: string): string => `${conversationId}:${requestId}`; + +/** Index of the first element greater than `value` in an ascending array. */ +const upperBound = (values: readonly number[], value: number): number => { + let low = 0; + let high = values.length; + while (low < high) { + const middle = (low + high) >>> 1; + if (values[middle] <= value) low = middle + 1; + else high = middle; + } + return low; +}; + +const writeId = (write: StorageWrite): Id | undefined => { + switch (write.type) { + case "conversation": + case "entry": + case "task": + case "submission": + return write.value.id; + case "document.create": + case "document.copy": + return write.record.id; + case "document.change": + case "document.retire": + return undefined; + } +}; + +/** + * Postgres implementation of the Pico storage contract. + * + * One process owns a Session at a time. `open()` writes a fresh owner token that every commit checks under a row + * lock, so a later open takes the Session over and a commit from the earlier owner is rejected, after which that + * storage stops accepting work. Nothing is held between operations: an open storage keeps no connection. + */ +export class PostgresStorage implements Storage { + private readonly pool: PostgresPool; + private readonly session: string; + private readonly owner: string; + private readonly tables: Tables; + private nextId: number; + /** The sequence the next commit takes. This storage is the session's only writer, so it advances by one per commit. */ + private nextSeq: number; + private readonly commitFunction: string; + private closed = false; + private lost: Error | undefined; + private closing: Promise | undefined; + /** + * Conversation records read so far. This storage is the session's only writer and conversation records are + * only ever inserted, so a record stays valid until a commit of this storage writes that conversation. + */ + private readonly conversations = new Map(); + // Read-through caches with the same justification: this storage is the Session's only writer, so a record it + // has read stays valid until it writes that record, and every write passes through `commit` below. Values are + // kept JSON-encoded so a caller gets a detached object each time, as from the database. + private readonly tasks = new Map(); + private readonly submissions = new Map(); + // A request key with no submission is remembered too: nothing else can create one, so a miss stays a miss + // until this storage commits that submission. + private readonly submissionsByRequest = new Map(); + private readonly entries = new Map(); + private readonly entryWindows = new Map(); + private readonly headMarkers = new Map(); + private readonly documentStates = new Map(); + private documentScans = new Map(); + private documentAddresses = new Map(); + private admitted = 0; + private drained: (() => void) | undefined; + + private constructor( + pool: PostgresPool, + options: { session: string; schema: string; owner: string; nextId: number; nextSeq: number }, + ) { + this.pool = pool; + this.session = options.session; + this.owner = options.owner; + this.nextId = options.nextId; + this.nextSeq = options.nextSeq; + const schema = quoteIdentifier(options.schema); + this.commitFunction = `${schema}.commit_batch`; + this.tables = { + sessions: `${schema}.sessions`, + recordIds: `${schema}.record_ids`, + conversations: `${schema}.conversations`, + entries: `${schema}.entries`, + tasks: `${schema}.tasks`, + submissions: `${schema}.submissions`, + documents: `${schema}.documents`, + revisions: `${schema}.document_revisions`, + }; + } + + /** Migrate the schema and take ownership of the Session, from any process that held it before. */ + static async open(pool: PostgresPool, options: PostgresStorageOptions): Promise { + if (options.session.length === 0) throw new TypeError("PostgresStorage requires a non-empty session"); + const schema = options.schema ?? DEFAULT_POSTGRES_SCHEMA; + if (options.migrate !== false) await applyPostgresMigrations(pool, schema); + const owner = globalThis.crypto.randomUUID(); + const metadata = ( + await pool.query( + `INSERT INTO ${quoteIdentifier(schema)}.sessions (session_id, next_id, next_seq, owner) + VALUES ($1, 2, 1, $2) + ON CONFLICT (session_id) DO UPDATE SET owner = excluded.owner + RETURNING next_id, next_seq`, + [options.session, owner], + ) + ).rows[0]; + if (metadata === undefined) throw new Error("Durable Postgres session metadata is missing"); + return new PostgresStorage(pool, { + session: options.session, + schema, + owner, + nextId: Number(metadata.next_id), + nextSeq: Number(metadata.next_seq), + }); + } + + commit(writes: readonly StorageWrite[], _context: Context): Promise { + return this.admit(async () => { + const documentActions = this.prepareDocumentActions(writes); + const candidateNextId = this.candidateNextId(writes); + this.checkBatchIds(writes); + const seq = seqFromNumber(this.nextSeq); + const documentWrites = await this.prepareDocumentWrites(documentActions, seq); + // One round trip: the commit function locks the session row, checks that this storage still owns it + // at the sequence it expects, validates the batch against the stored state, and writes every table. + try { + await this.pool.query(`SELECT ${this.commitFunction}($1, $2, $3, $4, $5::json)`, [ + this.session, + this.owner, + seq, + candidateNextId, + this.encodeBatch(writes, documentWrites), + ]); + } catch (error) { + throw this.commitError(error); + } + this.nextSeq = seq + 1; + this.nextId = Math.max(this.nextId, candidateNextId); + this.absorb(writes, documentActions, documentWrites, seq); + return seq; + }); + } + + /** Maps an error raised by the commit function to the adapter's own errors; other errors pass through. */ + private commitError(error: unknown): unknown { + const code = (error as { readonly code?: unknown }).code; + const message = (error as { readonly message?: unknown }).message; + if (typeof code !== "string" || typeof message !== "string" || !(code in commitErrorCodes)) return error; + if (code === "DU002") { + const lost = new Error(message); + this.loseOwnership(lost); + return lost; + } + return new Error(message); + } + + /** The batch as the commit function reads it. Records and content go as JSON text, cast to `json` on insert. */ + private encodeBatch(writes: readonly StorageWrite[], documents: DocumentWrites): string { + const claims = new Map, TableName>(); + const conversations: unknown[] = []; + const entries: unknown[] = []; + const tasks = new Map, unknown>(); + const submissions = new Map, unknown>(); + for (const write of writes) { + switch (write.type) { + case "conversation": + claims.set(write.value.id, "conversation"); + conversations.push({ + id: write.value.id, + owner_conversation_id: write.value.owner?.conversationId ?? null, + owner_task_id: write.value.owner?.taskId ?? null, + record: encodeJson(write.value), + }); + break; + case "entry": + claims.set(write.value.id, "entry"); + entries.push({ + id: write.value.id, + conversation_id: write.value.conversationId, + head: write.value.head ?? null, + record: encodeJson(write.value), + }); + break; + case "task": + claims.set(write.value.id, "task"); + tasks.set(write.value.id, { + id: write.value.id, + conversation_id: write.value.conversationId, + kind: encodeIndexedString(write.value.kind), + status: write.value.state.status, + abort_requested: write.value.abortRequested, + background: write.value.background, + record: encodeJson(write.value), + }); + break; + case "submission": + claims.set(write.value.id, "submission"); + submissions.set(write.value.id, { + id: write.value.id, + conversation_id: write.value.conversationId, + request_id: write.value.requestId === undefined ? null : encodeIndexedString(write.value.requestId), + status: write.value.status, + record: encodeJson(write.value), + }); + break; + case "document.create": + case "document.copy": + case "document.change": + case "document.retire": + break; + } + } + for (const id of documents.claims) claims.set(id, "document"); + return encodeJson({ + claims: [...claims].map(([id, type]) => ({ id, type })), + conversations, + entries, + tasks: [...tasks.values()], + submissions: [...submissions.values()], + documents: documents.rows, + }); + } + + async mintId>(): Promise { + this.assertOpen(); + if (!Number.isSafeInteger(this.nextId)) throw new Error("ID space is exhausted"); + return idFromNumber(this.nextId++); + } + + conversation(id: ConversationId, _context: Context): Promise { + return this.admit(() => this.readConversation(id)); + } + + scanConversations( + query: ConversationQuery, + limit: number, + cursor: Cursor | undefined, + _context: Context, + ): Promise> { + return this.admit(async () => { + const start = scanStart(query.order, cursor, "ascending"); + const scan = scanSql(start); + const params: PostgresValue[] = [this.session]; + const clauses = ["session_id = $1", `id ${scan.comparator} $${params.push(scan.param)}`]; + if (query.ownerConversationId !== undefined) { + clauses.push(`owner_conversation_id = $${params.push(query.ownerConversationId)}`); + } + if (query.ownerTaskId !== undefined) clauses.push(`owner_task_id = $${params.push(query.ownerTaskId)}`); + const rows = await this.records( + `SELECT record::text AS record FROM ${this.tables.conversations} WHERE ${clauses.join(" AND ")} + ORDER BY id ${scan.direction} LIMIT $${params.push(limit + 1)}`, + params, + ); + return page(rows.map(parseJson), limit, start.order); + }); + } + + entry(id: EntryId, context: Context): Promise<{ readonly entry: EntryRecord; readonly commitSeq: Seq } | undefined>; + entry( + conversationId: ConversationId, + id: EntryId, + context: Context, + ): Promise<{ readonly entry: EntryRecord; readonly commitSeq: Seq } | undefined>; + entry( + idOrConversationId: EntryId | ConversationId, + idOrContext: EntryId | Context, + context?: Context, + ): Promise<{ readonly entry: EntryRecord; readonly commitSeq: Seq } | undefined> { + return this.admit(() => this.readEntry(idOrConversationId, idOrContext, context)); + } + + findLatestHeadMarker( + conversationId: ConversationId, + atOrBeforeEntryId: EntryId | undefined, + _context: Context, + ): Promise<(EntryRecord & { readonly head: EntryId }) | undefined> { + return this.admit(() => this.readLatestHeadMarker(conversationId, atOrBeforeEntryId)); + } + + scanEntries( + query: EntryQuery, + limit: number, + cursor: Cursor | undefined, + _context: Context, + ): Promise> { + return this.admit(() => this.readEntries(query, limit, cursor)); + } + + private async readEntry( + idOrConversationId: EntryId | ConversationId, + idOrContext: EntryId | Context, + context?: Context, + ): Promise<{ readonly entry: EntryRecord; readonly commitSeq: Seq } | undefined> { + const id = + context === undefined + ? idFromNumber(idOrConversationId) + : typeof idOrContext === "number" + ? idFromNumber(idOrContext) + : undefined; + if (id === undefined) throw new TypeError("Storage.entry() requires an entry ID"); + let conversation: ConversationRecord | undefined; + if (context !== undefined) { + const conversationId = idFromNumber(idOrConversationId); + conversation = await this.readConversation(conversationId); + if (conversation === undefined) throw new Error(`Unknown conversation: ${conversationId}`); + } + let stored = this.entries.get(id); + if (stored === undefined) { + const row = ( + await this.pool.query( + `SELECT record::text AS record, commit_seq FROM ${this.tables.entries} WHERE session_id = $1 AND id = $2`, + [this.session, id], + ) + ).rows[0]; + if (row === undefined) return undefined; + stored = { record: row.record, commitSeq: seqFromNumber(Number(row.commit_seq)) }; + this.entries.set(id, stored); + } + const entry = parseJson(stored.record); + if (conversation !== undefined) { + let upperEntryId = Number.POSITIVE_INFINITY; + while (conversation.id !== entry.conversationId) { + if (conversation.parent === undefined) return undefined; + upperEntryId = Math.min(upperEntryId, conversation.parent.at); + conversation = (await this.readConversation(conversation.parent.conversationId))!; + } + if (entry.id > upperEntryId) return undefined; + } + return { entry, commitSeq: stored.commitSeq }; + } + + private async readLatestHeadMarker( + conversationId: ConversationId, + atOrBeforeEntryId: EntryId | undefined, + ): Promise<(EntryRecord & { readonly head: EntryId }) | undefined> { + let conversation = await this.readConversation(conversationId); + if (conversation === undefined) throw new Error(`Unknown conversation: ${conversationId}`); + let upper = atOrBeforeEntryId ?? Number.POSITIVE_INFINITY; + while (true) { + const ids = await this.headMarkerIds(conversation.id); + const index = upperBound(ids, upper) - 1; + if (index >= 0) + return parseJson(this.entries.get(ids[index])!.record); + if (conversation.parent === undefined) return undefined; + upper = Math.min(upper, conversation.parent.at); + conversation = (await this.readConversation(conversation.parent.conversationId))!; + } + } + + /** Every head marker of one conversation, ascending; loaded once, extended by commits. */ + private async headMarkerIds(conversationId: ConversationId): Promise { + let ids = this.headMarkers.get(conversationId); + if (ids !== undefined) return ids; + const rows = ( + await this.pool.query( + `SELECT record::text AS record, commit_seq FROM ${this.tables.entries} + WHERE session_id = $1 AND conversation_id = $2 AND head IS NOT NULL ORDER BY id`, + [this.session, conversationId], + ) + ).rows; + ids = []; + for (const row of rows) { + const entry = parseJson(row.record); + this.entries.set(entry.id, { record: row.record, commitSeq: seqFromNumber(Number(row.commit_seq)) }); + ids.push(entry.id); + } + this.headMarkers.set(conversationId, ids); + return ids; + } + + private async readEntries( + query: EntryQuery, + limit: number, + cursor: Cursor | undefined, + ): Promise> { + const { order, after } = scanStart(query.order, cursor, "descending"); + if (order === "ascending") return this.readEntriesAscending(query, limit, after); + let conversation = await this.readConversation(query.conversationId); + if (conversation === undefined) throw new StorageRequestError(`Unknown conversation: ${query.conversationId}`); + let upper: number | undefined = query.maxEntryId; + if (after !== undefined) upper = Math.min(upper ?? Number.MAX_SAFE_INTEGER, after - 1); + const values: EntryRecord[] = []; + while (true) { + const rows = await this.entriesBelow(conversation.id, query.minEntryId, upper, limit + 1 - values.length); + values.push(...rows); + if (values.length > limit || conversation.parent === undefined) break; + upper = upper === undefined ? conversation.parent.at : Math.min(upper, conversation.parent.at); + if (query.minEntryId !== undefined && upper < query.minEntryId) break; + conversation = (await this.readConversation(conversation.parent.conversationId))!; + } + return page(values, limit, "descending"); + } + + /** + * Oldest first: the fork chain's segments from the root conversation forward, each up to its fork point. Read + * from the database; the entry window serves the newest-first scans the harness runs on every request. + */ + private async readEntriesAscending( + query: EntryQuery, + limit: number, + after: number | undefined, + ): Promise> { + const segments: { readonly conversationId: ConversationId; readonly upper: number | undefined }[] = []; + let conversation = await this.readConversation(query.conversationId); + if (conversation === undefined) throw new StorageRequestError(`Unknown conversation: ${query.conversationId}`); + let upper: number | undefined = query.maxEntryId; + while (true) { + segments.push({ conversationId: conversation.id, upper }); + if (conversation.parent === undefined) break; + upper = upper === undefined ? conversation.parent.at : Math.min(upper, conversation.parent.at); + if (query.minEntryId !== undefined && upper < query.minEntryId) break; + conversation = (await this.readConversation(conversation.parent.conversationId))!; + } + let lower: number | undefined = query.minEntryId; + if (after !== undefined) lower = Math.max(lower ?? after + 1, after + 1); + const values: EntryRecord[] = []; + for (const segment of segments.reverse()) { + const params: PostgresValue[] = [this.session, segment.conversationId]; + const clauses = ["session_id = $1", "conversation_id = $2"]; + if (lower !== undefined) clauses.push(`id >= $${params.push(lower)}`); + if (segment.upper !== undefined) clauses.push(`id <= $${params.push(segment.upper)}`); + const rows = await this.records( + `SELECT record::text AS record FROM ${this.tables.entries} WHERE ${clauses.join(" AND ")} + ORDER BY id ASC LIMIT $${params.push(limit + 1 - values.length)}`, + params, + ); + values.push(...rows.map(parseJson)); + if (values.length > limit) break; + } + return page(values, limit, "ascending"); + } + + /** + * Up to `limit` of one conversation's own entries with ids in `[min, upper]`, newest first. Served from the + * conversation's window when it covers the range; otherwise the part below the window is read, and extends + * the window when it is contiguous with it. + */ + private async entriesBelow( + conversationId: ConversationId, + min: number | undefined, + upper: number | undefined, + limit: number, + ): Promise { + const window = await this.entryWindow(conversationId); + const out: EntryRecord[] = []; + let index = upper === undefined ? window.ids.length - 1 : upperBound(window.ids, upper) - 1; + for (; index >= 0 && out.length < limit; index--) { + const id = window.ids[index]; + if (min !== undefined && id < min) return out; + out.push(parseJson(this.entries.get(id)!.record)); + } + if (out.length >= limit || window.complete) return out; + const oldest = window.ids[0]; + const belowWindow = upper === undefined || oldest === undefined || upper >= oldest; + const dbUpper = belowWindow && oldest !== undefined ? oldest - 1 : upper; + const params: PostgresValue[] = [this.session, conversationId]; + const clauses = ["session_id = $1", "conversation_id = $2"]; + if (min !== undefined) clauses.push(`id >= $${params.push(min)}`); + if (dbUpper !== undefined) clauses.push(`id <= $${params.push(dbUpper)}`); + const wanted = limit - out.length; + const rows = ( + await this.pool.query( + `SELECT record::text AS record, commit_seq FROM ${this.tables.entries} WHERE ${clauses.join(" AND ")} + ORDER BY id DESC LIMIT $${params.push(wanted)}`, + params, + ) + ).rows; + const fetched = rows.map((row) => ({ entry: parseJson(row.record), row })); + if (belowWindow) { + // Contiguous with the window: keep these so a later scan of the same range needs no read. Fewer rows + // than asked for, with no lower bound cutting them off, means the conversation starts here. + for (const { entry, row } of fetched) { + this.entries.set(entry.id, { record: row.record, commitSeq: seqFromNumber(Number(row.commit_seq)) }); + } + window.ids.unshift(...fetched.map(({ entry }) => entry.id).reverse()); + if (fetched.length < wanted && min === undefined) window.complete = true; + } + out.push(...fetched.map(({ entry }) => entry)); + return out; + } + + /** The conversation's entry window, loaded from its newest entries on first use. */ + private async entryWindow(conversationId: ConversationId): Promise { + let window = this.entryWindows.get(conversationId); + if (window !== undefined) return window; + const rows = ( + await this.pool.query( + `SELECT record::text AS record, commit_seq FROM ${this.tables.entries} + WHERE session_id = $1 AND conversation_id = $2 ORDER BY id DESC LIMIT $3`, + [this.session, conversationId, ENTRY_WINDOW_LOAD + 1], + ) + ).rows; + const ids: EntryId[] = []; + for (const row of rows.slice(0, ENTRY_WINDOW_LOAD).reverse()) { + const entry = parseJson(row.record); + this.entries.set(entry.id, { record: row.record, commitSeq: seqFromNumber(Number(row.commit_seq)) }); + ids.push(entry.id); + } + window = { ids, complete: rows.length <= ENTRY_WINDOW_LOAD }; + this.entryWindows.set(conversationId, window); + return window; + } + + /** Reflects a committed batch in the caches. Called only after the write statement reported success. */ + private absorb( + writes: readonly StorageWrite[], + actions: ReadonlyMap, + documents: DocumentWrites, + seq: Seq, + ): void { + for (const write of writes) { + switch (write.type) { + case "conversation": + this.conversations.set(write.value.id, parseJson(encodeJson(write.value))); + break; + case "task": + this.tasks.set(write.value.id, encodeJson(write.value)); + break; + case "submission": + this.rememberSubmission(encodeJson(write.value)); + break; + case "entry": { + const record = encodeJson(write.value); + this.entries.set(write.value.id, { record, commitSeq: seq }); + const window = this.entryWindows.get(write.value.conversationId); + if (window !== undefined) { + window.ids.push(write.value.id); + if (window.ids.length > ENTRY_WINDOW_MAX) { + const heads = new Set(this.headMarkers.get(write.value.conversationId) ?? []); + for (const id of window.ids.splice(0, window.ids.length - ENTRY_WINDOW_MAX)) { + if (!heads.has(id)) this.entries.delete(id); + } + window.complete = false; + } + } + if (write.value.head !== undefined) + this.headMarkers.get(write.value.conversationId)?.push(write.value.id); + break; + } + default: + break; + } + } + if (actions.size === 0) return; + for (const [id, { created, retired, content }] of documents.applied) { + if (retired) { + this.documentStates.delete(id); + continue; + } + const state = this.documentStates.get(id); + const record = created !== undefined ? encodeJson(created) : state?.record; + if (content?.kind === "base" && record !== undefined) { + this.documentStates.set(id, { + record, + version: content.version, + value: parseJson(encodeJson(content.value)), + deltasSinceBase: 0, + }); + } else if (content?.kind === "delta" && state !== undefined && state.version === content.version) { + state.value = apply(state.value, parseJson(encodeJson(content.ops))) as JsonObject; + state.deltasSinceBase++; + } else if (content !== undefined) { + this.documentStates.delete(id); + } + } + // Creating or retiring a document changes what an address or scope scan in its scope resolves to. A + // retirement of a document whose scope is not known here invalidates every scope. + for (const [id, { created, retired }] of documents.applied) { + if (created === undefined && !retired) continue; + const scope = created !== undefined ? scopeColumns(created.scope) : this.cachedDocumentScope(id); + if (scope === undefined) { + this.documentScans = new Map(); + this.documentAddresses = new Map(); + break; + } + const scanPrefix = `${scope.scopeKind}:${scope.ownerId}:`; + for (const key of [...this.documentScans.keys()]) + if (key.startsWith(scanPrefix)) this.documentScans.delete(key); + const addressMark = `,"${scope.scopeKind}",${scope.ownerId},`; + for (const key of [...this.documentAddresses.keys()]) + if (key.includes(addressMark)) this.documentAddresses.delete(key); + } + } + + /** The scope of a document this storage has a record for, from its state or address caches. */ + private cachedDocumentScope(id: DocumentId): ScopeColumns | undefined { + const state = this.documentStates.get(id); + if (state !== undefined) return scopeColumns(parseJson(state.record).scope); + for (const record of this.documentAddresses.values()) { + const parsed = parseJson(record); + if (parsed.id === id) return scopeColumns(parsed.scope); + } + return undefined; + } + + private rememberSubmission(record: string): void { + const submission = parseJson(record); + this.submissions.set(submission.id, record); + if (submission.requestId !== undefined) { + this.submissionsByRequest.set(requestKey(submission.conversationId, submission.requestId), submission.id); + } + } + + task(id: TaskId, _context: Context): Promise { + return this.admit(async () => { + const cached = this.tasks.get(id); + if (cached !== undefined) return parseJson(cached); + const [record] = await this.records( + `SELECT record::text AS record FROM ${this.tables.tasks} WHERE session_id = $1 AND id = $2`, + [this.session, id], + ); + if (record !== undefined) this.tasks.set(id, record); + return record === undefined ? undefined : parseJson(record); + }); + } + + scanTasks( + query: TaskQuery, + limit: number, + cursor: Cursor | undefined, + _context: Context, + ): Promise> { + return this.admit(async () => { + const start = scanStart(query.order, cursor, "ascending"); + const scan = scanSql(start); + const params: PostgresValue[] = [this.session]; + const clauses = ["session_id = $1", `id ${scan.comparator} $${params.push(scan.param)}`]; + if (query.conversationId !== undefined) + clauses.push(`conversation_id = $${params.push(query.conversationId)}`); + if (query.kind !== undefined) clauses.push(`kind = $${params.push(encodeIndexedString(query.kind))}`); + if (query.status !== undefined) clauses.push(`status = $${params.push(query.status)}`); + if (query.abortRequested !== undefined) { + clauses.push(`abort_requested = $${params.push(query.abortRequested)}`); + } + if (query.background !== undefined) clauses.push(`background = $${params.push(query.background)}`); + const rows = await this.records( + `SELECT record::text AS record FROM ${this.tables.tasks} WHERE ${clauses.join(" AND ")} + ORDER BY id ${scan.direction} LIMIT $${params.push(limit + 1)}`, + params, + ); + return page(rows.map(parseJson), limit, start.order); + }); + } + + submission(id: SubmissionId, _context: Context): Promise { + return this.admit(async () => { + const cached = this.submissions.get(id); + if (cached !== undefined) return parseJson(cached); + const [record] = await this.records( + `SELECT record::text AS record FROM ${this.tables.submissions} WHERE session_id = $1 AND id = $2`, + [this.session, id], + ); + if (record !== undefined) this.rememberSubmission(record); + return record === undefined ? undefined : parseJson(record); + }); + } + + scanSubmissions( + query: SubmissionQuery, + limit: number, + cursor: Cursor | undefined, + _context: Context, + ): Promise> { + return this.admit(async () => { + const start = scanStart(query.order, cursor, "ascending"); + const scan = scanSql(start); + const params: PostgresValue[] = [this.session]; + const clauses = ["session_id = $1", `id ${scan.comparator} $${params.push(scan.param)}`]; + if (query.conversationId !== undefined) + clauses.push(`conversation_id = $${params.push(query.conversationId)}`); + if (query.status !== undefined) clauses.push(`status = $${params.push(query.status)}`); + const rows = await this.records( + `SELECT record::text AS record FROM ${this.tables.submissions} WHERE ${clauses.join(" AND ")} + ORDER BY id ${scan.direction} LIMIT $${params.push(limit + 1)}`, + params, + ); + return page(rows.map(parseJson), limit, start.order); + }); + } + + submissionByRequest( + conversationId: ConversationId, + requestId: string, + _context: Context, + ): Promise { + return this.admit(async () => { + const key = requestKey(conversationId, requestId); + const cachedId = this.submissionsByRequest.get(key); + if (cachedId === null) return undefined; + const cached = cachedId === undefined ? undefined : this.submissions.get(cachedId); + if (cached !== undefined) return parseJson(cached); + const [record] = await this.records( + `SELECT record::text AS record FROM ${this.tables.submissions} + WHERE session_id = $1 AND conversation_id = $2 AND request_id = $3`, + [this.session, conversationId, encodeIndexedString(requestId)], + ); + if (record !== undefined) this.rememberSubmission(record); + else this.submissionsByRequest.set(key, null); + return record === undefined ? undefined : parseJson(record); + }); + } + + findDocument(address: DocumentAddress, at: DocumentPoint, _context: Context): Promise { + return this.admit(async () => { + const key = at === "current" ? addressKey(address) : undefined; + if (key !== undefined) { + const cached = this.documentAddresses.get(key); + if (cached !== undefined) return parseJson(cached); + } + const parts = addressParts(address); + const params: PostgresValue[] = [ + this.session, + parts.kind, + parts.scopeKind, + parts.ownerId, + parts.family, + parts.keyValue, + ]; + const alive = + at === "current" + ? "retired_at IS NULL" + : `created_at <= $${params.push(at)} AND (retired_at IS NULL OR retired_at > $${params.length})`; + const [record] = await this.records( + `SELECT record::text AS record FROM ${this.tables.documents} + WHERE session_id = $1 AND kind = $2 AND scope_kind = $3 AND owner_id = $4 AND family = $5 + AND key_value = $6 AND ${alive} + ORDER BY created_at DESC LIMIT 1`, + params, + ); + if (key !== undefined && record !== undefined) this.documentAddresses.set(key, record); + return record === undefined ? undefined : parseJson(record); + }); + } + + document(id: DocumentId, at: DocumentPoint, _context: Context): Promise { + return this.admit(async () => { + if (at === "current") { + const state = this.documentStates.get(id); + if (state !== undefined) { + return { + record: parseJson(state.record), + version: state.version, + value: parseJson(encodeJson(state.value)), + deltasSinceBase: state.deltasSinceBase, + }; + } + } + const stored = await this.materializeDocument(id, at); + if (at === "current" && stored !== undefined) { + this.documentStates.set(id, { + record: encodeJson(stored.record), + version: stored.version, + value: parseJson(encodeJson(stored.value)), + deltasSinceBase: stored.deltasSinceBase, + }); + } + return stored; + }); + } + + scanDocuments( + query: DocumentQuery, + limit: number, + cursor: Cursor | undefined, + _context: Context, + ): Promise> { + return this.admit(async () => { + const scope = scopeColumns(query.scope); + const after = cursorId(cursor) ?? -1; + const key = query.at === "current" ? `${scope.scopeKind}:${scope.ownerId}:${query.kind ?? ""}` : undefined; + if (key !== undefined) { + const cached = this.documentScans.get(key); + if (cached !== undefined) { + const records = cached.map(parseJson).filter((record) => record.id > after); + return page(records.slice(0, limit + 1), limit); + } + } + const params: PostgresValue[] = [this.session, scope.scopeKind, scope.ownerId, after]; + const clauses = ["session_id = $1", "scope_kind = $2", "owner_id = $3", "id > $4"]; + if (query.kind !== undefined) clauses.push(`kind = $${params.push(encodeIndexedString(query.kind))}`); + if (query.at === "current") { + clauses.push("retired_at IS NULL"); + } else { + clauses.push( + `created_at <= $${params.push(query.at)}`, + `(retired_at IS NULL OR retired_at > $${params.length})`, + ); + } + // A current scan is cached whole: it is read from its start without a page limit, and a scope's live + // documents are few. + const whole = key !== undefined && after === -1; + const rows = await this.records( + `SELECT record::text AS record FROM ${this.tables.documents} WHERE ${clauses.join(" AND ")} + ORDER BY id${whole ? "" : ` LIMIT $${params.push(limit + 1)}`}`, + params, + ); + if (whole) this.documentScans.set(key!, rows); + return page(rows.slice(0, limit + 1).map(parseJson), limit); + }); + } + + /** Wait for admitted operations. The pool stays open. */ + close(_context: Context): Promise { + if (this.closing === undefined) { + this.closed = true; + this.closing = this.release(); + } + return this.closing; + } + + private async release(): Promise { + if (this.admitted > 0) { + await new Promise((resolve) => { + this.drained = resolve; + }); + } + } + + /** Run one storage operation. Close waits for admitted operations, so it resolves after every query and commit that started while the storage was open. */ + private async admit(operation: () => Promise): Promise { + this.assertOpen(); + this.admitted++; + try { + return await operation(); + } finally { + if (--this.admitted === 0) this.drained?.(); + } + } + + private loseOwnership(error: Error): void { + this.lost ??= error; + } + + private async records(sql: string, params: PostgresValue[], executor: PostgresQueryable = this.pool) { + return (await executor.query(sql, params)).rows.map((row) => row.record); + } + + private async recordById(table: string, id: Id, executor?: PostgresQueryable): Promise { + const [record] = await this.records( + `SELECT record::text AS record FROM ${table} WHERE session_id = $1 AND id = $2`, + [this.session, id], + executor, + ); + return record === undefined ? undefined : parseJson(record); + } + + private async readConversation(id: ConversationId): Promise { + const cached = this.conversations.get(id); + if (cached !== undefined) return cached; + const record = await this.recordById(this.tables.conversations, id); + if (record !== undefined) this.conversations.set(id, record); + return record; + } + + /** One statement reads the record and the revisions from the newest base at or before `at`. */ + private async materializeDocument(id: DocumentId, at: DocumentPoint): Promise { + const upper = at === "current" ? Number.MAX_SAFE_INTEGER : at; + const row = ( + await this.pool.query<{ readonly record: string | null; readonly revisions: string | null }>( + `SELECT (SELECT x.record::text FROM ${this.tables.documents} x WHERE x.session_id = $1 AND x.id = $2) AS record, + (SELECT json_agg(json_build_object('seq', r.seq, 'kind', r.kind, 'version', r.version, 'content', r.content) + ORDER BY r.seq)::text + FROM ${this.tables.revisions} r WHERE r.session_id = $1 AND r.document_id = $2 AND r.seq <= $3 + AND r.seq >= (SELECT b.seq FROM ${this.tables.revisions} b WHERE b.session_id = $1 AND b.document_id = $2 + AND b.kind = 'base' AND b.seq <= $3 ORDER BY b.seq DESC LIMIT 1)) AS revisions`, + [this.session, id, upper], + ) + ).rows[0]; + if (row?.record === null || row?.record === undefined) return undefined; + const record = parseJson(row.record); + if (at !== "current" && isCurrentOnly(record)) { + throw new Error(`Document ${id} does not retain historical content`); + } + if (!isAliveAt(record, at)) return undefined; + const revisions = + row.revisions === null + ? [] + : parseJson< + readonly { + readonly seq: number; + readonly kind: DocumentContent["kind"]; + readonly version: number; + readonly content: JsonValue; + }[] + >(row.revisions); + const base = revisions[0]; + if (base === undefined || base.kind !== "base") throw new Error(`Document ${id} is missing a required base`); + let value = base.content as JsonObject; + const tail = revisions.slice(1); + for (const revision of tail) { + if (revision.kind !== "delta" || revision.version !== base.version) { + throw new Error(`Document ${id} crosses a stored version boundary without a base`); + } + value = apply(value, revision.content as unknown as readonly Op[]) as JsonObject; + } + return { record, version: base.version, value, deltasSinceBase: tail.length }; + } + + private candidateNextId(writes: readonly StorageWrite[]): number { + let nextId = this.nextId; + for (const write of writes) { + const id = writeId(write); + if (id !== undefined) nextId = Math.max(nextId, id + 1); + } + return nextId; + } + + /** The batch-internal id rules; whether an id already belongs to a stored record is the commit function's check. */ + private checkBatchIds(writes: readonly StorageWrite[]): void { + const claimed = new Map, TableName>(); + for (const write of writes) { + if (write.type === "document.change" || write.type === "document.retire") continue; + const document = write.type === "document.create" || write.type === "document.copy"; + const table: TableName = document ? "document" : write.type; + const id = document ? write.record.id : write.value.id; + const earlier = claimed.get(id); + if (table === "conversation" || table === "entry" || table === "document") { + if (earlier !== undefined) throw new Error(`ID ${id} is written more than once`); + } else if (earlier !== undefined && earlier !== table) { + throw new Error(`ID ${id} is written as two record types`); + } + claimed.set(id, table); + } + } + + private prepareDocumentActions(writes: readonly StorageWrite[]): Map { + const actions = new Map(); + for (const write of writes) { + if ( + write.type !== "document.create" && + write.type !== "document.copy" && + write.type !== "document.change" && + write.type !== "document.retire" + ) { + continue; + } + const id = write.type === "document.create" || write.type === "document.copy" ? write.record.id : write.id; + let action = actions.get(id); + if (action === undefined) { + action = { retire: false }; + actions.set(id, action); + } + switch (write.type) { + case "document.create": + if (action.create !== undefined || action.content !== undefined || action.copy !== undefined) { + throw new Error(`Document ${id} has more than one content command`); + } + action.create = write.record; + action.content = write.content; + break; + case "document.copy": + if (action.create !== undefined || action.content !== undefined || action.copy !== undefined) { + throw new Error(`Document ${id} has more than one content command`); + } + action.create = write.record; + action.copy = write.source; + break; + case "document.change": + if (action.content !== undefined || action.copy !== undefined) { + throw new Error(`Document ${id} has more than one content command`); + } + action.content = write.content; + break; + case "document.retire": + if (action.retire) throw new Error(`Document ${id} is retired more than once`); + action.retire = true; + break; + } + } + return actions; + } + + /** + * Resolves the batch's document actions into the rows the commit function takes, reading copy sources here. + * A created record carries its sequence; an existing document's retirement and any history clearing are the + * function's to apply, since it holds the stored record. + */ + private async prepareDocumentWrites( + actions: ReadonlyMap, + seq: Seq, + ): Promise { + const writes: DocumentWrites = { claims: [], rows: [], applied: new Map() }; + for (const [id, action] of actions) { + if (action.copy !== undefined && actions.has(action.copy.id)) { + throw new StorageRejected(`Document copy ${id} source is changed in the copy batch`); + } + let content = action.content; + if (action.copy !== undefined) { + try { + const stored = await this.materializeDocument(action.copy.id, action.copy.at); + if (stored === undefined) throw new Error(`Fork source document ${action.copy.id} cannot be read`); + const create = action.create!; + if ( + stored.record.scope.kind !== "conversation" || + create.scope.kind !== "conversation" || + stored.record.kind !== create.kind || + stored.record.key !== create.key || + stored.record.history !== create.history || + stored.record.fork !== create.fork + ) { + throw new Error(`Fork source document ${action.copy.id} does not match the copied record`); + } + content = { kind: "base", version: stored.version, value: stored.value }; + } catch (error) { + if (error instanceof StorageRejected) throw error; + throw new StorageRejected(`Document copy ${id} was rejected`, { cause: error }); + } + } + let created: DocumentRecord | undefined; + let parts: ReturnType | undefined; + if (action.create !== undefined) { + created = { ...action.create, createdAt: seq, ...(action.retire ? { retiredAt: seq } : {}) }; + parts = addressParts(created); + writes.claims.push(id); + } + writes.rows.push({ + id, + is_create: created !== undefined, + retire: action.retire, + kind: parts?.kind ?? null, + family: parts?.family ?? null, + key_value: parts?.keyValue ?? null, + scope_kind: parts?.scopeKind ?? null, + owner_id: parts?.ownerId ?? null, + record: created === undefined ? null : encodeJson(created), + content_kind: content?.kind ?? null, + version: content?.version ?? null, + content: content === undefined ? null : encodeJson(content.kind === "base" ? content.value : content.ops), + }); + writes.applied.set(id, { created, retired: action.retire, content }); + } + return writes; + } + + private assertOpen(): void { + if (this.closed) throw new Error("PostgresStorage is closed"); + if (this.lost !== undefined) { + throw new Error(`PostgresStorage lost ownership of session ${this.session}`, { cause: this.lost }); + } + } +} diff --git a/packages/durable/test/postgres-storage.test.ts b/packages/durable/test/postgres-storage.test.ts new file mode 100644 index 000000000000..651344c2433e --- /dev/null +++ b/packages/durable/test/postgres-storage.test.ts @@ -0,0 +1,225 @@ +import { randomUUID } from "node:crypto"; +import type { Context } from "@earendil-works/chord"; +import { BACKGROUND_CONTEXT } from "@earendil-works/chord/context"; +import { registerStorageConformance } from "@earendil-works/pi-durable/testing"; +import pg from "pg"; +import { afterAll, afterEach, beforeAll, describe, expect, it } from "vitest"; +import { idFromNumber } from "../src/ids.ts"; +import { + applyPostgresMigrations, + CURRENT_POSTGRES_SCHEMA_VERSION, + PostgresStorage, +} from "../src/storage/postgres/index.ts"; +import type { ConversationId, EntryId, Id, Seq, Storage, StorageWrite } from "../src/types.ts"; +import { ROOT_CONVERSATION_ID } from "../src/types.ts"; + +// Set to a database the tests may create and drop schemas in, e.g. postgres://postgres:pg@localhost:5432/postgres. +const url = process.env.PI_DURABLE_POSTGRES_URL; +const describePostgres = describe.skipIf(url === undefined); +const context = BACKGROUND_CONTEXT; +const schema = `pi_durable_test_${randomUUID().replaceAll("-", "")}`; +const openStorages = new Set(); +let pool: pg.Pool; + +beforeAll(async () => { + if (url === undefined) return; + pool = new pg.Pool({ connectionString: url, max: 20 }); + await applyPostgresMigrations(pool, schema); +}); + +afterEach(async () => { + for (const storage of openStorages) await storage.close(context); + openStorages.clear(); +}); + +afterAll(async () => { + if (url === undefined) return; + await pool.query(`DROP SCHEMA IF EXISTS "${schema}" CASCADE`); + await pool.end(); +}); + +async function openStorage(session: string = randomUUID()): Promise { + const storage = await PostgresStorage.open(pool, { session, schema }); + openStorages.add(storage); + return storage; +} + +/** Reopens the Session after every commit, so each later call reads what an earlier process persisted. */ +class ReopeningStorage implements Storage { + private current: PostgresStorage; + private readonly session: string; + private closed = false; + + constructor(current: PostgresStorage, session: string) { + this.current = current; + this.session = session; + } + + async commit(writes: readonly StorageWrite[], commitContext: Context): Promise { + if (this.closed) return this.current.commit(writes, commitContext); + try { + return await this.current.commit(writes, commitContext); + } finally { + await this.current.close(context); + this.current = await PostgresStorage.open(pool, { session: this.session, schema }); + } + } + + mintId>(): Promise { + return this.current.mintId(); + } + conversation: Storage["conversation"] = (id, readContext) => this.current.conversation(id, readContext); + scanConversations: Storage["scanConversations"] = (query, limit, cursor, readContext) => + this.current.scanConversations(query, limit, cursor, readContext); + entry(id: EntryId, readContext: Context): ReturnType; + entry(conversationId: ConversationId, id: EntryId, readContext: Context): ReturnType; + entry(idOrConversationId: EntryId | ConversationId, idOrContext: EntryId | Context, readContext?: Context) { + if (readContext === undefined) { + return this.current.entry(idFromNumber(idOrConversationId), idOrContext as Context); + } + if (typeof idOrContext !== "number") throw new TypeError("Storage.entry() requires an entry ID"); + return this.current.entry( + idFromNumber(idOrConversationId), + idFromNumber(idOrContext), + readContext, + ); + } + findLatestHeadMarker: Storage["findLatestHeadMarker"] = (conversationId, at, readContext) => + this.current.findLatestHeadMarker(conversationId, at, readContext); + scanEntries: Storage["scanEntries"] = (query, limit, cursor, readContext) => + this.current.scanEntries(query, limit, cursor, readContext); + task: Storage["task"] = (id, readContext) => this.current.task(id, readContext); + scanTasks: Storage["scanTasks"] = (query, limit, cursor, readContext) => + this.current.scanTasks(query, limit, cursor, readContext); + submission: Storage["submission"] = (id, readContext) => this.current.submission(id, readContext); + scanSubmissions: Storage["scanSubmissions"] = (query, limit, cursor, readContext) => + this.current.scanSubmissions(query, limit, cursor, readContext); + submissionByRequest: Storage["submissionByRequest"] = (conversationId, requestId, readContext) => + this.current.submissionByRequest(conversationId, requestId, readContext); + findDocument: Storage["findDocument"] = (address, at, readContext) => + this.current.findDocument(address, at, readContext); + document: Storage["document"] = (id, at, readContext) => this.current.document(id, at, readContext); + scanDocuments: Storage["scanDocuments"] = (query, limit, cursor, readContext) => + this.current.scanDocuments(query, limit, cursor, readContext); + + async close(closeContext: Context): Promise { + if (this.closed) return; + this.closed = true; + await this.current.close(closeContext); + } +} + +registerStorageConformance({ describe: describePostgres, expect, it }, "PostgresStorage", async (use) => + use(await openStorage()), +); + +registerStorageConformance({ describe: describePostgres, expect, it }, "PostgresStorage across reopen", async (use) => { + const session = randomUUID(); + const storage = new ReopeningStorage(await PostgresStorage.open(pool, { session, schema }), session); + try { + await use(storage); + } finally { + await storage.close(context); + } +}); + +async function createRoot(storage: Storage): Promise { + return storage.commit([{ type: "conversation", value: { id: ROOT_CONVERSATION_ID } }], context); +} + +describePostgres("PostgresStorage ownership", () => { + it("keeps Sessions in one schema independent", async () => { + const first = await openStorage(); + const second = await openStorage(); + await createRoot(first); + expect(await first.mintId()).toBe(2); + expect(await second.mintId()).toBe(2); + expect(await second.conversation(ROOT_CONVERSATION_ID, context)).toBeUndefined(); + await expect(createRoot(second)).resolves.toBe(1); + }); + + it("hands a Session to a later open and rejects the earlier owner's commits", async () => { + const session = randomUUID(); + const first = await openStorage(session); + await createRoot(first); + const second = await openStorage(session); + const entryId = await first.mintId(); + await expect( + first.commit([{ type: "entry", value: { id: entryId, conversationId: ROOT_CONVERSATION_ID, kind: "message" } }], context), + ).rejects.toThrow("was opened by another process"); + await expect(first.mintId()).rejects.toThrow("lost ownership"); + await expect( + second.commit([{ type: "entry", value: { id: entryId, conversationId: ROOT_CONVERSATION_ID, kind: "message" } }], context), + ).resolves.toBe(2); + }); + + it("rejects commits after another process takes ownership", async () => { + const session = randomUUID(); + const storage = await openStorage(session); + await createRoot(storage); + await pool.query(`UPDATE "${schema}".sessions SET owner = 'another-process' WHERE session_id = $1`, [session]); + const entryId = await storage.mintId(); + await expect( + storage.commit( + [{ type: "entry", value: { id: entryId, conversationId: ROOT_CONVERSATION_ID, kind: "message" } }], + context, + ), + ).rejects.toThrow("was opened by another process"); + await expect(storage.conversation(ROOT_CONVERSATION_ID, context)).rejects.toThrow("lost ownership"); + }); + + it("persists the ID and commit sequence across reopen", async () => { + const session = randomUUID(); + const first = await openStorage(session); + await createRoot(first); + const minted = await first.mintId(); + await first.commit( + [{ type: "entry", value: { id: minted, conversationId: ROOT_CONVERSATION_ID, kind: "message" } }], + context, + ); + await first.close(context); + const second = await openStorage(session); + expect(await second.mintId()).toBe(minted + 1); + await expect(createRoot(second)).rejects.toThrow("already belongs to conversation"); + const entryId = await second.mintId(); + await expect( + second.commit( + [{ type: "entry", value: { id: entryId, conversationId: ROOT_CONVERSATION_ID, kind: "message" } }], + context, + ), + ).resolves.toBe(3); + }); +}); + +describePostgres("Postgres migrations", () => { + it("are idempotent and record the current version", async () => { + await applyPostgresMigrations(pool, schema); + const result = await pool.query<{ version: number }>(`SELECT version FROM "${schema}".durable_schema`); + expect(result.rows).toEqual([{ version: CURRENT_POSTGRES_SCHEMA_VERSION }]); + }); + + it("serialize concurrent first opens of a fresh schema", async () => { + const fresh = `pi_durable_test_${randomUUID().replaceAll("-", "")}`; + try { + const storages = await Promise.all( + Array.from({ length: 4 }, () => PostgresStorage.open(pool, { session: randomUUID(), schema: fresh })), + ); + for (const storage of storages) await storage.close(context); + } finally { + await pool.query(`DROP SCHEMA IF EXISTS "${fresh}" CASCADE`); + } + }); + + it("reject a schema newer than this package supports", async () => { + const fresh = `pi_durable_test_${randomUUID().replaceAll("-", "")}`; + try { + await applyPostgresMigrations(pool, fresh); + await pool.query(`UPDATE "${fresh}".durable_schema SET version = 99`); + await expect(PostgresStorage.open(pool, { session: randomUUID(), schema: fresh })).rejects.toThrow( + "is newer than supported version", + ); + } finally { + await pool.query(`DROP SCHEMA IF EXISTS "${fresh}" CASCADE`); + } + }); +}); From 00d83ad62d44c988fb029aae9bbacc9c6f93cd1c Mon Sep 17 00:00:00 2001 From: rgarcia <72655+rgarcia@users.noreply.github.com> Date: Sun, 1 Nov 2026 20:52:30 +0000 Subject: [PATCH 2/2] feat(durable): export the Postgres storage subpath; batched commits write one statement per table --- package-lock.json | 174 ++++++++++++++++ packages/durable/package.json | 9 +- .../durable/src/storage/postgres/database.ts | 6 +- .../src/storage/postgres/migrations.ts | 49 +++++ .../durable/src/storage/postgres/storage.ts | 188 ++++++++++++++---- .../durable/test/postgres-storage.test.ts | 119 ++++++++--- 6 files changed, 472 insertions(+), 73 deletions(-) diff --git a/package-lock.json b/package-lock.json index 74451ff7532a..9373b46ce8d6 100644 --- a/package-lock.json +++ b/package-lock.json @@ -2472,6 +2472,18 @@ "undici-types": "~6.21.0" } }, + "node_modules/@types/pg": { + "version": "8.23.1", + "resolved": "https://registry.npmjs.org/@types/pg/-/pg-8.23.1.tgz", + "integrity": "sha512-fKVHpikPdg4GKks3JuLEhvwSyvwzF23hnabPy6DD8ljVbC7+6J5dQzdv4arV6jqq57djnMgs1HKBxX4P8aBI3A==", + "dev": true, + "license": "MIT", + "dependencies": { + "@types/node": "*", + "pg-protocol": "*", + "pg-types": "^2.2.0" + } + }, "node_modules/@types/proper-lockfile": { "version": "4.1.4", "resolved": "https://registry.npmjs.org/@types/proper-lockfile/-/proper-lockfile-4.1.4.tgz", @@ -4533,6 +4545,103 @@ "dev": true, "license": "MIT" }, + "node_modules/pg": { + "version": "8.23.0", + "resolved": "https://registry.npmjs.org/pg/-/pg-8.23.0.tgz", + "integrity": "sha512-Ip2EQCngowJLGOfCwkFhPXU7/ljlhn6Rxlmy4XYfL2Y+vyRM59+8uR2xqRWKdYmbXmxCFOAmKxBuSUCdF34qLg==", + "dev": true, + "license": "MIT", + "dependencies": { + "pg-connection-string": "^2.14.0", + "pg-pool": "^3.14.0", + "pg-protocol": "^1.16.0", + "pg-types": "2.2.0", + "pgpass": "1.0.5" + }, + "engines": { + "node": ">= 16.0.0" + }, + "optionalDependencies": { + "pg-cloudflare": "^1.4.0" + }, + "peerDependencies": { + "pg-native": ">=3.0.1" + }, + "peerDependenciesMeta": { + "pg-native": { + "optional": true + } + } + }, + "node_modules/pg-cloudflare": { + "version": "1.4.1", + "resolved": "https://registry.npmjs.org/pg-cloudflare/-/pg-cloudflare-1.4.1.tgz", + "integrity": "sha512-6PQbsFWZcp9EmJEwy5cGQ2La+AMWpP46lgbb8X+U/XsHIUweYDNCpeuKck5RxL2MdVFi7krbbEi5nX4Zh7JhrQ==", + "dev": true, + "license": "MIT", + "optional": true + }, + "node_modules/pg-connection-string": { + "version": "2.14.1", + "resolved": "https://registry.npmjs.org/pg-connection-string/-/pg-connection-string-2.14.1.tgz", + "integrity": "sha512-qR3kGNPBLpCNtz0evbKA0Y/MRFXwSSdT+pTJvYp/bXTcReZbvX1kzF0IyTc1QnxqF7AZbOeBhNL8R5mYQZV/MA==", + "dev": true, + "license": "MIT" + }, + "node_modules/pg-int8": { + "version": "1.0.1", + "resolved": "https://registry.npmjs.org/pg-int8/-/pg-int8-1.0.1.tgz", + "integrity": "sha512-WCtabS6t3c8SkpDBUlb1kjOs7l66xsGdKpIPZsg4wR+B3+u9UAum2odSsF9tnvxg80h4ZxLWMy4pRjOsFIqQpw==", + "dev": true, + "license": "ISC", + "engines": { + "node": ">=4.0.0" + } + }, + "node_modules/pg-pool": { + "version": "3.14.0", + "resolved": "https://registry.npmjs.org/pg-pool/-/pg-pool-3.14.0.tgz", + "integrity": "sha512-gKtPkFdQPU3DksooVLi9LsjZxrsBUZIpa+7aVx+LV5pNh0KzP4Zleud2po+ConrxbuXGBJ6Hfer6hdgpIBpBaw==", + "dev": true, + "license": "MIT", + "peerDependencies": { + "pg": ">=8.0" + } + }, + "node_modules/pg-protocol": { + "version": "1.16.1", + "resolved": "https://registry.npmjs.org/pg-protocol/-/pg-protocol-1.16.1.tgz", + "integrity": "sha512-p9VOFMiHB/ZbJATetbg+99PxssTVSQRnyuPSQ67mN1+1KBOjZaZ83ZQzltnxPhJwSsC3nwVjJ10DVJlerbFzLg==", + "dev": true, + "license": "MIT" + }, + "node_modules/pg-types": { + "version": "2.2.0", + "resolved": "https://registry.npmjs.org/pg-types/-/pg-types-2.2.0.tgz", + "integrity": "sha512-qTAAlrEsl8s4OiEQY69wDvcMIdQN6wdz5ojQiOy6YRMuynxenON0O5oCpJI6lshc6scgAY8qvJ2On/p+CXY0GA==", + "dev": true, + "license": "MIT", + "dependencies": { + "pg-int8": "1.0.1", + "postgres-array": "~2.0.0", + "postgres-bytea": "~1.0.0", + "postgres-date": "~1.0.4", + "postgres-interval": "^1.1.0" + }, + "engines": { + "node": ">=4" + } + }, + "node_modules/pgpass": { + "version": "1.0.5", + "resolved": "https://registry.npmjs.org/pgpass/-/pgpass-1.0.5.tgz", + "integrity": "sha512-FdW9r/jQZhSeohs1Z3sI1yxFQNFvMcnmfuj4WBMUTxOrAyLMaTcE1aAMBiTlbMNaXvBCQuVi0R7hd8udDSP7ug==", + "dev": true, + "license": "MIT", + "dependencies": { + "split2": "^4.1.0" + } + }, "node_modules/pi-extension-custom-provider-anthropic": { "resolved": "packages/coding-agent/examples/extensions/custom-provider-anthropic", "link": true @@ -4621,6 +4730,49 @@ "node": "^10 || ^12 || >=14" } }, + "node_modules/postgres-array": { + "version": "2.0.0", + "resolved": "https://registry.npmjs.org/postgres-array/-/postgres-array-2.0.0.tgz", + "integrity": "sha512-VpZrUqU5A69eQyW2c5CA1jtLecCsN2U/bD6VilrFDWq5+5UIEVO7nazS3TEcHf1zuPYO/sqGvUvW62g86RXZuA==", + "dev": true, + "license": "MIT", + "engines": { + "node": ">=4" + } + }, + "node_modules/postgres-bytea": { + "version": "1.0.1", + "resolved": "https://registry.npmjs.org/postgres-bytea/-/postgres-bytea-1.0.1.tgz", + "integrity": "sha512-5+5HqXnsZPE65IJZSMkZtURARZelel2oXUEO8rH83VS/hxH5vv1uHquPg5wZs8yMAfdv971IU+kcPUczi7NVBQ==", + "dev": true, + "license": "MIT", + "engines": { + "node": ">=0.10.0" + } + }, + "node_modules/postgres-date": { + "version": "1.0.7", + "resolved": "https://registry.npmjs.org/postgres-date/-/postgres-date-1.0.7.tgz", + "integrity": "sha512-suDmjLVQg78nMK2UZ454hAG+OAW+HQPZ6n++TNDUX+L0+uUlLywnoxJKDou51Zm+zTCjrCl0Nq6J9C5hP9vK/Q==", + "dev": true, + "license": "MIT", + "engines": { + "node": ">=0.10.0" + } + }, + "node_modules/postgres-interval": { + "version": "1.2.0", + "resolved": "https://registry.npmjs.org/postgres-interval/-/postgres-interval-1.2.0.tgz", + "integrity": "sha512-9ZhXKM/rw350N1ovuWHbGxnGh/SNJ4cnxHiM0rxE4VN41wsg8P8zWn9hv/buK00RP4WvlOyr/RBDiptyxVbkZQ==", + "dev": true, + "license": "MIT", + "dependencies": { + "xtend": "^4.0.0" + }, + "engines": { + "node": ">=0.10.0" + } + }, "node_modules/prebuild-install": { "version": "7.1.3", "resolved": "https://registry.npmjs.org/prebuild-install/-/prebuild-install-7.1.3.tgz", @@ -5060,6 +5212,16 @@ "node": ">=0.10.0" } }, + "node_modules/split2": { + "version": "4.2.0", + "resolved": "https://registry.npmjs.org/split2/-/split2-4.2.0.tgz", + "integrity": "sha512-UcjcJOWknrNkF6PLX83qcHM6KHgVKNkV62Y8a5uYDVv9ydGQVwAHMKqHdJje1VTWpljG0WYpCDhrCdAOYH4TWg==", + "dev": true, + "license": "ISC", + "engines": { + "node": ">= 10.x" + } + }, "node_modules/ssh2": { "version": "1.17.0", "resolved": "https://registry.npmjs.org/ssh2/-/ssh2-1.17.0.tgz", @@ -5647,6 +5809,16 @@ } } }, + "node_modules/xtend": { + "version": "4.0.2", + "resolved": "https://registry.npmjs.org/xtend/-/xtend-4.0.2.tgz", + "integrity": "sha512-LKYU1iAXJXUgAXn9URjiu+MWhyUXHsvfp7mcuYm9dSUKK0/CjtrUwFAxD82/mCWbtLsGjFIad0wIsod4zrTAEQ==", + "dev": true, + "license": "MIT", + "engines": { + "node": ">=0.4" + } + }, "node_modules/yaml": { "version": "2.9.0", "resolved": "https://registry.npmjs.org/yaml/-/yaml-2.9.0.tgz", @@ -5979,7 +6151,9 @@ }, "devDependencies": { "@cloudflare/vitest-plugin": "1.4.4", + "@types/pg": "8.23.1", "fake-indexeddb": "6.2.6", + "pg": "8.23.0", "shx": "0.3.4", "vitest": "4.1.12" }, diff --git a/packages/durable/package.json b/packages/durable/package.json index 933fe6db75bb..e100698ea5f9 100644 --- a/packages/durable/package.json +++ b/packages/durable/package.json @@ -114,7 +114,12 @@ "import": "./dist/testing/index.js", "default": "./dist/testing/index.js" }, - "./package.json": "./package.json" + "./package.json": "./package.json", + "./storage/postgres": { + "source": "./src/storage/postgres/index.ts", + "types": "./dist/storage/postgres/index.d.ts", + "import": "./dist/storage/postgres/index.js" + } }, "sideEffects": false, "files": [ @@ -168,7 +173,9 @@ }, "devDependencies": { "@cloudflare/vitest-plugin": "1.4.4", + "@types/pg": "8.23.1", "fake-indexeddb": "6.2.6", + "pg": "8.23.0", "shx": "0.3.4", "vitest": "4.1.12" } diff --git a/packages/durable/src/storage/postgres/database.ts b/packages/durable/src/storage/postgres/database.ts index 19e61c43ebb9..1402fba84978 100644 --- a/packages/durable/src/storage/postgres/database.ts +++ b/packages/durable/src/storage/postgres/database.ts @@ -20,10 +20,8 @@ export interface PostgresPoolClient extends PostgresQueryable { /** * The connection pool `PostgresStorage` runs on. A `pg.Pool` satisfies it structurally, so this package - * never imports a driver. The pool belongs to the caller: closing the storage does not end it. - * - * The storage keeps one client checked out for its whole lifetime to hold the session lock, so the pool - * needs at least two connections. + * never imports a driver. The pool belongs to the caller: closing the storage does not end it, and an open + * storage holds no connection between operations. */ export interface PostgresPool extends PostgresQueryable { connect(): Promise; diff --git a/packages/durable/src/storage/postgres/migrations.ts b/packages/durable/src/storage/postgres/migrations.ts index 546e51df4d62..21f2f77744cd 100644 --- a/packages/durable/src/storage/postgres/migrations.ts +++ b/packages/durable/src/storage/postgres/migrations.ts @@ -260,10 +260,59 @@ BEGIN END $fn$`; +/** + * Stores a group of batches in one call. The session row is locked and its owner checked once; then `commit_batch` + * runs per batch in a subtransaction, so a refusal by one of the adapter's checks (DU003 to DU009) keeps the batches + * before it and ends the call with the refusal, while any other error, ownership included, raises and rolls back + * everything the call wrote. + */ +const COMMIT_BATCHES_FUNCTION = `CREATE OR REPLACE FUNCTION {schema}.commit_batches( + p_session TEXT, p_owner TEXT, p_seq BIGINT, p_next_id BIGINT, p_batches JSON +) RETURNS JSON LANGUAGE plpgsql AS $fn$ +DECLARE + v_owner TEXT; + v_seq BIGINT; + v_batch JSON; + v_stored INTEGER := 0; + v_code TEXT; + v_message TEXT; +BEGIN + SELECT owner, next_seq INTO v_owner, v_seq FROM {schema}.sessions WHERE session_id = p_session FOR UPDATE; + IF NOT FOUND THEN + RAISE EXCEPTION 'Durable Postgres session metadata is missing' USING ERRCODE = 'DU001'; + END IF; + IF v_owner <> p_owner OR v_seq <> p_seq THEN + RAISE EXCEPTION 'Durable Postgres session % was opened by another process', p_session USING ERRCODE = 'DU002'; + END IF; + FOR v_batch IN SELECT value FROM json_array_elements(p_batches) LOOP + BEGIN + PERFORM {schema}.commit_batch(p_session, p_owner, p_seq + v_stored, p_next_id, v_batch); + EXCEPTION + WHEN SQLSTATE 'DU003' OR SQLSTATE 'DU004' OR SQLSTATE 'DU005' OR SQLSTATE 'DU006' + OR SQLSTATE 'DU007' OR SQLSTATE 'DU008' OR SQLSTATE 'DU009' THEN + GET STACKED DIAGNOSTICS v_code = RETURNED_SQLSTATE, v_message = MESSAGE_TEXT; + RETURN json_build_object('stored', v_stored, 'code', v_code, 'message', v_message); + END; + v_stored := v_stored + 1; + END LOOP; + RETURN json_build_object('stored', v_stored); +END +$fn$`; + /** Immutable, ordered schema history. Append new migrations after the initial schema ships. */ export const POSTGRES_MIGRATIONS: readonly PostgresMigration[] = [ { version: 1, statements: INITIAL_SCHEMA }, { version: 2, statements: [COMMIT_FUNCTION] }, + // Version 3 adds the `sleeping` task status and the grouped commit function. + { + version: 3, + statements: [ + "ALTER TABLE {schema}.tasks DROP CONSTRAINT tasks_status_check", + `ALTER TABLE {schema}.tasks ADD CONSTRAINT tasks_status_check + CHECK (status IN ('pending', 'running', 'waiting', 'sleeping', 'completing', 'terminal'))`, + COMMIT_BATCHES_FUNCTION, + ], + }, ]; export const CURRENT_POSTGRES_SCHEMA_VERSION = POSTGRES_MIGRATIONS.at(-1)?.version ?? 0; diff --git a/packages/durable/src/storage/postgres/storage.ts b/packages/durable/src/storage/postgres/storage.ts index e95732a8b62e..fc05f1783083 100644 --- a/packages/durable/src/storage/postgres/storage.ts +++ b/packages/durable/src/storage/postgres/storage.ts @@ -1,6 +1,6 @@ import type { Context, JsonValue } from "@earendil-works/chord"; import { apply, type Op } from "@earendil-works/chord/delta"; -/** Thrown when the Session was opened by another process; any error from `commit` fails the Session. */ +/** A refusal by one of the adapter's own checks while preparing a batch: the batches before it stand. */ export class StorageRejected extends Error { constructor(message: string, options?: ErrorOptions) { super(message, options); @@ -8,9 +8,18 @@ export class StorageRejected extends Error { } } +/** A database request failed while a batch was being prepared: an unknown outcome, never a refusal. */ +class PostgresUnavailable extends Error { + constructor(cause: unknown) { + super("Postgres request failed", { cause }); + this.name = "PostgresUnavailable"; + } +} + import { StorageRequestError } from "../../errors.ts"; import { idFromNumber, seqFromNumber } from "../../ids.ts"; import type { + CommitResult, ConversationId, ConversationQuery, ConversationRecord, @@ -52,6 +61,15 @@ type MetadataRow = { readonly next_id: string; readonly next_seq: string }; /** A conversation's own entries over one contiguous id range ending at its newest entry, newest-last. */ type EntryWindow = { ids: EntryId[]; complete: boolean }; type DocumentState = { readonly record: string; readonly version: number; value: JsonObject; deltasSinceBase: number }; +/** One batch after the adapter's own checks, ready for the commit function. */ +type PreparedBatch = { + readonly writes: readonly StorageWrite[]; + readonly actions: Map; + readonly documents: DocumentWrites; + readonly encoded: string; + readonly candidateNextId: number; +}; +type CommitBatchesResult = { readonly stored: number; readonly code?: string; readonly message?: string }; type DocumentWrites = { readonly claims: DocumentId[]; readonly rows: unknown[]; @@ -66,7 +84,10 @@ type DocumentWrites = { >; }; -/** SQLSTATEs the commit function raises, each for one adapter error. */ +/** + * SQLSTATEs the commit functions raise, each for one adapter error. DU001 and DU002 reject the call; `commit_batches` + * catches the others per batch and returns them as the refusal. + */ const commitErrorCodes: Readonly> = { DU001: true, DU002: true, @@ -263,7 +284,7 @@ export class PostgresStorage implements Storage { this.nextId = options.nextId; this.nextSeq = options.nextSeq; const schema = quoteIdentifier(options.schema); - this.commitFunction = `${schema}.commit_batch`; + this.commitFunction = `${schema}.commit_batches`; this.tables = { sessions: `${schema}.sessions`, recordIds: `${schema}.record_ids`, @@ -301,33 +322,104 @@ export class PostgresStorage implements Storage { }); } - commit(writes: readonly StorageWrite[], _context: Context): Promise { + /** + * Stores `batches` in as few round trips as their content allows. Batches are prepared here (the adapter's own + * checks, document rows, encoding) and sent as one call of the commit function, which locks the session row, + * checks that this storage still owns it at the sequence it expects, and stores each batch against the state the + * earlier ones left, stopping at the first it refuses. Only a `document.copy` reads stored state while being + * prepared, so a batch with one starts a new call once the batches before it are durable; every other group is one + * round trip. A refusal here or in the function keeps the batches before it; any other error rejects the call. + */ + commit( + batches: readonly (readonly StorageWrite[])[], + receive: (index: number, seq: Seq) => void, + _context: Context, + ): Promise { return this.admit(async () => { - const documentActions = this.prepareDocumentActions(writes); - const candidateNextId = this.candidateNextId(writes); - this.checkBatchIds(writes); - const seq = seqFromNumber(this.nextSeq); - const documentWrites = await this.prepareDocumentWrites(documentActions, seq); - // One round trip: the commit function locks the session row, checks that this storage still owns it - // at the sequence it expects, validates the batch against the stored state, and writes every table. - try { - await this.pool.query(`SELECT ${this.commitFunction}($1, $2, $3, $4, $5::json)`, [ - this.session, - this.owner, - seq, - candidateNextId, - this.encodeBatch(writes, documentWrites), - ]); - } catch (error) { - throw this.commitError(error); + const stored: Seq[] = []; + const deliver = (): void => { + for (const [index, seq] of stored.entries()) receive(index, seq); + }; + let index = 0; + while (index < batches.length) { + const chunk: PreparedBatch[] = []; + let refused: { readonly error: unknown } | undefined; + let candidateNextId = this.nextId; + while (index < batches.length) { + const writes = batches[index]!; + if (chunk.length > 0 && writes.some((write) => write.type === "document.copy")) break; + let prepared: PreparedBatch; + try { + prepared = await this.prepareBatch(writes, seqFromNumber(this.nextSeq + chunk.length)); + } catch (error) { + if (error instanceof PostgresUnavailable) { + deliver(); + throw error.cause; + } + refused = { error }; + break; + } + candidateNextId = Math.max(candidateNextId, prepared.candidateNextId); + chunk.push(prepared); + index++; + } + if (chunk.length > 0) { + let result: CommitBatchesResult; + try { + const row = ( + await this.pool.query<{ readonly result: string }>( + `SELECT ${this.commitFunction}($1, $2, $3, $4, $5::json)::text AS result`, + [ + this.session, + this.owner, + this.nextSeq, + candidateNextId, + `[${chunk.map((batch) => batch.encoded).join(",")}]`, + ], + ) + ).rows[0]; + if (row === undefined) throw new Error("Durable Postgres commit returned no result"); + result = JSON.parse(row.result) as CommitBatchesResult; + } catch (error) { + deliver(); + throw this.commitError(error); + } + this.nextId = Math.max(this.nextId, candidateNextId); + for (const batch of chunk.slice(0, result.stored)) { + const seq = seqFromNumber(this.nextSeq); + this.nextSeq = seq + 1; + this.absorb(batch.writes, batch.actions, batch.documents, seq); + stored.push(seq); + } + if (result.code !== undefined) { + deliver(); + return { stored: stored.length, refused: new Error(result.message) }; + } + } + if (refused !== undefined) { + deliver(); + return { stored: stored.length, refused: refused.error }; + } } - this.nextSeq = seq + 1; - this.nextId = Math.max(this.nextId, candidateNextId); - this.absorb(writes, documentActions, documentWrites, seq); - return seq; + deliver(); + return { stored: stored.length }; }); } + /** Runs the adapter's own checks on one batch and encodes it. Throws the refusal, or `PostgresUnavailable`. */ + private async prepareBatch(writes: readonly StorageWrite[], seq: Seq): Promise { + const actions = this.prepareDocumentActions(writes); + this.checkBatchIds(writes); + const documents = await this.prepareDocumentWrites(actions, seq); + return { + writes, + actions, + documents, + encoded: this.encodeBatch(writes, documents), + candidateNextId: this.candidateNextId(writes), + }; + } + /** Maps an error raised by the commit function to the adapter's own errors; other errors pass through. */ private commitError(error: unknown): unknown { const code = (error as { readonly code?: unknown }).code; @@ -1039,20 +1131,23 @@ export class PostgresStorage implements Storage { return record; } - /** One statement reads the record and the revisions from the newest base at or before `at`. */ - private async materializeDocument(id: DocumentId, at: DocumentPoint): Promise { + /** + * Reads and materializes one document at `at` from one statement. With `preparing`, a failed database request is wrapped as + * `PostgresUnavailable`, so a commit tells it apart from a refusal. + */ + private async materializeDocument( + id: DocumentId, + at: DocumentPoint, + preparing = false, + ): Promise { const upper = at === "current" ? Number.MAX_SAFE_INTEGER : at; - const row = ( - await this.pool.query<{ readonly record: string | null; readonly revisions: string | null }>( - `SELECT (SELECT x.record::text FROM ${this.tables.documents} x WHERE x.session_id = $1 AND x.id = $2) AS record, - (SELECT json_agg(json_build_object('seq', r.seq, 'kind', r.kind, 'version', r.version, 'content', r.content) - ORDER BY r.seq)::text - FROM ${this.tables.revisions} r WHERE r.session_id = $1 AND r.document_id = $2 AND r.seq <= $3 - AND r.seq >= (SELECT b.seq FROM ${this.tables.revisions} b WHERE b.session_id = $1 AND b.document_id = $2 - AND b.kind = 'base' AND b.seq <= $3 ORDER BY b.seq DESC LIMIT 1)) AS revisions`, - [this.session, id, upper], - ) - ).rows[0]; + let row: { readonly record: string | null; readonly revisions: string | null } | undefined; + try { + row = (await this.queryDocumentRevisions(id, upper)).rows[0]; + } catch (error) { + if (preparing) throw new PostgresUnavailable(error); + throw error; + } if (row?.record === null || row?.record === undefined) return undefined; const record = parseJson(row.record); if (at !== "current" && isCurrentOnly(record)) { @@ -1083,6 +1178,19 @@ export class PostgresStorage implements Storage { return { record, version: base.version, value, deltasSinceBase: tail.length }; } + /** The document's record and the revisions from its newest base at or before `upper` through `upper`. */ + private queryDocumentRevisions(id: DocumentId, upper: number) { + return this.pool.query<{ readonly record: string | null; readonly revisions: string | null }>( + `SELECT (SELECT x.record::text FROM ${this.tables.documents} x WHERE x.session_id = $1 AND x.id = $2) AS record, + (SELECT json_agg(json_build_object('seq', r.seq, 'kind', r.kind, 'version', r.version, 'content', r.content) + ORDER BY r.seq)::text + FROM ${this.tables.revisions} r WHERE r.session_id = $1 AND r.document_id = $2 AND r.seq <= $3 + AND r.seq >= (SELECT b.seq FROM ${this.tables.revisions} b WHERE b.session_id = $1 AND b.document_id = $2 + AND b.kind = 'base' AND b.seq <= $3 ORDER BY b.seq DESC LIMIT 1)) AS revisions`, + [this.session, id, upper], + ); + } + private candidateNextId(writes: readonly StorageWrite[]): number { let nextId = this.nextId; for (const write of writes) { @@ -1174,7 +1282,7 @@ export class PostgresStorage implements Storage { let content = action.content; if (action.copy !== undefined) { try { - const stored = await this.materializeDocument(action.copy.id, action.copy.at); + const stored = await this.materializeDocument(action.copy.id, action.copy.at, true); if (stored === undefined) throw new Error(`Fork source document ${action.copy.id} cannot be read`); const create = action.create!; if ( @@ -1189,7 +1297,7 @@ export class PostgresStorage implements Storage { } content = { kind: "base", version: stored.version, value: stored.value }; } catch (error) { - if (error instanceof StorageRejected) throw error; + if (error instanceof StorageRejected || error instanceof PostgresUnavailable) throw error; throw new StorageRejected(`Document copy ${id} was rejected`, { cause: error }); } } diff --git a/packages/durable/test/postgres-storage.test.ts b/packages/durable/test/postgres-storage.test.ts index 651344c2433e..349bbc9e472a 100644 --- a/packages/durable/test/postgres-storage.test.ts +++ b/packages/durable/test/postgres-storage.test.ts @@ -1,16 +1,27 @@ import { randomUUID } from "node:crypto"; import type { Context } from "@earendil-works/chord"; import { BACKGROUND_CONTEXT } from "@earendil-works/chord/context"; -import { registerStorageConformance } from "@earendil-works/pi-durable/testing"; +import { commitWrites, registerStorageConformance } from "@earendil-works/pi-durable/testing"; import pg from "pg"; import { afterAll, afterEach, beforeAll, describe, expect, it } from "vitest"; import { idFromNumber } from "../src/ids.ts"; +import type { PostgresPool } from "../src/storage/postgres/database.ts"; import { applyPostgresMigrations, CURRENT_POSTGRES_SCHEMA_VERSION, PostgresStorage, } from "../src/storage/postgres/index.ts"; -import type { ConversationId, EntryId, Id, Seq, Storage, StorageWrite } from "../src/types.ts"; +import type { + CommitResult, + ConversationId, + DocumentCreate, + DocumentId, + EntryId, + Id, + Seq, + Storage, + StorageWrite, +} from "../src/types.ts"; import { ROOT_CONVERSATION_ID } from "../src/types.ts"; // Set to a database the tests may create and drop schemas in, e.g. postgres://postgres:pg@localhost:5432/postgres. @@ -55,10 +66,14 @@ class ReopeningStorage implements Storage { this.session = session; } - async commit(writes: readonly StorageWrite[], commitContext: Context): Promise { - if (this.closed) return this.current.commit(writes, commitContext); + async commit( + batches: readonly (readonly StorageWrite[])[], + receive: (index: number, seq: Seq) => void, + commitContext: Context, + ): Promise { + if (this.closed) return this.current.commit(batches, receive, commitContext); try { - return await this.current.commit(writes, commitContext); + return await this.current.commit(batches, receive, commitContext); } finally { await this.current.close(context); this.current = await PostgresStorage.open(pool, { session: this.session, schema }); @@ -124,7 +139,11 @@ registerStorageConformance({ describe: describePostgres, expect, it }, "Postgres }); async function createRoot(storage: Storage): Promise { - return storage.commit([{ type: "conversation", value: { id: ROOT_CONVERSATION_ID } }], context); + return commitWrites(storage, [{ type: "conversation", value: { id: ROOT_CONVERSATION_ID } }], context); +} + +function message(id: EntryId): StorageWrite { + return { type: "entry", value: { id, conversationId: ROOT_CONVERSATION_ID, kind: "message" } }; } describePostgres("PostgresStorage ownership", () => { @@ -144,13 +163,9 @@ describePostgres("PostgresStorage ownership", () => { await createRoot(first); const second = await openStorage(session); const entryId = await first.mintId(); - await expect( - first.commit([{ type: "entry", value: { id: entryId, conversationId: ROOT_CONVERSATION_ID, kind: "message" } }], context), - ).rejects.toThrow("was opened by another process"); + await expect(commitWrites(first, [message(entryId)], context)).rejects.toThrow("was opened by another process"); await expect(first.mintId()).rejects.toThrow("lost ownership"); - await expect( - second.commit([{ type: "entry", value: { id: entryId, conversationId: ROOT_CONVERSATION_ID, kind: "message" } }], context), - ).resolves.toBe(2); + await expect(commitWrites(second, [message(entryId)], context)).resolves.toBe(2); }); it("rejects commits after another process takes ownership", async () => { @@ -159,12 +174,7 @@ describePostgres("PostgresStorage ownership", () => { await createRoot(storage); await pool.query(`UPDATE "${schema}".sessions SET owner = 'another-process' WHERE session_id = $1`, [session]); const entryId = await storage.mintId(); - await expect( - storage.commit( - [{ type: "entry", value: { id: entryId, conversationId: ROOT_CONVERSATION_ID, kind: "message" } }], - context, - ), - ).rejects.toThrow("was opened by another process"); + await expect(commitWrites(storage, [message(entryId)], context)).rejects.toThrow("was opened by another process"); await expect(storage.conversation(ROOT_CONVERSATION_ID, context)).rejects.toThrow("lost ownership"); }); @@ -173,21 +183,74 @@ describePostgres("PostgresStorage ownership", () => { const first = await openStorage(session); await createRoot(first); const minted = await first.mintId(); - await first.commit( - [{ type: "entry", value: { id: minted, conversationId: ROOT_CONVERSATION_ID, kind: "message" } }], - context, - ); + await commitWrites(first, [message(minted)], context); await first.close(context); const second = await openStorage(session); expect(await second.mintId()).toBe(minted + 1); await expect(createRoot(second)).rejects.toThrow("already belongs to conversation"); const entryId = await second.mintId(); - await expect( - second.commit( - [{ type: "entry", value: { id: entryId, conversationId: ROOT_CONVERSATION_ID, kind: "message" } }], - context, - ), - ).resolves.toBe(3); + await expect(commitWrites(second, [message(entryId)], context)).resolves.toBe(3); + }); +}); + +describePostgres("PostgresStorage round trips", () => { + it("stores a group of batches in one call, and a copy once the batches before it are durable", async () => { + let queries = 0; + const counting: PostgresPool = { + query: (text, values) => { + queries++; + return pool.query(text, values); + }, + connect: () => pool.connect(), + }; + const storage = await PostgresStorage.open(counting, { session: randomUUID(), schema, migrate: false }); + openStorages.add(storage); + await createRoot(storage); + const ids = [await storage.mintId(), await storage.mintId(), await storage.mintId()]; + const receipts: number[] = []; + queries = 0; + const grouped = await storage.commit( + ids.map((id) => [message(id)]), + (index) => receipts.push(index), + context, + ); + expect(grouped).toEqual({ stored: 3 }); + expect(receipts).toEqual([0, 1, 2]); + expect(queries).toBe(1); + + const sourceId = await storage.mintId(); + const copyId = await storage.mintId(); + const childId = await storage.mintId(); + const record = (id: DocumentId, conversationId: ConversationId): DocumentCreate => ({ + id, + kind: "trip", + scope: { kind: "conversation", conversationId }, + history: "rewindable", + fork: "asOf", + }); + queries = 0; + const copied = await storage.commit( + [ + [ + { + type: "document.create", + record: record(sourceId, ROOT_CONVERSATION_ID), + content: { kind: "base", version: 1, value: { n: 1 } }, + }, + ], + [{ type: "document.change", id: sourceId, content: { kind: "delta", version: 1, ops: [["s", ["n"], 2]] } }], + [ + { type: "conversation", value: { id: childId } }, + { type: "document.copy", record: record(copyId, childId), source: { id: sourceId, at: "current" } }, + ], + ], + () => {}, + context, + ); + expect(copied).toEqual({ stored: 3 }); + // The two batches before the copy in one call, the copy's source read once they are durable, then the copy's call. + expect(queries).toBe(3); + expect((await storage.document(copyId, "current", context))?.value).toEqual({ n: 2 }); }); });