diff --git a/bun.lock b/bun.lock index 890ff50d..73de8d0f 100644 --- a/bun.lock +++ b/bun.lock @@ -92,6 +92,7 @@ "name": "typewriter", "dependencies": { "@effect/platform-bun": "4.0.3", + "@electric-sql/pglite": "0.5.8", "abipg": "workspace:*", "abitype": "catalog:", "drizzle-kit": "catalog:", diff --git a/packages/typewriter/package.json b/packages/typewriter/package.json index 05be78ba..4a781a28 100644 --- a/packages/typewriter/package.json +++ b/packages/typewriter/package.json @@ -13,6 +13,7 @@ }, "dependencies": { "@effect/platform-bun": "4.0.3", + "@electric-sql/pglite": "0.5.8", "abipg": "workspace:*", "abitype": "catalog:", "drizzle-kit": "catalog:", diff --git a/packages/typewriter/src/db-query.ts b/packages/typewriter/src/db-query.ts index e3a4f288..6f05250f 100644 --- a/packages/typewriter/src/db-query.ts +++ b/packages/typewriter/src/db-query.ts @@ -153,9 +153,9 @@ export function selectMaxExecutionIndex( sql` UNION ALL `, ); const [{ max = null } = {}] = yield* tx.query((db) => - db.execute<{ max: string | null }>( - sql`SELECT max(max) AS max FROM (${maxes}) AS mutations`, - ), + db + .select({ max: sql`max(max)::text` }) + .from(sql`(${maxes}) AS mutations`), ); return max === null ? undefined : BigInt(max); }); diff --git a/packages/typewriter/src/db-transport.ts b/packages/typewriter/src/db-transport.ts index 7fee6c61..96a813d4 100644 --- a/packages/typewriter/src/db-transport.ts +++ b/packages/typewriter/src/db-transport.ts @@ -1,17 +1,24 @@ +import { PGlite } from "@electric-sql/pglite"; import { SQL } from "bun"; import { DrizzleQueryError, type SQLWrapper } from "drizzle-orm"; -import { type BunSQLDatabase, drizzle } from "drizzle-orm/bun-sql"; -import { PgDialect } from "drizzle-orm/pg-core"; +import { drizzle as drizzleBunSql } from "drizzle-orm/bun-sql"; +import { + type PgDatabase, + PgDialect, + type PgQueryResultHKT, +} from "drizzle-orm/pg-core"; +import { drizzle as drizzlePglite } from "drizzle-orm/pglite"; import { Data, Deferred, Effect, Exit, Schedule, type Scope } from "effect"; export type PostgresOptions = { + kind?: "postgres"; url: string; maxConnections?: number; statementTimeoutMs?: number; idleInTransactionTimeoutMs?: number; }; -export type Drizzle = BunSQLDatabase; +export type Drizzle = PgDatabase; export type DbStatement = PromiseLike & SQLWrapper; @@ -57,16 +64,97 @@ export class DbError extends Data.TaggedError("DbError")<{ const dialect = new PgDialect(); +export function statementSql(statement: SQLWrapper): string { + return dialect.sqlToQuery(statement.getSQL()).sql; +} + +function makeQuery( + db: Drizzle, + notApplied: (cause: unknown) => boolean, +): DbQuery { + return (build) => + Effect.suspend(() => { + const statement = build(db); + return Effect.tryPromise({ + try: () => statement, + catch: (cause) => + new DbError({ + sql: statementSql(statement), + cause, + notApplied: notApplied(cause), + }), + }); + }); +} + +/** + * Runs `f` in the transaction that `begin` opens and commits once `run` + * resolves. `commitFailed` decides the result when the transaction fails + * after `f` succeeded. + */ +function runTransaction( + f: (tx: DbTransaction) => Effect.Effect, + begin: (run: (query: DbQuery) => Promise) => Promise, + commitFailed: (cause: unknown, value: a) => Effect.Effect, +): Effect.Effect { + return Effect.flatMap(Effect.context(), (context) => + Effect.callback((resume) => { + const controller = new AbortController(); + let exit: Exit.Exit | undefined; + // A failed statement aborts the transaction, even if `f` recovers. + let failed: DbError | undefined; + const done = begin(async (query) => { + exit = await Effect.runPromiseExitWith(context)( + Effect.suspend(() => + f({ + query: (build) => + query(build).pipe( + Effect.tapError((error) => + Effect.sync(() => { + failed ??= error; + }), + ), + ), + }), + ), + { signal: controller.signal }, + ); + if (Exit.isFailure(exit)) throw exit; + if (failed !== undefined) throw failed; + return exit.value; + }) + .then( + (value) => resume(Effect.succeed(value)), + (cause) => { + if (exit !== undefined && Exit.isFailure(exit)) { + return resume(exit); + } + if (failed !== undefined) return resume(Effect.fail(failed)); + if (exit === undefined) { + return resume( + Effect.fail( + new DbError({ sql: "BEGIN", cause, notApplied: true }), + ), + ); + } + resume(commitFailed(cause, exit.value)); + }, + ) + .finally(() => controller.abort()); + return Effect.promise(() => { + controller.abort(); + return done; + }); + }), + ); +} + const LOCK_SQL = "SELECT pg_advisory_lock($1::bigint)"; const COMMIT_STATUS_SQL = "select pg_xact_status($1::xid8) as status"; /** The transaction's `COMMIT` hasn't finished, so check its status again. */ class CommitInProgressError extends Data.TaggedError("CommitInProgressError") {} -export function statementSql(statement: SQLWrapper): string { - return dialect.sqlToQuery(statement.getSQL()).sql; -} - /** Codes after which a statement did not apply. */ const NOT_APPLIED = new Set([ "ERR_POSTGRES_CONNECTION_REFUSED", @@ -99,24 +187,6 @@ function hasCode(cause: unknown, codes: Set): boolean { ); } -function makeQuery(db: Drizzle, inTransaction: boolean): DbQuery { - return (build) => - Effect.suspend(() => { - const statement = build(db); - return Effect.tryPromise({ - try: () => statement, - catch: (cause) => - new DbError({ - sql: statementSql(statement), - cause, - notApplied: - hasCode(cause, NOT_APPLIED) || - (inTransaction && hasCode(cause, CONNECTION_LOST)), - }), - }); - }); -} - export function makePostgresTransport( options: PostgresOptions, ): Effect.Effect { @@ -142,7 +212,9 @@ export function makePostgresTransport( (client) => Effect.promise(() => client.close()), ); - const query = makeQuery(drizzle({ client }), false); + const query = makeQuery(drizzleBunSql({ client }), (cause) => + hasCode(cause, NOT_APPLIED), + ); // The status of a transaction whose `COMMIT` got no answer, or `undefined`. // It uses its own connection, since the pool may be busy or failing. @@ -171,105 +243,68 @@ export function makePostgresTransport( // Not `reserve`: after a lost connection, a reserved connection can run // later statements on another session. - const transaction: DbTransport["transaction"] = ( - f: (tx: DbTransaction) => Effect.Effect, - ) => - Effect.flatMap(Effect.context(), (context) => - Effect.callback((resume) => { - const controller = new AbortController(); - let exit: Exit.Exit | undefined; - // A failed statement aborts the transaction, even if `f` recovers. - let failed: DbError | undefined; - // Set before `COMMIT`. A transaction without an id wrote nothing. - let xid: string | null | undefined; - const done = client - .begin(async (tx) => { - const query = makeQuery(drizzle({ client: tx }), true); - exit = await Effect.runPromiseExitWith(context)( - Effect.suspend(() => - f({ - query: (build) => - query(build).pipe( - Effect.tapError((error) => - Effect.sync(() => { - failed ??= error; - }), - ), - ), - }), + const transaction: DbTransport["transaction"] = (f) => + Effect.suspend(() => { + // Set before `COMMIT`. A transaction without an id wrote nothing. + let xid: string | null | undefined; + return runTransaction( + f, + (run) => + client.begin(async (tx) => { + const value = await run( + makeQuery( + drizzleBunSql({ client: tx }), + (cause) => + hasCode(cause, NOT_APPLIED) || + hasCode(cause, CONNECTION_LOST), ), - { signal: controller.signal }, ); - if (Exit.isFailure(exit)) throw exit; - if (failed !== undefined) throw failed; [{ xid }] = await tx`select pg_current_xact_id_if_assigned()::text as xid`; - return exit.value; - }) - .then( - (value) => resume(Effect.succeed(value as a)), - (cause) => { - if (exit !== undefined && Exit.isFailure(exit)) { - return resume(exit); - } - if (failed !== undefined) return resume(Effect.fail(failed)); - // Postgres rolls back unless it ran `COMMIT`. - if (exit === undefined || xid === undefined) { - return resume( - Effect.fail( - new DbError({ - sql: - exit === undefined - ? "BEGIN" - : "select pg_current_xact_id_if_assigned()", - cause, - notApplied: true, - }), - ), - ); - } - // Postgres answered `COMMIT` with an error, so it rolled back. - if ( - hasCode(cause, SERVER_ERROR) && - !hasCode(cause, CONNECTION_LOST) - ) { - return resume( - Effect.fail( + return value; + }), + (cause, value) => { + // Postgres rolls back unless it ran `COMMIT`. + if (xid === undefined) { + return Effect.fail( + new DbError({ + sql: "select pg_current_xact_id_if_assigned()", + cause, + notApplied: true, + }), + ); + } + // Postgres answered `COMMIT` with an error, so it rolled back. + if ( + hasCode(cause, SERVER_ERROR) && + !hasCode(cause, CONNECTION_LOST) + ) { + return Effect.fail( + new DbError({ + sql: "COMMIT", + cause, + notApplied: hasCode(cause, NOT_APPLIED), + }), + ); + } + if (xid === null) return Effect.succeed(value); + // `COMMIT` may have run. + return commitStatus(xid).pipe( + Effect.flatMap((status) => + status === "committed" + ? Effect.succeed(value) + : Effect.fail( new DbError({ sql: "COMMIT", cause, - notApplied: hasCode(cause, NOT_APPLIED), + notApplied: status === "aborted", }), ), - ); - } - const value = exit.value; - if (xid === null) return resume(Effect.succeed(value)); - // `COMMIT` may have run. - resume( - commitStatus(xid).pipe( - Effect.flatMap((status) => - status === "committed" - ? Effect.succeed(value) - : Effect.fail( - new DbError({ - sql: "COMMIT", - cause, - notApplied: status === "aborted", - }), - ), - ), - ), - ); - }, - ) - .finally(() => controller.abort()); - return Effect.promise(() => { - controller.abort(); - return done; - }); - }), - ); + ), + ); + }, + ); + }); // The lock belongs to one session, so it has its own connection, which // Bun must not close. Closing it ends the session and releases the lock. @@ -302,10 +337,41 @@ export function makePostgresTransport( return Deferred.await(lost); }); + return { query, transaction, lock }; + }); +} + +/** + * An in-memory PGlite database, until the scope closes. + * + * @dev A text value that starts with U+FEFF (a BOM) reads back without it: + * PGlite's protocol reader decodes each field with a `TextDecoder` that strips + * a leading BOM. The stored value keeps it. + */ +export function makePgliteTransport(): Effect.Effect< + DbTransport, + never, + Scope.Scope +> { + return Effect.gen(function* () { + const client = yield* Effect.acquireRelease( + Effect.promise(() => PGlite.create()), + (client) => Effect.promise(() => client.close()), + ); + const db = drizzlePglite({ client }); + // PGlite has no connection to lose, so a retry would fail the same way. + const notApplied = () => false; + return { - query, - transaction, - lock, + query: makeQuery(db, notApplied), + transaction: (f) => + runTransaction( + f, + (run) => db.transaction((tx) => run(makeQuery(tx, notApplied))), + (cause) => Effect.fail(new DbError({ sql: "COMMIT", cause })), + ), + // Only this process can reach an in-memory database. + lock: () => Effect.succeed(Effect.never), }; }); } diff --git a/packages/typewriter/src/db.ts b/packages/typewriter/src/db.ts index fc282fec..bb9a0964 100644 --- a/packages/typewriter/src/db.ts +++ b/packages/typewriter/src/db.ts @@ -2,11 +2,16 @@ import { Context, Effect, Layer, Schedule, type Scope } from "effect"; import { DbError, type DbTransport, + makePgliteTransport, makePostgresTransport, type PostgresOptions, } from "./db-transport"; -export type DatabaseOptions = PostgresOptions; +/** + * A Postgres database, or an in-memory PGlite database that lasts as long as + * the process. + */ +export type DatabaseOptions = PostgresOptions | { kind: "pglite" }; /** Retries run a transaction again from the start, so `f` can run more than once. */ export type Db = DbTransport; @@ -19,6 +24,23 @@ export const layerDatabase = ( options: DatabaseOptions, ): Layer.Layer => Layer.effect(Database, makeDb(options)); +// Inside a transaction, a statement on `Db` instead of `tx` runs outside it on +// Postgres, and waits on it forever on PGlite. +const InTransaction = Context.Reference("typewriter/InTransaction", { + defaultValue: () => false, +}); + +function outsideTransaction(effect: Effect.Effect) { + return Effect.gen(function* () { + if (yield* InTransaction) { + return yield* Effect.die( + new Error("Used the database inside a transaction instead of `tx`"), + ); + } + return yield* effect; + }); +} + function retry(effect: Effect.Effect) { return effect.pipe( // A failure that stops the runtime is logged when it stops, so an attempt @@ -35,8 +57,15 @@ function retry(effect: Effect.Effect) { /** Retries the work of `transport` that did not apply. */ function makeDbClient(transport: DbTransport): Db { return { - query: (build) => retry(transport.query(build)), - transaction: (f) => retry(transport.transaction(f)), + query: (build) => outsideTransaction(retry(transport.query(build))), + transaction: (f) => + outsideTransaction( + retry( + transport.transaction((tx) => + Effect.provideService(f(tx), InTransaction, true), + ), + ), + ), lock: transport.lock, }; } @@ -44,5 +73,10 @@ function makeDbClient(transport: DbTransport): Db { export function makeDb( options: DatabaseOptions, ): Effect.Effect { - return Effect.map(makePostgresTransport(options), makeDbClient); + return Effect.map( + options.kind === "pglite" + ? makePgliteTransport() + : makePostgresTransport(options), + makeDbClient, + ); }