@prisma-next/sql-runtime 0.5.0-dev.60 → 0.5.0-dev.62

package/src/codecs/decoding.ts

@@ -7,14 +7,12 @@ import {
 import type {
   AnyQueryAst,
   Codec,
-  CodecRegistry,
   ContractCodecRegistry,
   ProjectionItem,
   SqlCodecCallContext,
 } from '@prisma-next/sql-relational-core/ast';
 import type { SqlExecutionPlan } from '@prisma-next/sql-relational-core/plan';
-import type { JsonSchemaValidatorRegistry } from '@prisma-next/sql-relational-core/query-lane-context';
-import { validateJsonValue } from './json-schema-validation';
+import { makeAliasResolver } from './alias-resolver';
 
 type ColumnRef = { table: string; column: string };
 
@@ -44,40 +42,39 @@ function projectionListFromAst(ast: AnyQueryAst): ReadonlyArray<ProjectionItem>
 /**
  * Resolve the per-cell codec for a projection item.
  *
- * Phase B: when a `(table, column)` ref is available for the projection,
- * prefer `contractCodecs.forColumn(table, column)` — that's the per-
- * instance resolved codec materialized from the codec descriptor's
- * factory at context-construction time (carries any per-instance state
- * such as the compiled JSON-Schema validator). When the projection
- * resolves to a non-`column-ref` expression (computed projections, raw
- * SQL aliases) but still carries a codec id (ADR 205 stamps every
- * `ProjectionItem` with the producer's codec id), fall back to the
- * codec-id-keyed `forCodecId(codecId)` lookup, which itself falls back
- * to the legacy `CodecRegistry` for codec ids the contract walk
- * couldn't resolve.
+ * When a `(table, column)` ref is available — either implicit on a `column-ref` expression or carried explicitly via `item.refs` for column-bound non-`column-ref` projections — prefer `contractCodecs.forColumn(table, column)`: that returns the per-instance codec materialized from the descriptor's factory for that column, encoding any per-instance state (typeParams like vector length, schema validators, etc.).
  *
- * Codec-registry-unification spec § AC-4.
+ * The wrong-instance risk for parameterized codecs is closed off structurally:
+ *
+ * 1. `buildContractCodecRegistry` pre-populates `byCodecId` with one canonical instance per non-parameterized descriptor; parameterized descriptors are intentionally absent. 2. `forCodecId` rejects ambiguous parameterized fallbacks (`ambiguousCodecIds`). 3. The non-ambiguous parameterized case stores the column-correct per-instance codec under `byCodecId`, so the fall-through still resolves to the right instance.
+ *
+ * The `forCodecId` fallback otherwise covers projections that are *not* column-bound (computed projections, raw SQL aliases) but still carry a `codecId` (ADR 205 stamps every `ProjectionItem` with the producer's codec id).
+ *
+ * Codec-registry-unification spec § AC-4 / AC-5.
  */
 function resolveProjectionCodec(
   item: ProjectionItem,
-  registry: CodecRegistry,
   contractCodecs: ContractCodecRegistry | undefined,
+  aliasResolver: (alias: string) => string,
 ): Codec | undefined {
-  if (item.expr.kind === 'column-ref' && contractCodecs) {
-    const byColumn = contractCodecs.forColumn(item.expr.table, item.expr.column);
-    if (byColumn) return byColumn;
+  if (contractCodecs) {
+    if (item.expr.kind === 'column-ref') {
+      const byColumn = contractCodecs.forColumn(aliasResolver(item.expr.table), item.expr.column);
+      // Only honour `byColumn` when its codec id agrees with `item.codecId`. They can legitimately disagree when an `OperationExpr`-shaped projection carries a single inner column-ref but transforms the value's codec (e.g. `cosineDistance(col, x)` projects `pg/float8@1` while the inner column-ref points at a `pg/vector@1` column).
+      if (byColumn && (item.codecId === undefined || byColumn.id === item.codecId)) return byColumn;
+    } else if (item.refs) {
+      const byColumn = contractCodecs.forColumn(aliasResolver(item.refs.table), item.refs.column);
+      if (byColumn && (item.codecId === undefined || byColumn.id === item.codecId)) return byColumn;
+    }
   }
   if (item.codecId) {
-    const fromContract = contractCodecs?.forCodecId(item.codecId);
-    if (fromContract) return fromContract;
-    return registry.get(item.codecId);
+    return contractCodecs?.forCodecId(item.codecId);
   }
   return undefined;
 }
 
 function buildDecodeContext(
   plan: SqlExecutionPlan,
-  registry: CodecRegistry,
   contractCodecs: ContractCodecRegistry | undefined,
 ): DecodeContext {
   if (!isAstBackedPlan(plan)) {
@@ -103,17 +100,26 @@ function buildDecodeContext(
   const codecs = new Map<string, Codec>();
   const columnRefs = new Map<string, ColumnRef>();
   const includeAliases = new Set<string>();
+  const aliasResolver = makeAliasResolver(plan.ast);
 
   for (const item of projection) {
     aliases.push(item.alias);
 
-    const codec = resolveProjectionCodec(item, registry, contractCodecs);
+    const codec = resolveProjectionCodec(item, contractCodecs, aliasResolver);
     if (codec) {
       codecs.set(item.alias, codec);
     }
 
     if (item.expr.kind === 'column-ref') {
-      columnRefs.set(item.alias, { table: item.expr.table, column: item.expr.column });
+      columnRefs.set(item.alias, {
+        table: aliasResolver(item.expr.table),
+        column: item.expr.column,
+      });
+    } else if (item.refs) {
+      columnRefs.set(item.alias, {
+        table: aliasResolver(item.refs.table),
+        column: item.refs.column,
+      });
     } else if (item.expr.kind === 'subquery' || item.expr.kind === 'json-array-agg') {
       includeAliases.add(item.alias);
     }
@@ -131,10 +137,6 @@ function previewWireValue(wireValue: unknown): string {
   return String(wireValue).substring(0, WIRE_PREVIEW_LIMIT);
 }
 
-function isJsonSchemaValidationError(error: unknown): boolean {
-  return isRuntimeError(error) && error.code === 'RUNTIME.JSON_SCHEMA_VALIDATION_FAILED';
-}
-
 function wrapDecodeFailure(
   error: unknown,
   alias: string,
@@ -197,25 +199,16 @@ function decodeIncludeAggregate(alias: string, wireValue: unknown): unknown {
 }
 
 /**
- * Decodes a single field. Single-armed: every cell takes the same path —
- * `codec.decode → await → JSON-Schema validate → return plain value` — so
- * sync- and async-authored codecs are indistinguishable to callers.
+ * Decodes a single field. Single-armed: every cell takes the same path — `codec.decode → await → return plain value` — so sync- and async-authored codecs are indistinguishable to callers. JSON-Schema validation, when required, lives inside the resolved codec's `decode` body (e.g. `arktype-json` validates against its rehydrated schema and throws `RUNTIME.JSON_SCHEMA_VALIDATION_FAILED` from `decode` directly); there is
+ * no separate validator-registry pass.
  *
- * The row-level `rowCtx` is repackaged into a per-cell
- * `SqlCodecCallContext` whose `column = { table, name }` is a structural
- * projection of the per-cell `ColumnRef = { table, column }` resolved from
- * the AST-backed `DecodeContext` (the same resolution `wrapDecodeFailure`
- * uses for envelope construction — one resolution per cell, two consumers).
- * Cells the runtime cannot resolve to a single underlying column (aggregate
- * aliases, computed projections without a simple ref) get
- * `column: undefined`, matching the spec contract that the runtime never
- * silently defaults this field.
+ * The row-level `rowCtx` is repackaged into a per-cell `SqlCodecCallContext` whose `column = { table, name }` is a structural projection of the per-cell `ColumnRef = { table, column }` resolved from the AST-backed `DecodeContext` (the same resolution `wrapDecodeFailure` uses for envelope construction — one resolution per cell, two consumers). Cells the runtime cannot resolve to a single underlying column (aggregate
+ * aliases, computed projections without a simple ref) get `column: undefined`, matching the spec contract that the runtime never silently defaults this field.
  */
 async function decodeField(
   alias: string,
   wireValue: unknown,
   decodeCtx: DecodeContext,
-  jsonValidators: JsonSchemaValidatorRegistry | undefined,
   rowCtx: SqlCodecCallContext,
 ): Promise<unknown> {
   if (wireValue === null) {
@@ -229,15 +222,8 @@ async function decodeField(
 
   const ref = decodeCtx.columnRefs.get(alias);
 
-  // Per-cell ctx: the cell-level `column` is a `SqlColumnRef = { table, name }`
-  // projection of the resolved `ColumnRef = { table, column }` (same
-  // resolution `wrapDecodeFailure` uses below — no double work). Cells the
-  // runtime cannot resolve (aggregate aliases, computed projections without
-  // a simple ref) drop the `column` field entirely — explicitly cleared so
-  // a previously-populated `rowCtx.column` cannot leak through to unrelated
-  // cells. Destructuring (rather than `column: undefined`) is required
-  // because `SqlCodecCallContext.column` is declared `column?: SqlColumnRef`
-  // under `exactOptionalPropertyTypes`.
+  // Per-cell ctx: the cell-level `column` is a `SqlColumnRef = { table, name }` projection of the resolved `ColumnRef = { table, column }` (same resolution `wrapDecodeFailure` uses below — no double work). Cells the runtime cannot resolve (aggregate aliases, computed projections without a simple ref) drop the `column` field entirely — explicitly cleared so a previously-populated `rowCtx.column` cannot leak through to
+  // unrelated cells. Destructuring (rather than `column: undefined`) is required because `SqlCodecCallContext.column` is declared `column?: SqlColumnRef` under `exactOptionalPropertyTypes`.
   let cellCtx: SqlCodecCallContext;
   if (ref) {
     cellCtx = { ...rowCtx, column: { table: ref.table, name: ref.column } };
@@ -246,55 +232,36 @@ async function decodeField(
     cellCtx = rowCtxWithoutColumn;
   }
 
-  let decoded: unknown;
   try {
-    decoded = await codec.decode(wireValue, cellCtx);
+    return await codec.decode(wireValue, cellCtx);
   } catch (error) {
-    wrapDecodeFailure(error, alias, ref, codec, wireValue);
-  }
-
-  if (jsonValidators && ref) {
-    try {
-      validateJsonValue(jsonValidators, ref.table, ref.column, decoded, 'decode', codec.id);
-    } catch (error) {
-      if (isJsonSchemaValidationError(error)) throw error;
-      wrapDecodeFailure(error, alias, ref, codec, wireValue);
+    // Codec-authored runtime envelopes (e.g. `RUNTIME.DECODE_FAILED` thrown from inside the codec body, or `RUNTIME.ABORTED` raised via `CodecCallContext.signal` per ADR 207) carry their own per-codec context — wrapping them again would erase that context and coerce the abort intent into a generic decode failure. Pass them through unchanged; only foreign errors get the `wrapDecodeFailure` envelope.
+    if (isRuntimeError(error)) {
+      throw error;
     }
+    wrapDecodeFailure(error, alias, ref, codec, wireValue);
   }
-
-  return decoded;
 }
 
 /**
- * Decodes a row by dispatching all per-cell codec calls concurrently via
- * `Promise.all`. Each cell follows the single-armed `decodeField` path.
- * Failures are wrapped in `RUNTIME.DECODE_FAILED` with `{ table, column,
- * codec }` (or `{ alias, codec }` when no column ref is resolvable) and the
- * original error attached on `cause`.
+ * Decodes a row by dispatching all per-cell codec calls concurrently via `Promise.all`. Each cell follows the single-armed `decodeField` path. Failures are wrapped in `RUNTIME.DECODE_FAILED` with `{ table, column, codec }` (or `{ alias, codec }` when no column ref is resolvable) and the original error attached on `cause`.
  *
  * When `rowCtx.signal` is provided:
  *
- * - **Already-aborted at entry** short-circuits with `RUNTIME.ABORTED`
- *   (`{ phase: 'decode' }`) before any `codec.decode` call is made.
- * - **Mid-flight aborts** race the per-cell `Promise.all` against the
- *   signal so the runtime returns promptly even when codec bodies ignore
- *   it. In-flight bodies that ignore the signal complete in the
- *   background (cooperative cancellation).
- * - Existing `RUNTIME.DECODE_FAILED` envelopes from codec bodies pass
- *   through unchanged (no double wrap).
+ * - **Already-aborted at entry** short-circuits with `RUNTIME.ABORTED` (`{ phase: 'decode' }`) before any `codec.decode` call is made.
+ * - **Mid-flight aborts** race the per-cell `Promise.all` against the signal so the runtime returns promptly even when codec bodies ignore it. In-flight bodies that ignore the signal complete in the background (cooperative cancellation).
+ * - Existing `RUNTIME.DECODE_FAILED` envelopes from codec bodies pass through unchanged (no double wrap).
  */
 export async function decodeRow(
   row: Record<string, unknown>,
   plan: SqlExecutionPlan,
-  registry: CodecRegistry,
-  jsonValidators: JsonSchemaValidatorRegistry | undefined,
   rowCtx: SqlCodecCallContext,
   contractCodecs?: ContractCodecRegistry,
 ): Promise<Record<string, unknown>> {
   checkAborted(rowCtx, 'decode');
   const signal = rowCtx.signal;
 
-  const decodeCtx = buildDecodeContext(plan, registry, contractCodecs);
+  const decodeCtx = buildDecodeContext(plan, contractCodecs);
 
   const aliases = decodeCtx.aliases ?? Object.keys(row);
 
@@ -323,13 +290,12 @@ export async function decodeRow(
       continue;
     }
 
-    tasks.push(decodeField(alias, wireValue, decodeCtx, jsonValidators, rowCtx));
+    tasks.push(decodeField(alias, wireValue, decodeCtx, rowCtx));
   }
 
   const settled = await raceAgainstAbort(Promise.all(tasks), signal, 'decode');
 
-  // Include aggregates are decoded synchronously after concurrent codec
-  // dispatch settles, so any decode failures upstream propagate first.
+  // Include aggregates are decoded synchronously after concurrent codec dispatch settles, so any decode failures upstream propagate first.
   for (const entry of includeIndices) {
     settled[entry.index] = decodeIncludeAggregate(entry.alias, entry.value);
   }

package/src/codecs/encoding.ts

@@ -5,51 +5,54 @@ import {
 } from '@prisma-next/framework-components/runtime';
 import {
   type Codec,
-  type CodecRegistry,
   type ContractCodecRegistry,
   collectOrderedParamRefs,
+  type ParamRefBindingRefs,
   type SqlCodecCallContext,
 } from '@prisma-next/sql-relational-core/ast';
 import type { SqlExecutionPlan } from '@prisma-next/sql-relational-core/plan';
+import { makeAliasResolver } from './alias-resolver';
 
 interface ParamMetadata {
   readonly codecId: string | undefined;
   readonly name: string | undefined;
+  readonly refs: ParamRefBindingRefs | undefined;
 }
 
-const NO_METADATA: ParamMetadata = Object.freeze({ codecId: undefined, name: undefined });
+const NO_METADATA: ParamMetadata = Object.freeze({
+  codecId: undefined,
+  name: undefined,
+  refs: undefined,
+});
 
 /**
  * Resolve the codec for an outgoing param.
  *
- * Phase B (and AC-5-deferred carve-out): `ParamRef` does not carry a
- * `(table, column)` ref today — every `ParamRef` carries `codecId` but
- * not the column it relates to. Encode-side dispatch therefore consults
- * `contractCodecs.forCodecId(codecId)` (which itself prefers the
- * contract-walk-derived shared codec, falling back to the legacy
- * `CodecRegistry.get` for parameterized codec ids whose contracts don't
- * have a column the walk could resolve through).
+ * Column-aware dispatch: when `metadata.refs` is populated by a column-bound construction site, prefer `contractCodecs.forColumn(refs.table, refs.column)` — that returns the per-instance codec the contract walk materialized for the `(table, column)` pair, encoding the column's typeParams (e.g. `vector(1024)` vs. `vector(1536)`).
  *
- * For the parameterized codecs shipped at Phase B (pgvector, postgres
- * json/jsonb), encode is per-instance-stateless w.r.t. params:
- * - pgvector formats `[v1,v2,...]` regardless of declared length;
- * - postgres json/jsonb encode is `JSON.stringify` regardless of schema.
+ * On a column-lookup miss the resolver falls through to `forCodecId`. The wrong-instance risk for parameterized codecs is closed off structurally:
  *
- * So the codec-id-keyed lookup yields a structurally equivalent encoder
- * even when the resolved per-instance codec carries extra state (e.g. a
- * compiled JSON-Schema validator used only by `decode`). TML-2357 retires
- * the fallback by threading `ParamRef.refs` through column-bound
- * construction sites.
+ * 1. `buildContractCodecRegistry` pre-populates `byCodecId` with one canonical instance per non-parameterized descriptor; parameterized descriptors are intentionally absent from this pre-population. 2. `forCodecId` rejects ambiguous parameterized fallbacks (`ambiguousCodecIds`) — if the contract walk resolved more than one distinct instance under a single parameterized id, the call throws rather than binding to
+ * whichever landed first. 3. For the non-ambiguous parameterized case (a single column with that id), `byCodecId` stores the column-correct per-instance codec, so the fall-through still resolves to the right instance.
+ *
+ * Refs-less fallback: ParamRefs constructed outside a column-bound site (literals, transient builder state) carry a non-parameterized `codecId` whose dispatch is ambiguity-free. The validator pass (`validateParamRefRefs`) already enforced refs on every parameterized ParamRef before encode runs.
  */
 function resolveParamCodec(
   metadata: ParamMetadata,
-  registry: CodecRegistry,
   contractCodecs: ContractCodecRegistry | undefined,
+  aliasResolver: (alias: string) => string,
 ): Codec | undefined {
   if (!metadata.codecId) return undefined;
-  const fromContract = contractCodecs?.forCodecId(metadata.codecId);
-  if (fromContract) return fromContract;
-  return registry.get(metadata.codecId);
+  if (metadata.refs && contractCodecs) {
+    const byColumn = contractCodecs.forColumn(
+      aliasResolver(metadata.refs.table),
+      metadata.refs.column,
+    );
+    // Only honour `byColumn` when its codec id agrees with the `ParamRef`'s declared `codecId`. They can legitimately disagree when a heuristic (e.g. the ORM's `refsFromLeft`) lifts column refs out of an `OperationExpr` that changed the codec id — e.g. `cosineDistance(p.embedding, x).lt(1)` carries `refs={post,embedding}` (a vector column) but the comparison side's codec is `pg/float8@1`. Trusting `byColumn` blindly would dispatch the float
+    // literal through the vector codec.
+    if (byColumn && byColumn.id === metadata.codecId) return byColumn;
+  }
+  return contractCodecs?.forCodecId(metadata.codecId);
 }
 
 function paramLabel(metadata: ParamMetadata, paramIndex: number): string {
@@ -74,34 +77,28 @@ function wrapEncodeFailure(
 }
 
 /**
- * Encodes a single parameter through its codec. Always awaits codec.encode so
- * a Promise can never leak into the driver, even if a sync-authored codec is
- * lifted to async by the codec() factory. Failures are wrapped in
- * `RUNTIME.ENCODE_FAILED` with `{ label, codec, paramIndex }` and the original
- * error attached on `cause`.
+ * Encodes a single parameter through its codec. Always awaits codec.encode so a Promise can never leak into the driver, even if a sync-authored codec is lifted to async by the codec factory. Failures are wrapped in `RUNTIME.ENCODE_FAILED` with `{ label, codec, paramIndex }` and the original error attached on `cause`.
  *
- * `ctx` is forwarded verbatim to `codec.encode` so codec authors who opt
- * into the `(value, ctx)` arity see the same `SqlCodecCallContext` the
- * runtime built for the surrounding `runtime.execute()` call. The ctx is
- * always present; its `signal` field may be `undefined`. Encode call
- * sites do not populate `ctx.column` — encode-time column context is the
- * middleware's domain.
+ * `ctx` is forwarded verbatim to `codec.encode` so codec authors who opt into the `(value, ctx)` arity see the same `SqlCodecCallContext` the runtime built for the surrounding `runtime.execute()` call. The ctx is always present; its `signal` field may be `undefined`. Encode call sites do not populate `ctx.column` — encode-time column context is the middleware's domain.
  */
 export async function encodeParam(
   value: unknown,
-  paramRef: { readonly codecId?: string; readonly name?: string },
+  paramRef: {
+    readonly codecId?: string;
+    readonly name?: string;
+    readonly refs?: ParamRefBindingRefs;
+  },
   paramIndex: number,
-  registry: CodecRegistry,
   ctx: SqlCodecCallContext,
   contractCodecs?: ContractCodecRegistry,
 ): Promise<unknown> {
   return encodeParamValue(
     value,
-    { codecId: paramRef.codecId, name: paramRef.name },
+    { codecId: paramRef.codecId, name: paramRef.name, refs: paramRef.refs },
     paramIndex,
-    registry,
     ctx,
     contractCodecs,
+    (alias) => alias,
   );
 }
 
@@ -109,15 +106,15 @@ async function encodeParamValue(
   value: unknown,
   metadata: ParamMetadata,
   paramIndex: number,
-  registry: CodecRegistry,
   ctx: SqlCodecCallContext,
   contractCodecs: ContractCodecRegistry | undefined,
+  aliasResolver: (alias: string) => string,
 ): Promise<unknown> {
   if (value === null || value === undefined) {
     return null;
   }
 
-  const codec = resolveParamCodec(metadata, registry, contractCodecs);
+  const codec = resolveParamCodec(metadata, contractCodecs, aliasResolver);
   if (!codec) {
     return value;
   }
@@ -130,27 +127,16 @@ async function encodeParamValue(
 }
 
 /**
- * Encodes all parameters concurrently via `Promise.all`. Per parameter, sync-
- * and async-authored codecs share the same path: `codec.encode → await →
- * return`. Param-level failures are wrapped in `RUNTIME.ENCODE_FAILED`.
+ * Encodes all parameters concurrently via `Promise.all`. Per parameter, sync-and async-authored codecs share the same path: `codec.encode → await → return`. Param-level failures are wrapped in `RUNTIME.ENCODE_FAILED`.
  *
  * When `ctx.signal` is provided:
  *
- * - **Already-aborted at entry** short-circuits with `RUNTIME.ABORTED`
- *   (`{ phase: 'encode' }`) before any `codec.encode` call is made — codecs
- *   can pin this with a per-call counter that stays at zero.
- * - **Mid-flight abort** races the per-param `Promise.all` against
- *   `abortable(ctx.signal)`. The runtime returns `RUNTIME.ABORTED` promptly
- *   even if codec bodies ignore the signal; the in-flight bodies are
- *   abandoned and run to completion in the background (cooperative
- *   cancellation, see ADR 204).
- * - Existing `RUNTIME.ENCODE_FAILED` envelopes that surface from a codec
- *   body before the runtime observes the abort pass through unchanged
- *   (no double wrap).
+ * - **Already-aborted at entry** short-circuits with `RUNTIME.ABORTED` (`{ phase: 'encode' }`) before any `codec.encode` call is made — codecs can pin this with a per-call counter that stays at zero.
+ * - **Mid-flight abort** races the per-param `Promise.all` against `abortable(ctx.signal)`. The runtime returns `RUNTIME.ABORTED` promptly even if codec bodies ignore the signal; the in-flight bodies are abandoned and run to completion in the background (cooperative cancellation, see ADR 204).
+ * - Existing `RUNTIME.ENCODE_FAILED` envelopes that surface from a codec body before the runtime observes the abort pass through unchanged (no double wrap).
  */
 export async function encodeParams(
   plan: SqlExecutionPlan,
-  registry: CodecRegistry,
   ctx: SqlCodecCallContext,
   contractCodecs?: ContractCodecRegistry,
 ): Promise<readonly unknown[]> {
@@ -169,20 +155,22 @@ export async function encodeParams(
     for (let i = 0; i < paramCount && i < refs.length; i++) {
       const ref = refs[i];
       if (ref) {
-        metadata[i] = { codecId: ref.codecId, name: ref.name };
+        metadata[i] = { codecId: ref.codecId, name: ref.name, refs: ref.refs };
       }
     }
   }
 
+  const aliasResolver = makeAliasResolver(plan.ast);
+
   const tasks: Promise<unknown>[] = new Array(paramCount);
   for (let i = 0; i < paramCount; i++) {
     tasks[i] = encodeParamValue(
       plan.params[i],
       metadata[i] ?? NO_METADATA,
       i,
-      registry,
       ctx,
       contractCodecs,
+      aliasResolver,
     );
   }
 

package/src/codecs/validation.ts

@@ -1,7 +1,6 @@
 import type { Contract } from '@prisma-next/contract/types';
 import { runtimeError } from '@prisma-next/framework-components/runtime';
 import type { SqlStorage } from '@prisma-next/sql-contract/types';
-import type { CodecRegistry } from '@prisma-next/sql-relational-core/ast';
 import type { CodecDescriptorRegistry } from '@prisma-next/sql-relational-core/query-lane-context';
 
 export function extractCodecIds(contract: Contract<SqlStorage>): Set<string> {
@@ -31,33 +30,15 @@ function extractCodecIdsFromColumns(contract: Contract<SqlStorage>): Map<string,
   return codecIds;
 }
 
-interface CodecLookupForValidation {
-  has(id: string): boolean;
-}
-
-function adaptDescriptorRegistry(registry: CodecDescriptorRegistry): CodecLookupForValidation {
-  return { has: (id: string) => registry.descriptorFor(id) !== undefined };
-}
-
-function isDescriptorRegistry(
-  registry: CodecRegistry | CodecDescriptorRegistry,
-): registry is CodecDescriptorRegistry {
-  return 'descriptorFor' in registry;
-}
-
 export function validateContractCodecMappings(
-  registry: CodecRegistry | CodecDescriptorRegistry,
+  registry: CodecDescriptorRegistry,
   contract: Contract<SqlStorage>,
 ): void {
-  const lookup: CodecLookupForValidation = isDescriptorRegistry(registry)
-    ? adaptDescriptorRegistry(registry)
-    : registry;
-
   const codecIds = extractCodecIdsFromColumns(contract);
   const invalidCodecs: Array<{ table: string; column: string; codecId: string }> = [];
 
   for (const [key, codecId] of codecIds.entries()) {
-    if (!lookup.has(codecId)) {
+    if (registry.descriptorFor(codecId) === undefined) {
       const parts = key.split('.');
       const table = parts[0] ?? '';
       const column = parts[1] ?? '';
@@ -80,7 +61,7 @@ export function validateContractCodecMappings(
 }
 
 export function validateCodecRegistryCompleteness(
-  registry: CodecRegistry | CodecDescriptorRegistry,
+  registry: CodecDescriptorRegistry,
   contract: Contract<SqlStorage>,
 ): void {
   validateContractCodecMappings(registry, contract);

package/src/middleware/budgets.ts

@@ -25,14 +25,17 @@ function hasAggregateWithoutGroupBy(ast: SelectAst): boolean {
   return ast.projection.some((item) => item.expr.kind === 'aggregate');
 }
 
-function primaryTableFromAst(ast: SelectAst): string | undefined {
+function primaryTableFromAst(ast: SelectAst): string {
   switch (ast.from.kind) {
     case 'table-source':
       return ast.from.name;
     case 'derived-table-source':
       return ast.from.alias;
+    // v8 ignore next 4
     default:
-      return undefined;
+      throw new Error(
+        `Unsupported source kind: ${(ast.from satisfies never as { kind: string }).kind}`,
+      );
   }
 }
 
@@ -41,17 +44,12 @@ function estimateRowsFromAst(
   tableRows: Record<string, number>,
   defaultTableRows: number,
   hasAggregateWithoutGroup: boolean,
-): number | null {
+): number {
   if (hasAggregateWithoutGroup) {
     return 1;
   }
 
-  const table = primaryTableFromAst(ast);
-  if (!table) {
-    return null;
-  }
-
-  const tableEstimate = tableRows[table] ?? defaultTableRows;
+  const tableEstimate = tableRows[primaryTableFromAst(ast)] ?? defaultTableRows;
 
   if (ast.limit !== undefined) {
     return Math.min(ast.limit, tableEstimate);
@@ -136,31 +134,19 @@ export function budgets(options?: BudgetsOptions): SqlMiddleware {
     const shouldBlock = rowSeverity === 'error' || ctx.mode === 'strict';
 
     if (isUnbounded) {
-      if (estimated !== null && estimated >= maxRows) {
-        emitBudgetViolation(
-          runtimeError('BUDGET.ROWS_EXCEEDED', 'Unbounded SELECT query exceeds budget', {
-            source: 'ast',
-            estimatedRows: estimated,
-            maxRows,
-          }),
-          shouldBlock,
-          ctx,
-        );
-        return;
-      }
-
+      const details =
+        estimated >= maxRows
+          ? { source: 'ast', estimatedRows: estimated, maxRows }
+          : { source: 'ast', maxRows };
       emitBudgetViolation(
-        runtimeError('BUDGET.ROWS_EXCEEDED', 'Unbounded SELECT query exceeds budget', {
-          source: 'ast',
-          maxRows,
-        }),
+        runtimeError('BUDGET.ROWS_EXCEEDED', 'Unbounded SELECT query exceeds budget', details),
         shouldBlock,
         ctx,
       );
       return;
     }
 
-    if (estimated !== null && estimated > maxRows) {
+    if (estimated > maxRows) {
       emitBudgetViolation(
         runtimeError('BUDGET.ROWS_EXCEEDED', 'Estimated row count exceeds budget', {
           source: 'ast',