feat(typewriter): add an in-memory PGlite database transport

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:",

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<string | null>`max(max)::text` })
+        .from(sql`(${maxes}) AS mutations`),
     );
     return max === null ? undefined : BigInt(max);
   });

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<PgQueryResultHKT>;
 
 export type DbStatement<result> = PromiseLike<result> & 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<a, e, r>(
+  f: (tx: DbTransaction) => Effect.Effect<a, e, r>,
+  begin: (run: (query: DbQuery) => Promise<a>) => Promise<a>,
+  commitFailed: (cause: unknown, value: a) => Effect.Effect<a, DbError>,
+): Effect.Effect<a, e | DbError, r> {
+  return Effect.flatMap(Effect.context<r>(), (context) =>
+    Effect.callback<a, e | DbError>((resume) => {
+      const controller = new AbortController();
+      let exit: Exit.Exit<a, e> | 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<string>): 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<DbTransport, never, Scope.Scope> {
@@ -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"] = <a, e, r>(
-      f: (tx: DbTransaction) => Effect.Effect<a, e, r>,
-    ) =>
-      Effect.flatMap(Effect.context<r>(), (context) =>
-        Effect.callback<a, e | DbError>((resume) => {
-          const controller = new AbortController();
-          let exit: Exit.Exit<a, e> | 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),
     };
   });
 }