fix(postgres): keep postgres-js transactions committable after a recovered savepoint refusal

.changeset/postgres-js-savepoint-recovery.md

@@ -0,0 +1,5 @@
+---
+"@nicia-ai/typegraph": patch
+---
+
+Keep a transaction committable on the postgres-js driver after TypeGraph recovers a refusal through a savepoint. postgres-js remembers the first failed statement of a transaction and rethrows it at commit unless the statement ran inside one of the driver's own savepoint scopes, so a `UniquenessError` caught inside `store.transaction(...)` after `bulkReplaceById` (or another fused write program) lost a unique key still failed the transaction at commit with the recovered driver error, and an `importGraph` whose per-row edge create was refused by the database (a deferred row losing its durable match identity) rolled back the whole import instead of reporting that row as an error. TypeGraph now runs the recovered work inside a postgres-js savepoint scope nested in its own savepoint, so the caller's transaction commits what it wrote outside the refused write on every PostgreSQL driver. A savepoint recovery that itself fails is reported as one `AggregateError` naming both failures.

packages/typegraph/src/backend/drizzle/execution/postgres-transaction-session.ts

@@ -0,0 +1,185 @@
+import { CompilerInvariantError } from "../../../errors";
+import { type DriverSavepointScope } from "../../session-savepoint";
+import {
+  type AnyPgDatabase,
+  createPostgresExecutionAdapter,
+  getPinnedPostgresTransactionClient,
+  isPostgresJsClient,
+  type PostgresExecutionAdapterOptions,
+} from "./postgres-execution";
+import type {
+  CompiledSqlQuery,
+  ExecutableSql,
+  PreparedSqlStatement,
+  SqlExecutionAdapter,
+} from "./types";
+
+type ScopeExecution = Readonly<{
+  db: AnyPgDatabase;
+  operation: SqlExecutionAdapter;
+  compiled: SqlExecutionAdapter;
+}>;
+
+/**
+ * The execution surfaces of one open PostgreSQL transaction session.
+ *
+ * `operation` keeps ordinary transaction work on Drizzle's session; `compiled`
+ * runs already-compiled statements on the pinned driver client.
+ * `enterDriverScope` is present when the driver attributes a failed statement
+ * to the scope whose handle issued it (postgres-js rethrows the first one when
+ * its `begin` or `savepoint` scope ends). While such a scope is open, both
+ * surfaces issue every statement through its handle: PostgreSQL places them
+ * inside the savepoint whichever handle sent them, and routing them through it
+ * keeps the driver's attribution in step with the server.
+ *
+ * Driver scopes on one session must nest strictly: the scope that exits is
+ * always the innermost one still open. Two scopes that overlap without nesting
+ * (both opened by concurrent work in one transaction, the outer one exiting
+ * first) have no single innermost scope to route to, so the out-of-order exit
+ * throws a `CompilerInvariantError`, and every later statement on the session
+ * refuses with the same error instead of running on a scope that has ended.
+ */
+type PostgresTransactionSession = Readonly<{
+  operation: SqlExecutionAdapter;
+  compiled: SqlExecutionAdapter;
+  enterDriverScope: DriverSavepointScope | undefined;
+}>;
+
+function bindScopeExecution(
+  db: AnyPgDatabase,
+  options: PostgresExecutionAdapterOptions | undefined,
+): ScopeExecution {
+  return {
+    db,
+    operation: createPostgresExecutionAdapter(db, options),
+    compiled: createPostgresExecutionAdapter(db, {
+      ...options,
+      useTransactionClient: true,
+    }),
+  };
+}
+
+function requireScopeMember<T>(member: T | undefined, name: string): T {
+  if (member !== undefined) return member;
+  throw new CompilerInvariantError(
+    `A PostgreSQL driver savepoint scope lost the ${name} execution member its transaction exposes.`,
+    { member: name },
+  );
+}
+
+/**
+ * An adapter shaped like `outer` whose statements run on whichever scope
+ * `current` names when they execute. A transaction handle exposes no atomic
+ * batch, exclusive window or serialization classifier of its own, so these
+ * four I/O members are its complete surface.
+ */
+function overlappingScopes(cause?: unknown): CompilerInvariantError {
+  return new CompilerInvariantError(
+    "Overlapping savepoint scopes in one PostgreSQL transaction: a driver savepoint scope ended while a scope opened after it was still open. Savepoint-protected writes on one transaction must not run concurrently.",
+    { code: "OVERLAPPING_SAVEPOINT_SCOPES" },
+    cause === undefined ? undefined : { cause },
+  );
+}
+
+function routeThroughCurrentScope(
+  outer: SqlExecutionAdapter,
+  current: () => SqlExecutionAdapter,
+): SqlExecutionAdapter {
+  return {
+    compile(query: ExecutableSql): CompiledSqlQuery {
+      return current().compile(query);
+    },
+    execute<TRow>(query: ExecutableSql): Promise<readonly TRow[]> {
+      return current().execute<TRow>(query);
+    },
+    ...(outer.executeCompiled === undefined ?
+      {}
+    : {
+        executeCompiled<TRow>(
+          compiledQuery: CompiledSqlQuery,
+        ): Promise<readonly TRow[]> {
+          return requireScopeMember(
+            current().executeCompiled,
+            "executeCompiled",
+          )<TRow>(compiledQuery);
+        },
+      }),
+    ...(outer.prepare === undefined ?
+      {}
+    : {
+        prepare(sqlText: string): PreparedSqlStatement {
+          return {
+            execute<TRow>(
+              params: readonly unknown[],
+            ): Promise<readonly TRow[]> {
+              return requireScopeMember(current().prepare, "prepare")(
+                sqlText,
+              ).execute<TRow>(params);
+            },
+          };
+        },
+      }),
+  };
+}
+
+export function createPostgresTransactionSession(
+  db: AnyPgDatabase,
+  options: PostgresExecutionAdapterOptions | undefined,
+): PostgresTransactionSession {
+  const outer = bindScopeExecution(db, options);
+  if (!isPostgresJsClient(getPinnedPostgresTransactionClient(db))) {
+    return {
+      operation: outer.operation,
+      compiled: outer.compiled,
+      enterDriverScope: undefined,
+    };
+  }
+
+  const openScopes: ScopeExecution[] = [];
+  const endedScopes = new WeakSet<ScopeExecution>();
+
+  function current(): ScopeExecution {
+    const innermost = openScopes.at(-1);
+    if (innermost === undefined) return outer;
+    if (endedScopes.has(innermost)) throw overlappingScopes();
+    return innermost;
+  }
+
+  function leave(scope: ScopeExecution, cause?: unknown): void {
+    endedScopes.add(scope);
+    if (openScopes.at(-1) !== scope) {
+      throw overlappingScopes(cause);
+    }
+    openScopes.pop();
+  }
+
+  // Drizzle's nested transaction on postgres-js is the driver's own
+  // `savepoint()` scope, bound to a Drizzle session over its handle.
+  async function enterDriverScope<T>(scope: () => Promise<T>): Promise<T> {
+    return current().db.transaction(async (nested) => {
+      const entered = bindScopeExecution(nested, options);
+      openScopes.push(entered);
+      let result: T;
+      try {
+        result = await scope();
+      } catch (error) {
+        leave(entered, error);
+        throw error;
+      }
+      leave(entered);
+      return result;
+    });
+  }
+
+  return {
+    operation: routeThroughCurrentScope(
+      outer.operation,
+      () => current().operation,
+    ),
+    compiled: routeThroughCurrentScope(
+      outer.compiled,
+      () => current().compiled,
+    ),
+    enterDriverScope,
+  };
+}

packages/typegraph/src/backend/drizzle/execution/session-atomic-batch.ts

@@ -3,78 +3,73 @@ import type {
   AtomicSqlBatchExecutor,
   CompiledAtomicSqlStatement,
 } from "../../capabilities/atomic-sql-program";
+import {
+  type DriverSavepointScope,
+  runSessionSavepoint,
+  SavepointRecoveryFailure,
+} from "../../session-savepoint";
 import type { SqlExecutionAdapter } from "./types";
 
 type SessionAtomicBatchExecution = Readonly<{
   executeCompiled: NonNullable<SqlExecutionAdapter["executeCompiled"]>;
   runExclusive: NonNullable<SqlExecutionAdapter["runExclusive"]>;
+  enterDriverScope?: DriverSavepointScope | undefined;
 }>;
 
 type SessionAtomicBatchAdapter = Readonly<{
   executeAtomicBatch: AtomicSqlBatchExecutor;
 }>;
 
 const ATOMIC_PROGRAM_SAVEPOINT = "typegraph_atomic_program";
-const BEGIN_ATOMIC_PROGRAM = {
-  sql: `SAVEPOINT ${ATOMIC_PROGRAM_SAVEPOINT}`,
-  params: [],
-} as const;
-const ROLLBACK_ATOMIC_PROGRAM = {
-  sql: `ROLLBACK TO SAVEPOINT ${ATOMIC_PROGRAM_SAVEPOINT}`,
-  params: [],
-} as const;
-const RELEASE_ATOMIC_PROGRAM = {
-  sql: `RELEASE SAVEPOINT ${ATOMIC_PROGRAM_SAVEPOINT}`,
-  params: [],
-} as const;
 
 /**
  * Adapts one already-open, exclusively reservable transaction session to the
  * closed atomic-program transport.
  *
  * `runExclusive` is load-bearing: queueing each statement independently would
  * let unrelated work on the same transaction interleave with a program. The
- * returned batch owns one queue slot for the whole sequence, while the outer
- * transaction owns commit and rollback. A savepoint keeps deliberate
+ * returned batch owns one queue slot for the whole savepoint scope, while the
+ * outer transaction owns commit and rollback. A savepoint keeps deliberate
  * constraint refusals from aborting that outer transaction before Store
  * diagnostics can translate them to typed errors.
  */
 export function createSessionAtomicBatchAdapter(
   execution: SessionAtomicBatchExecution,
 ): SessionAtomicBatchAdapter {
-  const { executeCompiled, runExclusive } = execution;
+  const { executeCompiled, runExclusive, enterDriverScope } = execution;
+  const port = {
+    executeControl: (statement: string) =>
+      executeCompiled({ sql: statement, params: [] }),
+    enterDriverScope,
+  };
 
   return {
     async executeAtomicBatch<TRow>(
       statements: readonly CompiledAtomicSqlStatement[],
     ) {
       if (statements.length === 0) return [];
       return runExclusive(async () => {
-        await executeCompiled(BEGIN_ATOMIC_PROGRAM);
-        const results: (readonly TRow[])[] = [];
         try {
-          for (const statement of statements) {
-            results.push(await executeCompiled<TRow>(statement));
-          }
-          await executeCompiled(RELEASE_ATOMIC_PROGRAM);
-          return results;
+          // PostgreSQL protocol- and SQL-level prepared statements survive
+          // ROLLBACK TO SAVEPOINT (verified on PostgreSQL 18), so recovery
+          // restores data state without invalidating the session's cache.
+          return await runSessionSavepoint(
+            port,
+            ATOMIC_PROGRAM_SAVEPOINT,
+            async () => {
+              const results: (readonly TRow[])[] = [];
+              for (const statement of statements) {
+                results.push(await executeCompiled<TRow>(statement));
+              }
+              return { action: "release", value: results };
+            },
+          );
         } catch (error) {
-          try {
-            // PostgreSQL protocol- and SQL-level prepared statements survive
-            // ROLLBACK TO SAVEPOINT (verified on PostgreSQL 18), so recovery
-            // restores data state without invalidating the session's cache.
-            await executeCompiled(ROLLBACK_ATOMIC_PROGRAM);
-            await executeCompiled(RELEASE_ATOMIC_PROGRAM);
-          } catch (recoveryError) {
-            const failures = new AggregateError(
-              [error, recoveryError],
-              "The atomic program failed and its savepoint recovery also failed.",
-              { cause: error },
-            );
+          if (error instanceof SavepointRecoveryFailure) {
             throw new CompilerInvariantError(
               "An atomic transaction session could not restore its savepoint after a failed program.",
               {},
-              { cause: failures },
+              { cause: error },
             );
           }
           throw error;

packages/typegraph/src/backend/drizzle/postgres.ts

@@ -151,6 +151,10 @@ import {
   nowIso,
   POSTGRES_ROW_MAPPER_CONFIG,
 } from "../row-mappers";
+import {
+  registerDriverSavepointScope,
+  runSessionSavepoint,
+} from "../session-savepoint";
 import {
   resolveDeclaredBackendResource,
   type SerializedResourceDeclaration,
@@ -250,6 +254,7 @@ import {
   type PostgresExecutionAdapter,
   type PostgresExecutionAdapterOptions,
 } from "./execution/postgres-execution";
+import { createPostgresTransactionSession } from "./execution/postgres-transaction-session";
 import { createSessionAtomicBatchAdapter } from "./execution/session-atomic-batch";
 import { createSerialExecutionAdapter } from "./execution/statement-queue";
 import { type ExecutableSql, toDrizzleSql } from "./execution/types";
@@ -2456,12 +2461,17 @@ function buildPostgresEngineProfileInternal(
                 // SAVEPOINT refuses outside an explicit transaction. Release it
                 // immediately: this proves the literal pinned session is active,
                 // without opening another transaction or changing application data.
+                // The scope issues nothing, so no driver scope is needed to
+                // attribute a failure inside it.
                 try {
-                  await externalTx.execute(
-                    sql`SAVEPOINT typegraph_schema_adoption_probe`,
-                  );
-                  await externalTx.execute(
-                    sql`RELEASE SAVEPOINT typegraph_schema_adoption_probe`,
+                  await runSessionSavepoint(
+                    {
+                      executeControl: (statement) =>
+                        externalTx.execute(sql.raw(statement)),
+                    },
+                    "typegraph_schema_adoption_probe",
+                    () =>
+                      Promise.resolve({ action: "release", value: undefined }),
                   );
                 } catch (error) {
                   throw new ConfigurationError(
@@ -4000,28 +4010,29 @@ function createTransactionBackend(
   // The wrap is unconditional because the overlap is created inside TypeGraph
   // (the node write pipeline syncs embeddings and fulltext concurrently), not
   // only by user code.
-  const operationExecutionAdapter = createPostgresExecutionAdapter(
+  const session = createPostgresTransactionSession(
     options.db,
     options.adapterOptions,
   );
-  const compiledExecutionAdapter = createPostgresExecutionAdapter(options.db, {
-    ...options.adapterOptions,
-    useTransactionClient: true,
-  });
-  const executeCompiled = compiledExecutionAdapter.executeCompiled;
+  const executeCompiled = session.compiled.executeCompiled;
+  const { enterDriverScope } = session;
   // Preserve the established Drizzle-routed execution surface for ordinary
   // transaction work. Only a closed atomic program needs the pinned client's
   // compiled-text seam; changing every statement to that seam would also
   // bypass driver logging and other Drizzle session behavior merely because
   // this transaction can execute programs.
   const txExecutionAdapter = createSerialExecutionAdapter(
-    operationExecutionAdapter,
+    session.operation,
     getPinnedPostgresTransactionClient(options.db) ?? options.db,
   );
   const runExclusive = txExecutionAdapter.runExclusive;
   const sessionAtomicBatchAdapter =
     executeCompiled === undefined ? undefined : (
-      createSessionAtomicBatchAdapter({ executeCompiled, runExclusive })
+      createSessionAtomicBatchAdapter({
+        executeCompiled,
+        runExclusive,
+        enterDriverScope,
+      })
     );
   const sessionAtomicSqlProgramExecutor =
     sessionAtomicBatchAdapter === undefined ? undefined : (
@@ -4089,6 +4100,9 @@ function createTransactionBackend(
     options.isFirstParty ?
       markFirstPartyFactory(txOperationBackend)
     : txOperationBackend;
+  if (enterDriverScope !== undefined) {
+    registerDriverSavepointScope(backend, enterDriverScope);
+  }
 
   if (sessionAtomicBatchAdapter !== undefined) {
     registerAtomicSqlProgram(backend, sessionAtomicBatchAdapter);