collection-dispatch.test.ts

import { AsyncIterableResult } from '@prisma-next/framework-components/runtime';
import type { RuntimeScope } from '@prisma-next/sql-relational-core/types';
import { describe, expect, it } from 'vitest';
import { dispatchCollectionRows, stitchIncludes } from '../src/collection-dispatch';
import type { IncludeExpr } from '../src/types';
import { emptyState } from '../src/types';
import { createCollectionFor } from './collection-fixtures';
import type { MockRuntime, TestContract } from './helpers';
import { createMockRuntime, getTestContract, withCapabilities } from './helpers';

function withSingleQueryCapabilities(contract: TestContract) {
  return withCapabilities(contract, {
    ...contract.capabilities,
    [contract.targetFamily]: {
      ...(contract.capabilities[contract.targetFamily] ?? {}),
      jsonAgg: true,
    },
    [contract.target]: {
      ...(contract.capabilities[contract.target] ?? {}),
      jsonAgg: true,
      lateral: true,
    },
  });
}

/**
 * Mirrors the shape produced by the contract emitter: capability flags
 * nested under the family + target namespaces, with no top-level entries.
 * Used to assert "single-query path is selected for an emitted-shape
 * contract" — the regression scenario the principled namespaced lookup
 * was introduced to handle.
 */
function withEmittedSqlCapabilities(contract: TestContract) {
  return withCapabilities(contract, {
    sql: { jsonAgg: true, returning: true },
    postgres: { jsonAgg: true, lateral: true, returning: true },
  });
}

function addConnection(
  runtime: MockRuntime,
  onRelease: () => void,
): MockRuntime & {
  connection: () => Promise<{
    execute: MockRuntime['execute'];
    release: () => Promise<void>;
  }>;
} {
  return Object.assign(runtime, {
    async connection() {
      return {
        execute: runtime.execute.bind(runtime),
        async release() {
          onRelease();
        },
      };
    },
  });
}

function cloneInclude(include: IncludeExpr, overrides: Partial<IncludeExpr>): IncludeExpr {
  return {
    ...include,
    ...overrides,
  };
}

function emptyScope(): RuntimeScope {
  return {
    execute() {
      return new AsyncIterableResult((async function* () {})());
    },
  };
}

describe('collection-dispatch', () => {
  it('dispatchCollectionRows() maps rows when includes are absent', async () => {
    const { collection, runtime } = createCollectionFor('User');
    runtime.setNextResults([[{ id: 1, name: 'Alice', email: 'alice@example.com' }]]);

    const rows = await dispatchCollectionRows<Record<string, unknown>>({
      contract: collection.ctx.context.contract,
      runtime,
      state: collection.state,
      tableName: collection.tableName,
      modelName: collection.modelName,
    }).toArray();

    expect(rows).toEqual([{ id: 1, name: 'Alice', email: 'alice@example.com' }]);
  });

  it('dispatchCollectionRows() depth-1 include with emitted-shape capabilities fires a single SQL execution (regression guard for namespaced capability lookup)', async () => {
    // Guards against regressing the fix that taught `selectIncludeStrategy`
    // to read capability flags from the contract's `targetFamily` and
    // `target` namespaces. Prior to that fix, every emitted contract fell
    // back to multi-query for nested includes — silently, because
    // functional correctness was unaffected. This test fails fast if the
    // regression returns: an emitted-shape contract should resolve a
    // depth-1 include in one SQL execution, not two.
    const contract = withEmittedSqlCapabilities(getTestContract());
    const { collection, runtime } = createCollectionFor('User', contract);
    const scoped = collection.select('name').include('posts');
    runtime.setNextResults([
      [{ id: 1, name: 'Alice', posts: '[{"id":10,"title":"Post A","user_id":1,"views":3}]' }],
    ]);

    const rows = await dispatchCollectionRows<Record<string, unknown>>({
      contract,
      runtime,
      state: scoped.state,
      tableName: scoped.tableName,
      modelName: scoped.modelName,
    }).toArray();

    expect(rows).toEqual([
      { name: 'Alice', posts: [{ id: 10, title: 'Post A', userId: 1, views: 3 }] },
    ]);
    // The point of the test: 1 execution, not N+1.
    expect(runtime.executions).toHaveLength(1);
  });

  it('dispatchCollectionRows() single-query path returns empty rows and releases scope', async () => {
    const contract = withSingleQueryCapabilities(getTestContract());
    const { collection, runtime } = createCollectionFor('User', contract);
    const scoped = collection.include('posts');
    runtime.setNextResults([[]]);

    let released = false;
    const runtimeWithConnection = addConnection(runtime, () => {
      released = true;
    });

    const rows = await dispatchCollectionRows<Record<string, unknown>>({
      contract,
      runtime: runtimeWithConnection,
      state: scoped.state,
      tableName: scoped.tableName,
      modelName: scoped.modelName,
    }).toArray();

    expect(rows).toEqual([]);
    expect(released).toBe(true);
  });

  it('dispatchCollectionRows() single-query path parses include payloads and strips hidden join columns', async () => {
    const contract = withSingleQueryCapabilities(getTestContract());
    const { collection, runtime } = createCollectionFor('User', contract);
    const scoped = collection.select('name').include('posts');
    runtime.setNextResults([
      [
        {
          id: 1,
          name: 'Alice',
          posts: '[{"id":10,"title":"Post A","user_id":1,"views":3},42,null]',
        },
        {
          id: 2,
          name: 'Bob',
          posts: 'not-json',
        },
        {
          id: 3,
          name: 'Cara',
          posts: null,
        },
        {
          id: 4,
          name: 'Drew',
          posts: '{"id":99}',
        },
      ],
    ]);

    const rows = await dispatchCollectionRows<Record<string, unknown>>({
      contract,
      runtime,
      state: scoped.state,
      tableName: scoped.tableName,
      modelName: scoped.modelName,
    }).toArray();

    expect(rows).toEqual([
      {
        name: 'Alice',
        posts: [{ id: 10, title: 'Post A', userId: 1, views: 3 }],
      },
      {
        name: 'Bob',
        posts: [],
      },
      {
        name: 'Cara',
        posts: [],
      },
      {
        name: 'Drew',
        posts: [],
      },
    ]);
  });

  it('dispatchCollectionRows() single-query to-one include returns mapped row or null', async () => {
    const contract = withSingleQueryCapabilities(getTestContract());
    const { collection, runtime } = createCollectionFor('Post', contract);
    const scoped = collection.select('title').include('author');
    runtime.setNextResults([
      [
        {
          user_id: 1,
          title: 'Has Author',
          author: '[{"id":1,"name":"Alice","email":"alice@example.com"}]',
        },
        {
          user_id: null,
          title: 'No Author',
          author: '[]',
        },
      ],
    ]);

    const rows = await dispatchCollectionRows<Record<string, unknown>>({
      contract,
      runtime,
      state: scoped.state,
      tableName: scoped.tableName,
      modelName: scoped.modelName,
    }).toArray();

    expect(rows).toEqual([
      {
        title: 'Has Author',
        author: {
          id: 1,
          name: 'Alice',
          email: 'alice@example.com',
        },
      },
      {
        title: 'No Author',
        author: null,
      },
    ]);
  });

  it('dispatchCollectionRows() multi-query path stitches includes, strips hidden fields, and releases scope', async () => {
    // Force multi-query strategy by clearing capabilities. Otherwise
    // the base test contract's postgres.lateral / postgres.jsonAgg
    // would route to single-query lateral.
    const contract = withCapabilities(getTestContract(), {});
    const { collection, runtime } = createCollectionFor('User', contract);
    const scoped = collection.select('name').include('posts', (posts) => posts.select('title'));

    runtime.setNextResults([
      [
        { id: 1, name: 'Alice' },
        { id: 2, name: 'Bob' },
      ],
      [
        { user_id: 1, title: 'One' },
        { user_id: 1, title: 'Two' },
      ],
    ]);

    let released = false;
    const runtimeWithConnection = addConnection(runtime, () => {
      released = true;
    });

    const rows = await dispatchCollectionRows<Record<string, unknown>>({
      contract,
      runtime: runtimeWithConnection,
      state: scoped.state,
      tableName: scoped.tableName,
      modelName: scoped.modelName,
    }).toArray();

    expect(rows).toEqual([
      {
        name: 'Alice',
        posts: [{ title: 'One' }, { title: 'Two' }],
      },
      {
        name: 'Bob',
        posts: [],
      },
    ]);
    expect(released).toBe(true);
  });

  it('dispatchCollectionRows() multi-query path handles empty parent result sets', async () => {
    // Force multi-query strategy so the empty-parent early return inside
    // `dispatchWithMultiQueryIncludes` is actually exercised. The base
    // contract's postgres.lateral / postgres.jsonAgg would otherwise
    // route to single-query lateral.
    const contract = withCapabilities(getTestContract(), {});
    const { collection, runtime } = createCollectionFor('User', contract);
    const scoped = collection.include('posts');

    runtime.setNextResults([[]]);

    const rows = await dispatchCollectionRows<Record<string, unknown>>({
      contract,
      runtime,
      state: scoped.state,
      tableName: scoped.tableName,
      modelName: scoped.modelName,
    }).toArray();

    expect(rows).toEqual([]);
  });

  it('stitchIncludes() assigns empty values for row, scalar, and combine descriptors', async () => {
    const contract = getTestContract();
    const { collection } = createCollectionFor('User', contract);
    const rowInclude = collection.include('posts').state.includes[0]!;
    const scalarInclude = collection.include('posts', (posts) => posts.sum('views' as never)).state
      .includes[0]!;
    const combineInclude = collection.include('posts', (posts) =>
      posts.combine({
        rows: posts.take(1),
        total: posts.sum('views' as never),
      }),
    ).state.includes[0]!;

    const parentRows = [
      { raw: {}, mapped: {} as Record<string, unknown> },
      { raw: {}, mapped: {} as Record<string, unknown> },
    ];

    await stitchIncludes(emptyScope(), contract, parentRows, [
      cloneInclude(rowInclude, { relationName: 'rowBranch' }),
      cloneInclude(scalarInclude, { relationName: 'scalarBranch' }),
      cloneInclude(combineInclude, { relationName: 'combineBranch' }),
    ]);

    expect(parentRows).toEqual([
      {
        raw: {},
        mapped: {
          rowBranch: [],
          scalarBranch: null,
          combineBranch: {
            rows: [],
            total: null,
          },
        },
      },
      {
        raw: {},
        mapped: {
          rowBranch: [],
          scalarBranch: null,
          combineBranch: {
            rows: [],
            total: null,
          },
        },
      },
    ]);
  });

  it('stitchIncludes() computes scalar aggregates with numeric coercion and unknown selectors', async () => {
    const contract = getTestContract();
    const runtime = createMockRuntime();
    const baseInclude = createCollectionFor('User', contract).collection.include('posts').state
      .includes[0]!;

    const sumSelector = {
      kind: 'includeScalar',
      fn: 'sum',
      column: 'views',
      state: emptyState(),
    } as IncludeExpr['scalar'];
    const noColumnSelector = {
      kind: 'includeScalar',
      fn: 'sum',
      state: emptyState(),
    } as IncludeExpr['scalar'];
    const unknownSelector = {
      kind: 'includeScalar',
      fn: 'median' as never,
      column: 'views',
      state: emptyState(),
    } as unknown as IncludeExpr['scalar'];

    runtime.setNextResults([
      [
        { user_id: 1, views: 3 },
        { user_id: 1, views: 10n },
        { user_id: 1, views: '20' },
        { user_id: 1, views: 'bad' },
        { user_id: 1, views: null },
        { user_id: 1, views: {} },
        { user_id: 2, views: 'bad' },
      ],
      [{ user_id: 1, views: 99 }],
      [{ user_id: 1, views: 5 }],
    ]);

    const parentRows = [
      { raw: { id: 1 }, mapped: {} as Record<string, unknown> },
      { raw: { id: 2 }, mapped: {} as Record<string, unknown> },
    ];

    await stitchIncludes(runtime, contract, parentRows, [
      cloneInclude(baseInclude, {
        relationName: 'sumViews',
        scalar: sumSelector,
      }),
      cloneInclude(baseInclude, {
        relationName: 'noColumn',
        scalar: noColumnSelector,
      }),
      cloneInclude(baseInclude, {
        relationName: 'unknownFn',
        scalar: unknownSelector,
      }),
    ]);

    expect(parentRows).toEqual([
      {
        raw: { id: 1 },
        mapped: {
          sumViews: 33,
          noColumn: null,
          unknownFn: null,
        },
      },
      {
        raw: { id: 2 },
        mapped: {
          sumViews: null,
          noColumn: null,
          unknownFn: null,
        },
      },
    ]);
  });

  it('stitchIncludes() returns null for empty to-one row includes', async () => {
    const contract = getTestContract();
    const include = createCollectionFor('Post', contract).collection.include('author').state
      .includes[0]!;

    const parentRows = [{ raw: {}, mapped: {} as Record<string, unknown> }];

    await stitchIncludes(emptyScope(), contract, parentRows, [include]);

    expect(parentRows[0]?.mapped['author']).toBeNull();
  });

  // ---------------------------------------------------------------------------
  // Single-query include child-row codec decoding — DEFERRED follow-up.
  //
  // The three `it.skip` blocks below are placeholders for the case where the
  // single-query include strategy (lateral / correlated jsonb_agg payload)
  // routes embedded child rows through the codec registry and surfaces
  // decoded values (or wrapped failures) on each child cell. The titles
  // describe what each case would assert under the single-path always-await
  // runtime; the bodies are stubbed and not carried over verbatim from any
  // historical implementation.
  //
  // The deferral is structural: the current `dispatchCollectionRows`
  // single-query path (packages/3-extensions/sql-orm-client/src/
  // collection-dispatch.ts) only JSON.parses the include payload and
  // applies field-name mapping; it does not invoke codec query-time methods
  // on child cells (`rg 'codec\.(encode|decode)' packages/3-extensions/
  // sql-orm-client/src` returns zero matches). Adding child-row codec
  // decoding to the single-query include path is a separate piece of ORM
  // work, orthogonal to the codec async-shape decision tracked in ADR 204.
  // ---------------------------------------------------------------------------
  it.skip('dispatchCollectionRows() single-query include decodes async child fields and validates decoded values', async () => {
    // Activates when child-row codec decoding is added to the single-query
    // include path; assertions will express the single-path always-await
    // contract (plain decoded values, no Promises).
  });

  it.skip('dispatchCollectionRows() single-query include preserves JSON schema validation failures for async child decodes', async () => {
    // Activates with the orm-include-aggregate-codec-dispatch follow-up;
    // will assert that JSON schema validation failures on async child cells
    // are reported via the runtime envelope (RUNTIME.VALIDATION_FAILED).
  });

  it.skip('dispatchCollectionRows() single-query include wraps async child decode failures with codec context', async () => {
    // Activates with the orm-include-aggregate-codec-dispatch follow-up;
    // will assert that decode rejections on child cells are wrapped with
    // codec id + lane context (RUNTIME.DECODE_FAILED).
  });
});