diff --git a/package-lock.json b/package-lock.json index e65704b94447..56bd4ace2dfd 100644 --- a/package-lock.json +++ b/package-lock.json @@ -1731,6 +1731,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", @@ -4124,6 +4136,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.0", + "resolved": "https://registry.npmjs.org/pg-cloudflare/-/pg-cloudflare-1.4.0.tgz", + "integrity": "sha512-Vo7z/6rrQYxpNRylp4Tlob2elzbh+N/MOQbxFVWCxS7oEx6jF53GTJFxK2WWpKuBRkmiin4Mt+xofFDjx09R0A==", + "dev": true, + "license": "MIT", + "optional": true + }, + "node_modules/pg-connection-string": { + "version": "2.14.0", + "resolved": "https://registry.npmjs.org/pg-connection-string/-/pg-connection-string-2.14.0.tgz", + "integrity": "sha512-XwWDGcLRGCXAR8F/AM5bG7Q+A3Wm2s6QeEjlOKZLlH3UYcguiqCWKyWXVag5TLTIjR7oOJUY8kcADaZgWPyLeg==", + "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.0", + "resolved": "https://registry.npmjs.org/pg-protocol/-/pg-protocol-1.16.0.tgz", + "integrity": "sha512-sILXutLVjCLjcDuOmvhX5e2Z4cS5qG/6Bu3VkpFwdf/633ElGLpEh9bgmuI5I4sqKqkifQiGyiCcx1HdtrK7tg==", + "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 @@ -4193,6 +4302,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", @@ -4649,6 +4801,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", @@ -5249,6 +5411,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", @@ -5578,6 +5750,8 @@ "typebox": "1.3.27" }, "devDependencies": { + "@types/pg": "8.23.1", + "pg": "8.23.0", "shx": "0.4.0", "vitest": "4.1.11" }, diff --git a/packages/durable/CHANGELOG.md b/packages/durable/CHANGELOG.md index 5d587d0f47cc..5e275ec11723 100644 --- a/packages/durable/CHANGELOG.md +++ b/packages/durable/CHANGELOG.md @@ -2,6 +2,13 @@ ## [Unreleased] +### Changed + +- Postgres storage is built for a database that is a network hop away. A commit is two round trips: one statement reads the session row, the batch's claimed ids, its documents, and their latest revision versions as JSON; one data-modifying CTE statement writes every table's rows, clears and adds revisions, and advances the session, each part gated on the session still carrying this owner at the sequence the read observed. A single statement is atomic, so no transaction brackets a commit, and a takeover by another process writes nothing and is reported as before. +- Postgres storage serves repeated reads from memory. It is the Session's only writer, so a record it has read stays valid until it writes that record: conversation records, tasks and submissions by id and request, each conversation's newest entries and head markers, current document scans and addresses, and the materialized current content of documents are cached and kept current by every commit. Measured on one faux-model turn, admitting an input drops from 22 queries to 3, creating a conversation from 37 to 11, a plain turn from 61 to 9, and a turn with one tool call from 189 to 31. +- `PostgresStorage.open` takes `migrate: false` to skip the schema check for deployments that migrate in an explicit step; an open is then two statements. + + ### Breaking Changes - The portable SQLite facade in `@earendil-works/pi-durable/storage/sqlite` is asynchronous: `SqliteDatabase` extends the new `SqliteExecutor` (`exec`, `run`, `get`, `all` by SQL text), `prepare` and `SqliteStatement` are removed, `transaction` takes an async callback that receives a transaction handle, and `close()` returns a promise. Custom adapters must be rewritten ([#10232](https://github.com/earendil-works/pi/pull/10232) by [@christianklotz](https://github.com/christianklotz)). @@ -16,6 +23,7 @@ ### Added +- `PostgresStorage` in `@earendil-works/pi-durable/storage/postgres`: Sessions in a Postgres schema over a caller-owned pool such as `pg.Pool`, with a per-Session advisory lock and an owner token checked by every commit. The package adds no runtime dependency. - `CodingTools` extension in `@earendil-works/pi-durable/tools` with `read`, `write`, `edit`, and `bash`. - `AgentDoc`, `configure()`, `DEFAULT_RETRY_POLICY`, and `DEFAULT_COMPACTION_POLICY` exports. - `HarnessOptions.conversationCreated(tx, conversation)` runs in every commit that creates or forks a conversation, after the built-in documents, so applications can create their own documents in every conversation. diff --git a/packages/durable/README.md b/packages/durable/README.md index bf1d169d9a50..2d1174f827f1 100644 --- a/packages/durable/README.md +++ b/packages/durable/README.md @@ -523,8 +523,9 @@ const usage = await harness.usage(context); // { models: { "openai/gpt-6-sol": U | Memory | `MemoryStorage` from the package root | Nothing is persisted. | | SQLite | `openNodeSqliteStorage(file)` from `@earendil-works/pi-durable/storage/sqlite/node` | One database file. WAL mode with `synchronous = NORMAL`: commits survive process crashes; the newest may be lost on power or host failure. | | JSONL | `openNodeJsonlStorage(directory, context)` from `@earendil-works/pi-durable/storage/jsonl/node` | Append-only files in one directory. Pass `{ fsync: true }` to flush before each commit marker. | +| Postgres | `PostgresStorage.open(pool, { session })` from `@earendil-works/pi-durable/storage/postgres` | One Session per `session` name; many Sessions share one schema. Locked to one process. Built for a remote database: a commit is two round trips and repeated reads are served from memory. Pass `migrate: false` when the schema is migrated in an explicit step, so an open costs two statements. | -One process owns a storage at a time; there is no cross-process locking. The portable SQLite and JSONL cores (`/storage/sqlite`, `/storage/jsonl`) run without Node APIs, for example on Bun or in Cloudflare Durable Objects, given an asynchronous `SqliteDatabase` facade or a `FileSystem` from `@earendil-works/pi-durable/env`. +One process owns a storage at a time; SQLite and JSONL do no cross-process locking. The portable SQLite and JSONL cores (`/storage/sqlite`, `/storage/jsonl`) run without Node APIs, for example on Bun or in Cloudflare Durable Objects, given an asynchronous `SqliteDatabase` facade or a `FileSystem` from `@earendil-works/pi-durable/env`. SQLite adapters implement promise-based `exec`, `run`, `get`, `all`, `transaction`, and `close`. `run`, `get`, and `all` take SQL text plus positional bindings; adapters may cache prepared statements by SQL text. A transaction callback receives a transaction handle; all work in the transaction must use it, and the handle expires when the callback settles. Adapters must queue unrelated operations and other transactions until the transaction finishes, so calling `database` itself inside the callback never settles: @@ -535,6 +536,16 @@ await database.transaction(async (transaction) => { }); ``` +Postgres storage runs on a pool you create and own; a `pg.Pool` fits as is, and this package imports no driver. The first open creates the `pi_durable` schema (or `schema`) and its tables. Each storage keeps one pooled client checked out to hold its Session's advisory lock, so give the pool at least two connections. A second open of a held Session rejects. Every commit also checks the owner token written by the latest open, so a process that lost its lock can no longer write. Requires Postgres 11 or newer. + +```typescript +import pg from "pg"; +import { PostgresStorage } from "@earendil-works/pi-durable/storage/postgres"; + +const pool = new pg.Pool({ connectionString: process.env.DATABASE_URL }); +const harness = await Harness.open(await PostgresStorage.open(pool, { session: "support-bot" }), { models, registry }, context); +``` + Custom backends can run the shared conformance suite with any Vitest- or Jest-compatible runner: ```typescript @@ -590,6 +601,8 @@ Examples that call OpenAI need `OPENAI_API_KEY`; most use the faux provider othe - [`docs/pico-v5-handoff.md`](docs/pico-v5-handoff.md): the implementation plan - [`docs/pico-v5-chord-usage.md`](docs/pico-v5-chord-usage.md): how the package uses Chord +Postgres storage tests run when `PI_DURABLE_POSTGRES_URL` names a database they may create and drop schemas in; otherwise they are skipped. + Benchmarks: `npm run bench:storage`, `npm run bench:storage:memory`, and `npm run bench:tool-output`. ## License diff --git a/packages/durable/package.json b/packages/durable/package.json index 63fb5c0b13e2..d4eb6767f83a 100644 --- a/packages/durable/package.json +++ b/packages/durable/package.json @@ -51,6 +51,11 @@ "types": "./dist/storage/sqlite/node.d.ts", "import": "./dist/storage/sqlite/node.js" }, + "./storage/postgres": { + "source": "./src/storage/postgres/index.ts", + "types": "./dist/storage/postgres/index.d.ts", + "import": "./dist/storage/postgres/index.js" + }, "./testing": { "source": "./src/testing/index.ts", "types": "./dist/testing/index.d.ts", @@ -96,6 +101,8 @@ "typebox": "1.3.27" }, "devDependencies": { + "@types/pg": "8.23.1", + "pg": "8.23.0", "shx": "0.4.0", "vitest": "4.1.11" } 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..014e3cb008c3 --- /dev/null +++ b/packages/durable/src/storage/postgres/storage.ts @@ -0,0 +1,1222 @@ +import type { Context, JsonValue } from "@earendil-works/chord"; +import { apply, type Op } from "@earendil-works/chord/delta"; +import { StorageRejected } 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, + Seq, + Storage, + StorageWrite, + StoredDocument, + SubmissionId, + SubmissionQuery, + SubmissionRecord, + TaskId, + TaskQuery, + TaskRecord, +} from "../../types.ts"; +import type { PostgresPool, PostgresPoolClient, 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 TypeError("Invalid storage cursor"); + return idFromNumber(after); +}; + +const page = }>(values: readonly T[], limit: number): Page => { + const items = values.slice(0, limit); + if (values.length <= limit) return { items }; + return { items, next: { after: items.at(-1)!.id } }; +}; + +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; + } +}; + +const lockKey = (schema: string, session: string): string => `pi-durable:session:${schema}:${session}`; + +/** + * Postgres implementation of the Pico storage contract. + * + * One process owns a Session at a time. `open()` takes a session-level advisory lock on a client it keeps checked out, + * and writes a fresh owner token that every commit checks under a row lock. If the lock connection drops, the storage + * stops accepting work; if another process has since opened the Session, a commit from this one is rejected. + */ +export class PostgresStorage implements Storage { + private readonly pool: PostgresPool; + private readonly lockClient: PostgresPoolClient; + private readonly onLockError: (error: Error) => void; + private readonly session: string; + private readonly lockKey: 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, + lockClient: PostgresPoolClient, + onLockError: (error: Error) => void, + options: { session: string; schema: string; lockKey: string; owner: string; nextId: number; nextSeq: number }, + ) { + this.pool = pool; + this.lockClient = lockClient; + this.onLockError = onLockError; + this.session = options.session; + this.lockKey = options.lockKey; + 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, lock the Session, and take ownership of it. Rejects when another process holds the Session. + */ + 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 lockClient = await pool.connect(); + let storage: PostgresStorage | undefined; + let lockError: Error | undefined; + const onLockError = (error: Error) => { + lockError = error; + storage?.loseLock(error); + }; + lockClient.on("error", onLockError); + const key = lockKey(schema, options.session); + let locked = false; + try { + const lock = await lockClient.query<{ readonly locked: boolean }>( + "SELECT pg_try_advisory_lock(hashtextextended($1, 0)) AS locked", + [key], + ); + if (lock.rows[0]?.locked !== true) { + throw new Error(`Durable Postgres session ${options.session} is open in another process`); + } + locked = true; + const owner = globalThis.crypto.randomUUID(); + const metadata = ( + await lockClient.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"); + storage = new PostgresStorage(pool, lockClient, onLockError, { + session: options.session, + schema, + lockKey: key, + owner, + nextId: Number(metadata.next_id), + nextSeq: Number(metadata.next_seq), + }); + if (lockError !== undefined) storage.loseLock(lockError); + return storage; + } catch (error) { + let discard = lockError; + if (locked && discard === undefined) { + try { + await lockClient.query("SELECT pg_advisory_unlock(hashtextextended($1, 0))", [key]); + } catch (unlockError) { + // Preserve the open failure; discarding the client closes its connection, which releases the lock. + discard = unlockError as Error; + } + } + lockClient.off("error", onLockError); + lockClient.release(discard); + throw error; + } + } + + 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.loseLock(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 params: PostgresValue[] = [this.session]; + const clauses = ["session_id = $1", `id > $${params.push(cursorId(cursor) ?? -1)}`]; + 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 LIMIT $${params.push(limit + 1)}`, + params, + ); + return page(rows.map(parseJson), limit); + }); + } + + 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> { + let conversation = await this.readConversation(query.conversationId); + if (conversation === undefined) throw new Error(`Unknown conversation: ${query.conversationId}`); + const after = cursorId(cursor); + 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); + } + + /** + * 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 params: PostgresValue[] = [this.session]; + const clauses = ["session_id = $1", `id > $${params.push(cursorId(cursor) ?? -1)}`]; + 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 LIMIT $${params.push(limit + 1)}`, + params, + ); + return page(rows.map(parseJson), limit); + }); + } + + 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 params: PostgresValue[] = [this.session]; + const clauses = ["session_id = $1", `id > $${params.push(cursorId(cursor) ?? -1)}`]; + 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 LIMIT $${params.push(limit + 1)}`, + params, + ); + return page(rows.map(parseJson), limit); + }); + } + + 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, release the session lock, and return its client. 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; + }); + } + let failure = this.lost; + if (failure === undefined) { + try { + await this.lockClient.query("SELECT pg_advisory_unlock(hashtextextended($1, 0))", [this.lockKey]); + } catch (error) { + failure = error as Error; + } + } + this.lockClient.off("error", this.onLockError); + // Discarding a client whose unlock failed closes its connection, which releases the lock. + this.lockClient.release(failure); + } + + /** + * Run one storage operation. Close waits for admitted operations, so the session lock outlives 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 loseLock(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..f76ef96c19fa --- /dev/null +++ b/packages/durable/test/postgres-storage.test.ts @@ -0,0 +1,247 @@ +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); +} + +async function lockHolder(session: string): Promise { + const result = await pool.query<{ pid: number }>( + `SELECT pid FROM pg_locks WHERE locktype = 'advisory' AND objsubid = 1 AND granted + AND ((classid::bigint << 32) | objid::bigint) = hashtextextended($1, 0)`, + [`pi-durable:session:${schema}:${session}`], + ); + return result.rows[0]?.pid; +} + +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("rejects a second open of a Session until the first closes", async () => { + const session = randomUUID(); + const first = await openStorage(session); + await expect(PostgresStorage.open(pool, { session, schema })).rejects.toThrow("is open in another process"); + await first.close(context); + const second = await openStorage(session); + await expect(createRoot(second)).resolves.toBe(1); + }); + + it("stops accepting work when the lock connection drops", async () => { + const session = randomUUID(); + const storage = await openStorage(session); + const pid = await lockHolder(session); + expect(pid).toBeDefined(); + await pool.query("SELECT pg_terminate_backend($1)", [pid]); + await expect + .poll(() => + storage.mintId().then( + () => "open", + (error: Error) => error.message, + ), + ) + .toContain("lost ownership"); + await expect(createRoot(storage)).rejects.toThrow("lost ownership"); + const replacement = await openStorage(session); + await expect(createRoot(replacement)).resolves.toBe(1); + }); + + 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`); + } + }); +}); diff --git a/packages/durable/test/storage-runtime-boundary.test.ts b/packages/durable/test/storage-runtime-boundary.test.ts index 84ef2d51252f..8a148e9a5801 100644 --- a/packages/durable/test/storage-runtime-boundary.test.ts +++ b/packages/durable/test/storage-runtime-boundary.test.ts @@ -28,6 +28,14 @@ describe("durable storage runtime boundaries", () => { expect([...graph].some((path) => path.includes("/env/"))).toBe(false); expect([...graph].some((path) => path.includes("/storage/jsonl/"))).toBe(false); expect([...graph].some((path) => path.includes("/storage/sqlite/"))).toBe(false); + expect([...graph].some((path) => path.includes("/storage/postgres/"))).toBe(false); + }); + + it("keeps the Postgres subpath free of Node imports and drivers", async () => { + const graph = await sourceGraph("storage/postgres/index.ts"); + for (const path of graph) { + expect(await readFile(path, "utf8"), `${path} imports a driver`).not.toMatch(/from\s+["']pg["']/); + } }); it("keeps the portable SQLite subpath free of Node imports", async () => {