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

package/test/sql-runtime.test.ts

@@ -9,19 +9,28 @@ import {
 import type { SqlStorage } from '@prisma-next/sql-contract/types';
 import type {
   CodecRegistry,
-  SelectAst,
   SqlDriver,
   SqlExecuteRequest,
 } from '@prisma-next/sql-relational-core/ast';
-import { codec, createCodecRegistry } from '@prisma-next/sql-relational-core/ast';
+import {
+  BinaryExpr,
+  ColumnRef,
+  codec,
+  createCodecRegistry,
+  LiteralExpr,
+  SelectAst,
+  TableSource,
+} from '@prisma-next/sql-relational-core/ast';
+import type { SqlQueryPlan } from '@prisma-next/sql-relational-core/plan';
 import { describe, expect, it, vi } from 'vitest';
+import type { SqlMiddleware } from '../src/middleware/sql-middleware';
 import type {
   SqlRuntimeAdapterDescriptor,
   SqlRuntimeAdapterInstance,
   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 +44,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 +90,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 +123,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 +142,11 @@ function createMockDriver(): MockSqlDriver {
       rootExecute,
       connectionExecute,
       transactionExecute,
+      connectionRelease: connection.release,
+      connectionDestroy: connection.destroy,
+      transactionCommit: transaction.commit,
+      transactionRollback: transaction.rollback,
+      driverClose,
     },
   });
 }
@@ -287,6 +306,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({
@@ -352,16 +389,404 @@ describe('createRuntime', () => {
 
   it('rejects a Mongo middleware with a clear error', () => {
     const { stackInstance, context, driver } = createTestSetup();
+    // Simulate a caller bypassing the SqlMiddleware type constraint (e.g. dynamically-loaded
+    // middleware). Static typing already rejects familyId: 'mongo'; this tests the runtime guard.
+    const mongoMiddleware = { name: 'mongo-mw', familyId: 'mongo' } as unknown as SqlMiddleware;
     expect(() =>
       createRuntime({
         stackInstance,
         context,
         driver,
         verify: { mode: 'onFirstUse', requireMarker: false },
-        middleware: [{ name: 'mongo-mw', familyId: 'mongo' }],
+        middleware: [mongoMiddleware],
       }),
     ).toThrow(
       "Middleware 'mongo-mw' requires family 'mongo' but the runtime is configured for family 'sql'",
     );
   });
+
+  it('invokes beforeCompile and lowers the rewritten AST', async () => {
+    const { stackInstance, context, driver } = createTestSetup();
+    const debug = vi.fn();
+    const softDeletePredicate = BinaryExpr.eq(
+      ColumnRef.of('users', 'deleted_at'),
+      LiteralExpr.of(null),
+    );
+    const softDelete: SqlMiddleware = {
+      name: 'softDelete',
+      familyId: 'sql',
+      async beforeCompile(draft) {
+        if (draft.ast.kind !== 'select') return;
+        return { ...draft, ast: draft.ast.withWhere(softDeletePredicate) };
+      },
+    };
+
+    const runtime = createRuntime({
+      stackInstance,
+      context,
+      driver,
+      verify: { mode: 'onFirstUse', requireMarker: false },
+      middleware: [softDelete],
+      log: { info: vi.fn(), warn: vi.fn(), error: vi.fn(), debug },
+    });
+
+    const queryPlan: SqlQueryPlan = {
+      ast: SelectAst.from(TableSource.named('users')).withProjection([]),
+      params: [],
+      meta: {
+        target: 'postgres',
+        storageHash: testContract.storage.storageHash,
+        lane: 'dsl',
+        paramDescriptors: [],
+      },
+    };
+
+    await runtime.execute(queryPlan).toArray();
+
+    expect(driver.__spies.rootExecute).toHaveBeenCalledTimes(1);
+    const request = driver.__spies.rootExecute.mock.calls[0]?.[0] as SqlExecuteRequest;
+    expect(request.sql).toContain('deleted_at');
+    expect(debug).toHaveBeenCalledWith({
+      event: 'middleware.rewrite',
+      middleware: 'softDelete',
+      lane: 'dsl',
+    });
+  });
+
+  it('invokes adapter.lower exactly once per execute regardless of chain length', async () => {
+    const adapter = createStubAdapter();
+    const lowerSpy = vi.spyOn(adapter, 'lower');
+    const driver = createMockDriver();
+
+    const targetDescriptor = createTestTargetDescriptor();
+    const adapterDescriptor = createTestAdapterDescriptor(adapter);
+    const stack = createSqlExecutionStack({
+      target: targetDescriptor,
+      adapter: adapterDescriptor,
+      extensionPacks: [],
+    });
+    const stackInstance = instantiateExecutionStack(stack) as ExecutionStackInstance<
+      'sql',
+      'postgres',
+      SqlRuntimeAdapterInstance<'postgres'>,
+      RuntimeDriverInstance<'sql', 'postgres'>,
+      RuntimeExtensionInstance<'sql', 'postgres'>
+    >;
+    const context = createExecutionContext({
+      contract: testContract,
+      stack: { target: targetDescriptor, adapter: adapterDescriptor, extensionPacks: [] },
+    });
+
+    const rewriteA: SqlMiddleware = {
+      name: 'rewriteA',
+      familyId: 'sql',
+      async beforeCompile(draft) {
+        if (draft.ast.kind !== 'select') return undefined;
+        return {
+          ...draft,
+          ast: draft.ast.withWhere(BinaryExpr.eq(ColumnRef.of('users', 'a'), LiteralExpr.of(1))),
+        };
+      },
+    };
+    const rewriteB: SqlMiddleware = {
+      name: 'rewriteB',
+      familyId: 'sql',
+      async beforeCompile(draft) {
+        if (draft.ast.kind !== 'select') return undefined;
+        return {
+          ...draft,
+          ast: draft.ast.withWhere(BinaryExpr.eq(ColumnRef.of('users', 'b'), LiteralExpr.of(2))),
+        };
+      },
+    };
+
+    const runtime = createRuntime({
+      stackInstance,
+      context,
+      driver,
+      verify: { mode: 'onFirstUse', requireMarker: false },
+      middleware: [rewriteA, rewriteB],
+    });
+
+    const queryPlan: SqlQueryPlan = {
+      ast: SelectAst.from(TableSource.named('users')).withProjection([]),
+      params: [],
+      meta: {
+        target: 'postgres',
+        storageHash: testContract.storage.storageHash,
+        lane: 'dsl',
+        paramDescriptors: [],
+      },
+    };
+
+    await runtime.execute(queryPlan).toArray();
+
+    expect(lowerSpy).toHaveBeenCalledTimes(1);
+    const loweredAst = lowerSpy.mock.calls[0]?.[0] as SelectAst;
+    expect(loweredAst.where?.kind).toBe('binary');
+  });
+
+  it('skips beforeCompile for raw execution plans with no AST', async () => {
+    const { stackInstance, context, driver } = createTestSetup();
+    const debug = vi.fn();
+    const beforeCompile = vi.fn();
+    const observer: SqlMiddleware = {
+      name: 'observer',
+      familyId: 'sql',
+      beforeCompile,
+    };
+
+    const runtime = createRuntime({
+      stackInstance,
+      context,
+      driver,
+      verify: { mode: 'onFirstUse', requireMarker: false },
+      middleware: [observer],
+      log: { info: vi.fn(), warn: vi.fn(), error: vi.fn(), debug },
+    });
+
+    await runtime.execute(createRawExecutionPlan()).toArray();
+
+    expect(beforeCompile).not.toHaveBeenCalled();
+    expect(debug).not.toHaveBeenCalled();
+  });
+});
+
+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] : [] });
     },
   };
 }