@prisma-next/sql-runtime 0.4.0-dev.9 → 0.5.0-dev.1

package/src/lower-sql-plan.ts

@@ -21,11 +21,9 @@ export function lowerSqlPlan<Row>(
     params: queryPlan.params,
   });
 
-  const body = lowered.body;
-
   return Object.freeze({
-    sql: body.sql,
-    params: body.params ?? queryPlan.params,
+    sql: lowered.sql,
+    params: lowered.params ?? queryPlan.params,
     ast: queryPlan.ast,
     meta: queryPlan.meta,
   });

package/src/sql-runtime.ts

@@ -13,7 +13,11 @@ import type {
   RuntimeVerifyOptions,
   TelemetryOutcome,
 } from '@prisma-next/runtime-executor';
-import { AsyncIterableResult, createRuntimeCore } from '@prisma-next/runtime-executor';
+import {
+  AsyncIterableResult,
+  createRuntimeCore,
+  runtimeError,
+} from '@prisma-next/runtime-executor';
 import type { SqlStorage } from '@prisma-next/sql-contract/types';
 import type {
   Adapter,
@@ -73,7 +77,28 @@ export interface Runtime extends RuntimeQueryable {
 
 export interface RuntimeConnection extends RuntimeQueryable {
   transaction(): Promise<RuntimeTransaction>;
+  /**
+   * Returns the connection to the pool for reuse. Only call this when the
+   * connection is known to be in a clean state. If a transaction
+   * commit/rollback failed or the connection is otherwise suspect, call
+   * `destroy(reason)` instead.
+   */
   release(): Promise<void>;
+  /**
+   * Evicts the connection so it is never reused. Call this when the
+   * connection may be in an indeterminate state (e.g. a failed rollback
+   * leaving an open transaction, or a broken socket).
+   *
+   * If teardown fails the error is propagated and the connection remains
+   * retryable, so the caller can decide whether to swallow the failure or
+   * retry cleanup. Calling destroy() or release() more than once after a
+   * successful teardown is caller error.
+   *
+   * `reason` is advisory context only. It may be surfaced to driver-level
+   * observability hooks (e.g. pg-pool's `'release'` event) but does not
+   * influence eviction behavior and is not rethrown.
+   */
+  destroy(reason?: unknown): Promise<void>;
 }
 
 export interface RuntimeTransaction extends RuntimeQueryable {
@@ -87,6 +112,10 @@ export interface RuntimeQueryable {
   ): AsyncIterableResult<Row>;
 }
 
+export interface TransactionContext extends RuntimeQueryable {
+  readonly invalidated: boolean;
+}
+
 interface CoreQueryable {
   execute<Row = Record<string, unknown>>(plan: ExecutionPlan<Row>): AsyncIterableResult<Row>;
 }
@@ -206,6 +235,7 @@ class SqlRuntimeImpl<TContract extends Contract<SqlStorage> = Contract<SqlStorag
         };
       },
       release: coreConn.release.bind(coreConn),
+      destroy: coreConn.destroy.bind(coreConn),
       execute<Row = Record<string, unknown>>(
         plan: ExecutionPlan<Row> | SqlQueryPlan<Row>,
       ): AsyncIterableResult<Row> {
@@ -224,6 +254,114 @@ class SqlRuntimeImpl<TContract extends Contract<SqlStorage> = Contract<SqlStorag
   }
 }
 
+function transactionClosedError(): Error {
+  return runtimeError(
+    'RUNTIME.TRANSACTION_CLOSED',
+    'Cannot read from a query result after the transaction has ended. Await the result or call .toArray() inside the transaction callback.',
+    {},
+  );
+}
+
+export async function withTransaction<R>(
+  runtime: Runtime,
+  fn: (tx: TransactionContext) => PromiseLike<R>,
+): Promise<R> {
+  const connection = await runtime.connection();
+  const transaction = await connection.transaction();
+
+  let invalidated = false;
+  const txContext: TransactionContext = {
+    get invalidated() {
+      return invalidated;
+    },
+    execute<Row = Record<string, unknown>>(
+      plan: ExecutionPlan<Row> | SqlQueryPlan<Row>,
+    ): AsyncIterableResult<Row> {
+      if (invalidated) {
+        throw transactionClosedError();
+      }
+      const inner = transaction.execute(plan);
+      const guarded = async function* (): AsyncGenerator<Row, void, unknown> {
+        for await (const row of inner) {
+          if (invalidated) {
+            throw transactionClosedError();
+          }
+          yield row;
+        }
+      };
+      return new AsyncIterableResult(guarded());
+    },
+  };
+
+  let connectionDisposed = false;
+  const destroyConnection = async (reason: unknown): Promise<void> => {
+    if (connectionDisposed) return;
+    connectionDisposed = true;
+    // SqlConnection.destroy() propagates teardown errors so callers can
+    // decide what to do with them. Here, we're already about to throw a
+    // more informative error describing why we're evicting the connection
+    // (rollback/commit failure), so swallowing the teardown error is the
+    // right call — surfacing it would mask the original cause.
+    await connection.destroy(reason).catch(() => undefined);
+  };
+
+  try {
+    let result: R;
+    try {
+      result = await fn(txContext);
+    } catch (error) {
+      try {
+        await transaction.rollback();
+      } catch (rollbackError) {
+        await destroyConnection(rollbackError);
+        const wrapped = runtimeError(
+          'RUNTIME.TRANSACTION_ROLLBACK_FAILED',
+          'Transaction rollback failed after callback error',
+          { rollbackError },
+        );
+        wrapped.cause = error;
+        throw wrapped;
+      }
+      throw error;
+    } finally {
+      invalidated = true;
+    }
+
+    try {
+      await transaction.commit();
+    } catch (commitError) {
+      // After a failed COMMIT the server-side transaction may be: (a) already
+      // committed (error on response path), (b) already rolled back (deferred
+      // constraint / serialization failure), or (c) still open (COMMIT never
+      // reached the server). Attempt a best-effort rollback to cover (c) and
+      // confirm the protocol is healthy.
+      //
+      // If rollback succeeds, the server is definitely no longer in a
+      // transaction (no-op in (a)/(b), real cleanup in (c)) and we've just
+      // proved the connection round-trips correctly — it's safe to return
+      // to the pool. If rollback fails, the connection state is ambiguous
+      // (broken socket, protocol desync, etc.) and we must destroy it.
+      try {
+        await transaction.rollback();
+      } catch {
+        await destroyConnection(commitError);
+      }
+      const wrapped = runtimeError(
+        'RUNTIME.TRANSACTION_COMMIT_FAILED',
+        'Transaction commit failed',
+        { commitError },
+      );
+      wrapped.cause = commitError;
+      throw wrapped;
+    }
+    return result;
+  } finally {
+    if (!connectionDisposed) {
+      await connection.release();
+    }
+  }
+}
+
 export function createRuntime<TContract extends Contract<SqlStorage>, TTargetId extends string>(
   options: CreateRuntimeOptions<TContract, TTargetId>,
 ): Runtime {

package/test/execution-stack.test.ts

@@ -41,10 +41,7 @@ function createStubAdapterDescriptor(): SqlRuntimeAdapterDescriptor<'postgres'>
             readMarkerStatement: () => ({ sql: '', params: [] }),
           },
           lower() {
-            return {
-              profileId: 'test-profile',
-              body: Object.freeze({ sql: '', params: [] }),
-            };
+            return Object.freeze({ sql: '', params: [] });
           },
         },
       );

package/test/sql-runtime.test.ts

@@ -21,7 +21,7 @@ import type {
   SqlRuntimeTargetDescriptor,
 } from '../src/sql-context';
 import { createExecutionContext, createSqlExecutionStack } from '../src/sql-context';
-import { createRuntime } from '../src/sql-runtime';
+import { createRuntime, withTransaction } from '../src/sql-runtime';
 
 const testContract: Contract<SqlStorage> = {
   targetFamily: 'sql',
@@ -35,13 +35,18 @@ const testContract: Contract<SqlStorage> = {
   meta: {},
 };
 
-interface DriverExecuteSpies {
+interface DriverMockSpies {
   rootExecute: ReturnType<typeof vi.fn>;
   connectionExecute: ReturnType<typeof vi.fn>;
   transactionExecute: ReturnType<typeof vi.fn>;
+  connectionRelease: ReturnType<typeof vi.fn>;
+  connectionDestroy: ReturnType<typeof vi.fn>;
+  transactionCommit: ReturnType<typeof vi.fn>;
+  transactionRollback: ReturnType<typeof vi.fn>;
+  driverClose: ReturnType<typeof vi.fn>;
 }
 
-type MockSqlDriver = SqlDriver & { __spies: DriverExecuteSpies };
+type MockSqlDriver = SqlDriver & { __spies: DriverMockSpies };
 
 function createStubCodecs(): CodecRegistry {
   const registry = createCodecRegistry();
@@ -76,10 +81,7 @@ function createStubAdapter() {
       },
     },
     lower(ast: SelectAst) {
-      return {
-        profileId: 'test-profile',
-        body: Object.freeze({ sql: JSON.stringify(ast), params: [] }),
-      };
+      return Object.freeze({ sql: JSON.stringify(ast), params: [] });
     },
   };
 }
@@ -112,15 +114,18 @@ function createMockDriver(): MockSqlDriver {
     execute: connectionExecute,
     query,
     release: vi.fn().mockResolvedValue(undefined),
+    destroy: vi.fn().mockResolvedValue(undefined),
     beginTransaction: vi.fn().mockResolvedValue(transaction),
   };
 
+  const driverClose = vi.fn().mockResolvedValue(undefined);
+
   const driver: SqlDriver = {
     execute: rootExecute,
     query,
     connect: vi.fn().mockImplementation(async (_binding?: undefined) => undefined),
     acquireConnection: vi.fn().mockResolvedValue(connection),
-    close: vi.fn().mockResolvedValue(undefined),
+    close: driverClose,
   };
 
   return Object.assign(driver, {
@@ -128,6 +133,11 @@ function createMockDriver(): MockSqlDriver {
       rootExecute,
       connectionExecute,
       transactionExecute,
+      connectionRelease: connection.release,
+      connectionDestroy: connection.destroy,
+      transactionCommit: transaction.commit,
+      transactionRollback: transaction.rollback,
+      driverClose,
     },
   });
 }
@@ -287,6 +297,24 @@ describe('createRuntime', () => {
     await connection.release();
   });
 
+  it('delegates connection.destroy() to the driver connection', async () => {
+    const { stackInstance, context, driver } = createTestSetup();
+    const runtime = createRuntime({
+      stackInstance,
+      context,
+      driver,
+      verify: { mode: 'onFirstUse', requireMarker: false },
+    });
+
+    const connection = await runtime.connection();
+    const reason = new Error('bad state');
+    await connection.destroy(reason);
+
+    expect(driver.__spies.connectionDestroy).toHaveBeenCalledOnce();
+    expect(driver.__spies.connectionDestroy).toHaveBeenCalledWith(reason);
+    expect(driver.__spies.connectionRelease).not.toHaveBeenCalled();
+  });
+
   it('uses transaction queryable for transaction.execute', async () => {
     const { stackInstance, context, driver } = createTestSetup();
     const runtime = createRuntime({
@@ -365,3 +393,242 @@ describe('createRuntime', () => {
     );
   });
 });
+
+describe('withTransaction', () => {
+  function createRuntimeForTransaction() {
+    const { stackInstance, context, driver } = createTestSetup();
+    const runtime = createRuntime({
+      stackInstance,
+      context,
+      driver,
+      verify: { mode: 'onFirstUse', requireMarker: false },
+    });
+    return { runtime, driver };
+  }
+
+  it('commits on successful callback and returns the result', async () => {
+    const { runtime, driver } = createRuntimeForTransaction();
+
+    const result = await withTransaction(runtime, async (tx) => {
+      await tx.execute(createRawExecutionPlan()).toArray();
+      return 42;
+    });
+
+    expect(result).toBe(42);
+    expect(driver.__spies.transactionCommit).toHaveBeenCalledOnce();
+    expect(driver.__spies.transactionRollback).not.toHaveBeenCalled();
+    expect(driver.__spies.connectionRelease).toHaveBeenCalledOnce();
+  });
+
+  it('rolls back on callback error and re-throws', async () => {
+    const { runtime, driver } = createRuntimeForTransaction();
+    const error = new Error('test error');
+
+    await expect(
+      withTransaction(runtime, async () => {
+        throw error;
+      }),
+    ).rejects.toBe(error);
+
+    expect(driver.__spies.transactionRollback).toHaveBeenCalledOnce();
+    expect(driver.__spies.transactionCommit).not.toHaveBeenCalled();
+    expect(driver.__spies.connectionRelease).toHaveBeenCalledOnce();
+  });
+
+  it('releases connection after commit', async () => {
+    const { runtime, driver } = createRuntimeForTransaction();
+
+    await withTransaction(runtime, async () => 'ok');
+
+    expect(driver.__spies.connectionRelease).toHaveBeenCalledOnce();
+  });
+
+  it('releases connection after rollback', async () => {
+    const { runtime, driver } = createRuntimeForTransaction();
+
+    await withTransaction(runtime, async () => {
+      throw new Error('fail');
+    }).catch(() => {});
+
+    expect(driver.__spies.connectionRelease).toHaveBeenCalledOnce();
+  });
+
+  it('wraps commit failure and exposes the original error as cause', async () => {
+    const { runtime, driver } = createRuntimeForTransaction();
+    const commitError = new Error('commit failed');
+    driver.__spies.transactionCommit.mockRejectedValueOnce(commitError);
+
+    const result = withTransaction(runtime, async () => 'value');
+
+    await expect(result).rejects.toMatchObject({
+      code: 'RUNTIME.TRANSACTION_COMMIT_FAILED',
+      cause: commitError,
+    });
+  });
+
+  it('attempts best-effort rollback after commit fails and releases when it succeeds', async () => {
+    const { runtime, driver } = createRuntimeForTransaction();
+    const commitError = new Error('commit failed');
+    driver.__spies.transactionCommit.mockRejectedValueOnce(commitError);
+
+    await withTransaction(runtime, async () => 'value').catch(() => {});
+
+    expect(driver.__spies.transactionCommit).toHaveBeenCalledOnce();
+    expect(driver.__spies.transactionRollback).toHaveBeenCalledOnce();
+    // A successful rollback after a failed commit means the server is no
+    // longer in a transaction and the connection round-tripped cleanly, so
+    // it is safe to return to the pool rather than evict it.
+    expect(driver.__spies.connectionRelease).toHaveBeenCalledOnce();
+    expect(driver.__spies.connectionDestroy).not.toHaveBeenCalled();
+  });
+
+  it('forwards the callback return value', async () => {
+    const { runtime } = createRuntimeForTransaction();
+
+    const result = await withTransaction(runtime, async () => ({
+      name: 'test',
+      count: 3,
+    }));
+
+    expect(result).toEqual({ name: 'test', count: 3 });
+  });
+
+  it('executes queries against the transaction', async () => {
+    const { runtime, driver } = createRuntimeForTransaction();
+
+    await withTransaction(runtime, async (tx) => {
+      await tx.execute(createRawExecutionPlan()).toArray();
+    });
+
+    expect(driver.__spies.transactionExecute).toHaveBeenCalledOnce();
+    expect(driver.__spies.rootExecute).not.toHaveBeenCalled();
+    expect(driver.__spies.connectionExecute).not.toHaveBeenCalled();
+  });
+
+  it('throws on execute after commit (invalidation)', async () => {
+    const { runtime } = createRuntimeForTransaction();
+    let savedTx: { execute: (plan: ExecutionPlan) => unknown } | undefined;
+
+    await withTransaction(runtime, async (tx) => {
+      savedTx = tx;
+    });
+
+    expect(() => savedTx!.execute(createRawExecutionPlan())).toThrow(
+      'Cannot read from a query result after the transaction has ended',
+    );
+  });
+
+  it('throws on iteration of escaped AsyncIterableResult after commit', async () => {
+    const { runtime } = createRuntimeForTransaction();
+
+    const escaped = await withTransaction(runtime, async (tx) => {
+      return { result: tx.execute(createRawExecutionPlan()) };
+    });
+
+    await expect(escaped.result.toArray()).rejects.toThrow(
+      'Cannot read from a query result after the transaction has ended',
+    );
+  });
+
+  it('sets invalidated flag after commit', async () => {
+    const { runtime } = createRuntimeForTransaction();
+    let txRef: { invalidated: boolean } | undefined;
+
+    await withTransaction(runtime, async (tx) => {
+      expect(tx.invalidated).toBe(false);
+      txRef = tx;
+    });
+
+    expect(txRef!.invalidated).toBe(true);
+  });
+
+  it('wraps original error when rollback fails', async () => {
+    const { runtime, driver } = createRuntimeForTransaction();
+    const callbackError = new Error('callback failed');
+    const rollbackError = new Error('rollback failed');
+    driver.__spies.transactionRollback.mockRejectedValueOnce(rollbackError);
+
+    const rejection = withTransaction(runtime, async () => {
+      throw callbackError;
+    });
+
+    await expect(rejection).rejects.toThrow('Transaction rollback failed after callback error');
+    await expect(rejection).rejects.toMatchObject({
+      code: 'RUNTIME.TRANSACTION_ROLLBACK_FAILED',
+      cause: callbackError,
+      details: { rollbackError },
+    });
+    expect(driver.__spies.connectionDestroy).toHaveBeenCalledOnce();
+    expect(driver.__spies.connectionRelease).not.toHaveBeenCalled();
+  });
+
+  it('destroys connection when rollback fails even if destroy also fails', async () => {
+    const { runtime, driver } = createRuntimeForTransaction();
+    const callbackError = new Error('callback failed');
+    const rollbackError = new Error('rollback failed');
+    const destroyError = new Error('destroy failed');
+    driver.__spies.transactionRollback.mockRejectedValueOnce(rollbackError);
+    driver.__spies.connectionDestroy.mockRejectedValueOnce(destroyError);
+
+    const rejection = withTransaction(runtime, async () => {
+      throw callbackError;
+    });
+
+    await expect(rejection).rejects.toMatchObject({
+      code: 'RUNTIME.TRANSACTION_ROLLBACK_FAILED',
+      cause: callbackError,
+      details: { rollbackError },
+    });
+    expect(driver.__spies.connectionDestroy).toHaveBeenCalledOnce();
+    expect(driver.__spies.connectionRelease).not.toHaveBeenCalled();
+  });
+
+  it('destroys connection when commit fails and best-effort rollback also fails', async () => {
+    const { runtime, driver } = createRuntimeForTransaction();
+    const commitError = new Error('commit failed');
+    const rollbackError = new Error('rollback also failed');
+    driver.__spies.transactionCommit.mockRejectedValueOnce(commitError);
+    driver.__spies.transactionRollback.mockRejectedValueOnce(rollbackError);
+
+    const rejection = withTransaction(runtime, async () => 'value');
+
+    await expect(rejection).rejects.toMatchObject({
+      code: 'RUNTIME.TRANSACTION_COMMIT_FAILED',
+      cause: commitError,
+    });
+    expect(driver.__spies.connectionDestroy).toHaveBeenCalledOnce();
+    expect(driver.__spies.connectionRelease).not.toHaveBeenCalled();
+  });
+
+  it('sets invalidated flag after rollback', async () => {
+    const { runtime } = createRuntimeForTransaction();
+    let txRef: { invalidated: boolean } | undefined;
+
+    await withTransaction(runtime, async (tx) => {
+      txRef = tx;
+      throw new Error('fail');
+    }).catch(() => {});
+
+    expect(txRef!.invalidated).toBe(true);
+  });
+
+  it('releases connection independently across sequential transactions', async () => {
+    const { runtime, driver } = createRuntimeForTransaction();
+
+    await withTransaction(runtime, async (tx) => {
+      await tx.execute(createRawExecutionPlan()).toArray();
+    });
+
+    await withTransaction(runtime, async (tx) => {
+      await tx.execute(createRawExecutionPlan()).toArray();
+    });
+
+    await withTransaction(runtime, async () => {
+      throw new Error('fail');
+    }).catch(() => {});
+
+    expect(driver.__spies.connectionRelease).toHaveBeenCalledTimes(3);
+    expect(driver.__spies.transactionCommit).toHaveBeenCalledTimes(2);
+    expect(driver.__spies.transactionRollback).toHaveBeenCalledTimes(1);
+  });
+});

package/test/utils.ts

@@ -132,7 +132,7 @@ export function createTestAdapterDescriptor(
     codecs: () => codecRegistry,
     parameterizedCodecs: () => [],
     mutationDefaultGenerators: createTestMutationDefaultGenerators,
-    create(): SqlRuntimeAdapterInstance<'postgres'> {
+    create(_stack): SqlRuntimeAdapterInstance<'postgres'> {
       return Object.assign({ familyId: 'sql' as const, targetId: 'postgres' as const }, adapter);
     },
   };
@@ -256,10 +256,7 @@ export function createStubAdapter(): Adapter<SelectAst, Contract<SqlStorage>, Lo
     },
     lower(ast: SelectAst, ctx: { contract: Contract<SqlStorage>; params?: readonly unknown[] }) {
       const sqlText = JSON.stringify(ast);
-      return {
-        profileId: this.profile.id,
-        body: Object.freeze({ sql: sqlText, params: ctx.params ? [...ctx.params] : [] }),
-      };
+      return Object.freeze({ sql: sqlText, params: ctx.params ? [...ctx.params] : [] });
     },
   };
 }