@prisma-next/sql-runtime 0.3.0-pr.99.6 → 0.4.0-dev.1

package/test/budgets.test.ts

@@ -0,0 +1,431 @@
+import type { ExecutionPlan, PlanMeta } from '@prisma-next/contract/types';
+import type { AfterExecuteResult, MiddlewareContext } from '@prisma-next/runtime-executor';
+import {
+  AggregateExpr,
+  ColumnRef,
+  DeleteAst,
+  ProjectionItem,
+  SelectAst,
+  TableSource,
+} from '@prisma-next/sql-relational-core/ast';
+import { timeouts } from '@prisma-next/test-utils';
+import { describe, expect, it, vi } from 'vitest';
+import { budgets } from '../src/middleware/budgets';
+
+const userTable = TableSource.named('user');
+const idCol = ColumnRef.of('user', 'id');
+
+function createMiddlewareContext(
+  overrides?: Partial<MiddlewareContext<unknown>>,
+): MiddlewareContext<unknown> {
+  return {
+    contract: {},
+    mode: 'strict' as const,
+    now: () => Date.now(),
+    log: {
+      info: vi.fn(),
+      warn: vi.fn(),
+      error: vi.fn(),
+    },
+    ...overrides,
+  };
+}
+
+const baseMeta: PlanMeta = {
+  target: 'postgres',
+  storageHash: 'sha256:test',
+  lane: 'dsl',
+  paramDescriptors: [],
+};
+
+type PlanOverrides = Partial<Omit<ExecutionPlan, 'meta'>> & { meta?: Partial<PlanMeta> };
+
+function createPlan(overrides: PlanOverrides): ExecutionPlan {
+  const { meta: metaOverrides, ...rest } = overrides;
+  return {
+    sql: 'SELECT 1',
+    params: [],
+    meta: { ...baseMeta, ...(metaOverrides ?? {}) } as unknown as PlanMeta,
+    ...rest,
+  } as unknown as ExecutionPlan;
+}
+
+describe('budgets middleware', () => {
+  describe('heuristic row budget (no AST)', () => {
+    it(
+      'throws for unbounded raw SELECT exceeding budget',
+      async () => {
+        const plan = createPlan({
+          sql: 'SELECT id, email FROM "user"',
+          meta: { refs: { tables: ['user'] } },
+        });
+        const mw = budgets({ maxRows: 50, defaultTableRows: 10_000 });
+        const ctx = createMiddlewareContext();
+
+        await expect(mw.beforeExecute?.(plan, ctx)).rejects.toMatchObject({
+          code: 'BUDGET.ROWS_EXCEEDED',
+          category: 'BUDGET',
+        });
+      },
+      timeouts.default,
+    );
+
+    it(
+      'throws for unbounded raw SELECT even without table refs',
+      async () => {
+        const plan = createPlan({
+          sql: 'SELECT 1',
+        });
+        const mw = budgets({ maxRows: 50 });
+        const ctx = createMiddlewareContext();
+
+        await expect(mw.beforeExecute?.(plan, ctx)).rejects.toMatchObject({
+          code: 'BUDGET.ROWS_EXCEEDED',
+          category: 'BUDGET',
+        });
+      },
+      timeouts.default,
+    );
+
+    it(
+      'allows bounded raw SELECT with limit annotation within budget',
+      async () => {
+        const plan = createPlan({
+          sql: 'SELECT id FROM "user" LIMIT 5',
+          meta: {
+            refs: { tables: ['user'] },
+            annotations: { limit: 5 },
+          },
+        });
+        const mw = budgets({ maxRows: 10_000, defaultTableRows: 10_000 });
+        const ctx = createMiddlewareContext();
+
+        await mw.beforeExecute?.(plan, ctx);
+        expect(ctx.log.warn).not.toHaveBeenCalled();
+      },
+      timeouts.default,
+    );
+
+    it(
+      'throws when estimated rows exceed budget for bounded query',
+      async () => {
+        const plan = createPlan({
+          sql: 'SELECT id FROM "user" LIMIT 500',
+          meta: {
+            refs: { tables: ['user'] },
+            annotations: { limit: 500 },
+          },
+        });
+        const mw = budgets({ maxRows: 50, defaultTableRows: 10_000 });
+        const ctx = createMiddlewareContext();
+
+        await expect(mw.beforeExecute?.(plan, ctx)).rejects.toMatchObject({
+          code: 'BUDGET.ROWS_EXCEEDED',
+          category: 'BUDGET',
+          details: expect.objectContaining({ source: 'heuristic' }),
+        });
+      },
+      timeouts.default,
+    );
+
+    it(
+      'uses tableRows config for estimation',
+      async () => {
+        const plan = createPlan({
+          sql: 'SELECT id FROM "user"',
+          meta: { refs: { tables: ['user'] } },
+        });
+        const mw = budgets({ maxRows: 100, tableRows: { user: 50 } });
+        const ctx = createMiddlewareContext();
+
+        await expect(mw.beforeExecute?.(plan, ctx)).rejects.toMatchObject({
+          code: 'BUDGET.ROWS_EXCEEDED',
+        });
+      },
+      timeouts.default,
+    );
+
+    it(
+      'does not check row budget for non-SELECT statements',
+      async () => {
+        const plan = createPlan({
+          sql: 'INSERT INTO "user" (id, email) VALUES ($1, $2)',
+        });
+        const mw = budgets({ maxRows: 1 });
+        const ctx = createMiddlewareContext();
+
+        await mw.beforeExecute?.(plan, ctx);
+      },
+      timeouts.default,
+    );
+  });
+
+  describe('observed row count (onRow)', () => {
+    it(
+      'throws when observed rows exceed budget',
+      async () => {
+        const mw = budgets({ maxRows: 2 });
+        const plan = createPlan({
+          sql: 'INSERT INTO "user" (id) VALUES ($1)',
+        });
+        const ctx = createMiddlewareContext();
+
+        await mw.beforeExecute?.(plan, ctx);
+        await mw.onRow?.({}, plan, ctx);
+        await mw.onRow?.({}, plan, ctx);
+        await expect(mw.onRow?.({}, plan, ctx)).rejects.toMatchObject({
+          code: 'BUDGET.ROWS_EXCEEDED',
+          details: expect.objectContaining({ source: 'observed' }),
+        });
+      },
+      timeouts.default,
+    );
+
+    it(
+      'tracks row counts independently per execution plan',
+      async () => {
+        const mw = budgets({ maxRows: 2 });
+        const planA = createPlan({ sql: 'INSERT INTO "user" (id) VALUES ($1)' });
+        const planB = createPlan({ sql: 'INSERT INTO "user" (id) VALUES ($2)' });
+        const ctxA = createMiddlewareContext();
+        const ctxB = createMiddlewareContext();
+
+        await mw.beforeExecute?.(planA, ctxA);
+        await mw.beforeExecute?.(planB, ctxB);
+
+        await mw.onRow?.({}, planA, ctxA);
+        await mw.onRow?.({}, planB, ctxB);
+        await mw.onRow?.({}, planA, ctxA);
+        await mw.onRow?.({}, planB, ctxB);
+
+        await expect(mw.onRow?.({}, planA, ctxA)).rejects.toMatchObject({
+          code: 'BUDGET.ROWS_EXCEEDED',
+        });
+        await expect(mw.onRow?.({}, planB, ctxB)).rejects.toMatchObject({
+          code: 'BUDGET.ROWS_EXCEEDED',
+        });
+      },
+      timeouts.default,
+    );
+  });
+
+  describe('latency budget (afterExecute)', () => {
+    it(
+      'warns when latency exceeds budget in non-strict mode',
+      async () => {
+        const mw = budgets({ maxLatencyMs: 100, severities: { latency: 'warn' } });
+        const plan = createPlan({ sql: 'SELECT 1', meta: { annotations: { limit: 1 } } });
+        const ctx = createMiddlewareContext({ mode: 'permissive' });
+        const result: AfterExecuteResult = { rowCount: 1, latencyMs: 200, completed: true };
+
+        await mw.afterExecute?.(plan, result, ctx);
+        expect(ctx.log.warn).toHaveBeenCalledWith(
+          expect.objectContaining({ code: 'BUDGET.TIME_EXCEEDED' }),
+        );
+      },
+      timeouts.default,
+    );
+
+    it(
+      'throws when latency exceeds budget in strict mode with error severity',
+      async () => {
+        const mw = budgets({ maxLatencyMs: 100, severities: { latency: 'error' } });
+        const plan = createPlan({ sql: 'SELECT 1', meta: { annotations: { limit: 1 } } });
+        const ctx = createMiddlewareContext({ mode: 'strict' });
+        const result: AfterExecuteResult = { rowCount: 1, latencyMs: 200, completed: true };
+
+        await expect(mw.afterExecute?.(plan, result, ctx)).rejects.toMatchObject({
+          code: 'BUDGET.TIME_EXCEEDED',
+          category: 'BUDGET',
+        });
+      },
+      timeouts.default,
+    );
+
+    it(
+      'throws when latency exceeds budget in strict mode even with warn severity',
+      async () => {
+        const mw = budgets({ maxLatencyMs: 100, severities: { latency: 'warn' } });
+        const plan = createPlan({ sql: 'SELECT 1', meta: { annotations: { limit: 1 } } });
+        const ctx = createMiddlewareContext({ mode: 'strict' });
+        const result: AfterExecuteResult = { rowCount: 1, latencyMs: 200, completed: true };
+
+        await expect(mw.afterExecute?.(plan, result, ctx)).rejects.toMatchObject({
+          code: 'BUDGET.TIME_EXCEEDED',
+          category: 'BUDGET',
+        });
+      },
+      timeouts.default,
+    );
+
+    it(
+      'does not warn when latency is within budget',
+      async () => {
+        const mw = budgets({ maxLatencyMs: 1000 });
+        const plan = createPlan({ sql: 'SELECT 1', meta: { annotations: { limit: 1 } } });
+        const ctx = createMiddlewareContext();
+        const result: AfterExecuteResult = { rowCount: 1, latencyMs: 50, completed: true };
+
+        await mw.afterExecute?.(plan, result, ctx);
+        expect(ctx.log.warn).not.toHaveBeenCalled();
+      },
+      timeouts.default,
+    );
+  });
+
+  describe('severity configuration', () => {
+    it(
+      'warns instead of throwing when rowCount severity is warn and mode is permissive',
+      async () => {
+        const plan = createPlan({
+          sql: 'SELECT id FROM "user"',
+          meta: { refs: { tables: ['user'] } },
+        });
+        const mw = budgets({
+          maxRows: 50,
+          defaultTableRows: 10_000,
+          severities: { rowCount: 'warn' },
+        });
+        const ctx = createMiddlewareContext({ mode: 'permissive' });
+
+        await mw.beforeExecute?.(plan, ctx);
+        expect(ctx.log.warn).toHaveBeenCalledWith(
+          expect.objectContaining({ code: 'BUDGET.ROWS_EXCEEDED' }),
+        );
+      },
+      timeouts.default,
+    );
+  });
+
+  describe('AST-based row budget', () => {
+    it(
+      'allows bounded SelectAst with limit within budget',
+      async () => {
+        const ast = SelectAst.from(userTable)
+          .withProjection([ProjectionItem.of('id', idCol)])
+          .withLimit(5);
+        const plan = createPlan({
+          ast,
+          meta: { refs: { tables: ['user'] } },
+        });
+        const mw = budgets({ maxRows: 10_000, defaultTableRows: 10_000 });
+        const ctx = createMiddlewareContext();
+
+        await mw.beforeExecute?.(plan, ctx);
+        expect(ctx.log.warn).not.toHaveBeenCalled();
+      },
+      timeouts.default,
+    );
+
+    it(
+      'throws for unbounded SelectAst without limit',
+      async () => {
+        const ast = SelectAst.from(userTable).withProjection([ProjectionItem.of('id', idCol)]);
+        const plan = createPlan({
+          ast,
+          meta: { refs: { tables: ['user'] } },
+        });
+        const mw = budgets({ maxRows: 50, defaultTableRows: 10_000 });
+        const ctx = createMiddlewareContext();
+
+        await expect(mw.beforeExecute?.(plan, ctx)).rejects.toMatchObject({
+          code: 'BUDGET.ROWS_EXCEEDED',
+          category: 'BUDGET',
+          details: expect.objectContaining({ source: 'ast' }),
+        });
+      },
+      timeouts.default,
+    );
+
+    it(
+      'throws for unbounded SelectAst without table refs',
+      async () => {
+        const ast = SelectAst.from(userTable).withProjection([ProjectionItem.of('id', idCol)]);
+        const plan = createPlan({ ast });
+        const mw = budgets({ maxRows: 50 });
+        const ctx = createMiddlewareContext();
+
+        await expect(mw.beforeExecute?.(plan, ctx)).rejects.toMatchObject({
+          code: 'BUDGET.ROWS_EXCEEDED',
+          details: expect.objectContaining({ source: 'ast' }),
+        });
+      },
+      timeouts.default,
+    );
+
+    it(
+      'reads limit from AST, not from annotations',
+      async () => {
+        const ast = SelectAst.from(userTable)
+          .withProjection([ProjectionItem.of('id', idCol)])
+          .withLimit(5);
+        const plan = createPlan({
+          ast,
+          meta: {
+            refs: { tables: ['user'] },
+            annotations: { limit: 99999 },
+          },
+        });
+        const mw = budgets({ maxRows: 10_000, defaultTableRows: 10_000 });
+        const ctx = createMiddlewareContext();
+
+        await mw.beforeExecute?.(plan, ctx);
+        expect(ctx.log.warn).not.toHaveBeenCalled();
+      },
+      timeouts.default,
+    );
+
+    it(
+      'does not check row budget for non-SelectAst (e.g. DeleteAst)',
+      async () => {
+        const ast = DeleteAst.from(userTable);
+        const plan = createPlan({ ast });
+        const mw = budgets({ maxRows: 1 });
+        const ctx = createMiddlewareContext();
+
+        await mw.beforeExecute?.(plan, ctx);
+      },
+      timeouts.default,
+    );
+
+    it(
+      'estimates 1 row for aggregate without GROUP BY',
+      async () => {
+        const ast = SelectAst.from(userTable).withProjection([
+          ProjectionItem.of('count', AggregateExpr.count()),
+        ]);
+        const plan = createPlan({
+          ast,
+          meta: { refs: { tables: ['user'] } },
+        });
+        const mw = budgets({ maxRows: 1, defaultTableRows: 10_000 });
+        const ctx = createMiddlewareContext();
+
+        await mw.beforeExecute?.(plan, ctx);
+        expect(ctx.log.warn).not.toHaveBeenCalled();
+      },
+      timeouts.default,
+    );
+
+    it(
+      'does not reduce estimate for aggregate with GROUP BY',
+      async () => {
+        const ast = SelectAst.from(userTable)
+          .withProjection([ProjectionItem.of('count', AggregateExpr.count())])
+          .withGroupBy([idCol]);
+        const plan = createPlan({
+          ast,
+          meta: { refs: { tables: ['user'] } },
+        });
+        const mw = budgets({ maxRows: 50, defaultTableRows: 10_000 });
+        const ctx = createMiddlewareContext();
+
+        await expect(mw.beforeExecute?.(plan, ctx)).rejects.toMatchObject({
+          code: 'BUDGET.ROWS_EXCEEDED',
+          details: expect.objectContaining({ source: 'ast' }),
+        });
+      },
+      timeouts.default,
+    );
+  });
+});

package/test/context.types.test-d.ts

@@ -1,10 +1,12 @@
-import type { SqlContract, SqlMappings, SqlStorage } from '@prisma-next/sql-contract/types';
+import type { Contract, StorageHashBase } from '@prisma-next/contract/types';
+import type { SqlStorage } from '@prisma-next/sql-contract/types';
 import { expectTypeOf, test } from 'vitest';
-import type { RuntimeContext, TypeHelperRegistry } from '../src/sql-context';
+import type { ExecutionContext, TypeHelperRegistry } from '../src/sql-context';
 
 // Contract type with storage.types using literal types (matching emission output)
-type TestContract = SqlContract<
+type TestContract = Contract<
   {
+    readonly storageHash: StorageHashBase<string>;
     readonly tables: {
       readonly document: {
         readonly columns: {
@@ -28,28 +30,24 @@ type TestContract = SqlContract<
       };
     };
   },
-  Record<string, never>,
-  Record<string, never>,
-  SqlMappings
+  Record<string, never>
 >;
 
-test('RuntimeContext.types is TypeHelperRegistry', () => {
-  // RuntimeContext.types is intentionally loose (Record<string, unknown>)
+test('ExecutionContext.types is TypeHelperRegistry', () => {
+  // ExecutionContext.types is intentionally loose (Record<string, unknown>)
   // The strong typing comes from schema(context).types via ExtractSchemaTypes
-  expectTypeOf<RuntimeContext<TestContract>['types']>().toEqualTypeOf<
-    TypeHelperRegistry | undefined
-  >();
+  expectTypeOf<ExecutionContext<TestContract>['types']>().toEqualTypeOf<TypeHelperRegistry>();
 
   // TypeHelperRegistry allows any values - the actual type depends on init hooks
   expectTypeOf<TypeHelperRegistry>().toEqualTypeOf<Record<string, unknown>>();
 });
 
-test('RuntimeContext preserves contract type parameter', () => {
-  // Verify the contract type is preserved in RuntimeContext
-  expectTypeOf<RuntimeContext<TestContract>['contract']>().toEqualTypeOf<TestContract>();
+test('ExecutionContext preserves contract type parameter', () => {
+  // Verify the contract type is preserved in ExecutionContext
+  expectTypeOf<ExecutionContext<TestContract>['contract']>().toEqualTypeOf<TestContract>();
 
   // Verify we can access storage.types through the context's contract
-  type ContractStorageTypes = RuntimeContext<TestContract>['contract']['storage']['types'];
+  type ContractStorageTypes = ExecutionContext<TestContract>['contract']['storage']['types'];
   expectTypeOf<ContractStorageTypes>().toExtend<
     | {
         readonly Vector1536: {
@@ -62,9 +60,9 @@ test('RuntimeContext preserves contract type parameter', () => {
   >();
 });
 
-test('RuntimeContext accepts generic SqlContract', () => {
-  // Verify RuntimeContext defaults work
-  type DefaultContext = RuntimeContext;
-  expectTypeOf<DefaultContext['contract']>().toExtend<SqlContract<SqlStorage>>();
-  expectTypeOf<DefaultContext['types']>().toEqualTypeOf<TypeHelperRegistry | undefined>();
+test('ExecutionContext accepts generic Contract', () => {
+  // Verify ExecutionContext defaults work
+  type DefaultContext = ExecutionContext;
+  expectTypeOf<DefaultContext['contract']>().toExtend<Contract<SqlStorage>>();
+  expectTypeOf<DefaultContext['types']>().toEqualTypeOf<TypeHelperRegistry>();
 });

package/test/execution-stack.test.ts

@@ -0,0 +1,164 @@
+import { createExecutionStack } from '@prisma-next/framework-components/execution';
+import { codec, createCodecRegistry } from '@prisma-next/sql-relational-core/ast';
+import { describe, expect, it } from 'vitest';
+import { createExecutionContext, createSqlExecutionStack } from '../src/exports';
+import type {
+  ExecutionContext,
+  SqlRuntimeAdapterDescriptor,
+  SqlRuntimeExtensionDescriptor,
+  SqlRuntimeTargetDescriptor,
+} from '../src/sql-context';
+import { createTestContract } from './utils';
+
+function createStubAdapterDescriptor(): SqlRuntimeAdapterDescriptor<'postgres'> {
+  const registry = createCodecRegistry();
+  registry.register(
+    codec({
+      typeId: 'pg/text@1',
+      targetTypes: ['text'],
+      encode: (value: string) => value,
+      decode: (wire: string) => wire,
+    }),
+  );
+
+  return {
+    kind: 'adapter',
+    id: 'test-adapter',
+    version: '0.0.1',
+    familyId: 'sql' as const,
+    targetId: 'postgres' as const,
+    codecs: () => registry,
+    parameterizedCodecs: () => [],
+    create() {
+      return Object.assign(
+        { familyId: 'sql' as const, targetId: 'postgres' as const },
+        {
+          profile: {
+            id: 'test-profile',
+            target: 'postgres',
+            capabilities: {},
+            codecs: () => registry,
+            readMarkerStatement: () => ({ sql: '', params: [] }),
+          },
+          lower() {
+            return {
+              profileId: 'test-profile',
+              body: Object.freeze({ sql: '', params: [] }),
+            };
+          },
+        },
+      );
+    },
+  };
+}
+
+function createStubTargetDescriptor(): SqlRuntimeTargetDescriptor<'postgres'> {
+  return {
+    kind: 'target',
+    id: 'postgres',
+    version: '0.0.1',
+    familyId: 'sql' as const,
+    targetId: 'postgres' as const,
+    codecs: () => createCodecRegistry(),
+    parameterizedCodecs: () => [],
+    create() {
+      return { familyId: 'sql' as const, targetId: 'postgres' as const };
+    },
+  };
+}
+
+function createStubExtensionDescriptor(): SqlRuntimeExtensionDescriptor<'postgres'> {
+  const registry = createCodecRegistry();
+  registry.register(
+    codec({
+      typeId: 'pg/uuid@1',
+      targetTypes: ['uuid'],
+      encode: (value: string) => value,
+      decode: (wire: string) => wire,
+    }),
+  );
+
+  const operations = [
+    {
+      method: 'example',
+      args: [] as const,
+      returns: { codecId: 'pg/text@1', nullable: false },
+      lowering: {
+        targetFamily: 'sql' as const,
+        strategy: 'function' as const,
+        template: 'example({args})',
+      },
+    },
+  ];
+
+  return {
+    kind: 'extension',
+    id: 'test-extension',
+    version: '0.0.1',
+    familyId: 'sql' as const,
+    targetId: 'postgres' as const,
+    codecs: () => registry,
+    queryOperations: () => operations,
+    parameterizedCodecs: () => [],
+    create() {
+      return {
+        familyId: 'sql' as const,
+        targetId: 'postgres' as const,
+      };
+    },
+  };
+}
+
+describe('createExecutionStack', () => {
+  it('defaults driver to undefined and extensions to empty', () => {
+    const stack = createExecutionStack({
+      target: createStubTargetDescriptor(),
+      adapter: createStubAdapterDescriptor(),
+    });
+
+    expect(stack.driver).toBeUndefined();
+    expect(stack.extensionPacks).toEqual([]);
+  });
+
+  it('creates an execution context from descriptors-only stack', () => {
+    const contract = createTestContract({
+      storage: { tables: {} },
+    });
+
+    const context = createExecutionContext({
+      contract,
+      stack: {
+        target: createStubTargetDescriptor(),
+        adapter: createStubAdapterDescriptor(),
+        extensionPacks: [createStubExtensionDescriptor()],
+      },
+    }) as ExecutionContext<typeof contract>;
+
+    expect(context.contract).toBe(contract);
+    expect(context.codecs.get('pg/text@1')).toBeDefined();
+    expect(context.codecs.get('pg/uuid@1')).toBeDefined();
+    expect(context.queryOperations.entries()['example']).toBeDefined();
+    expect(context.types).toEqual({});
+  });
+});
+
+describe('createSqlExecutionStack', () => {
+  it('preserves descriptor references and defaults extensions', () => {
+    const target = createStubTargetDescriptor();
+    const adapter = createStubAdapterDescriptor();
+    const stack = createSqlExecutionStack({ target, adapter });
+
+    expect(stack.target).toBe(target);
+    expect(stack.adapter).toBe(adapter);
+    expect(stack.extensionPacks).toEqual([]);
+  });
+
+  it('keeps extension packs intact', () => {
+    const target = createStubTargetDescriptor();
+    const adapter = createStubAdapterDescriptor();
+    const extension = createStubExtensionDescriptor();
+    const stack = createSqlExecutionStack({ target, adapter, extensionPacks: [extension] });
+
+    expect(stack.extensionPacks).toEqual([extension]);
+  });
+});