query-plan-select.ts

import type { SqlContract, SqlStorage } from '@prisma-next/sql-contract/types';
import {
  AndExpr,
  type AnyExpression,
  type AstRewriter,
  BinaryExpr,
  type BinaryOp,
  ColumnRef,
  DerivedTableSource,
  EqColJoinOn,
  JoinAst,
  JsonArrayAggExpr,
  JsonObjectExpr,
  ListExpression,
  OrderByItem,
  OrExpr,
  ParamRef,
  ProjectionItem,
  SelectAst,
  SubqueryExpr,
  TableSource,
} from '@prisma-next/sql-relational-core/ast';
import type { SqlQueryPlan } from '@prisma-next/sql-relational-core/plan';
import { resolveColumnCodecId } from './collection-contract';
import { buildOrmQueryPlan, deriveParamsFromAst, resolveTableColumns } from './query-plan-meta';
import type { CollectionState, IncludeExpr, OrderExpr } from './types';
import { combineWhereExprs } from './where-utils';

type CursorOrderEntry = OrderExpr & {
  readonly value: unknown;
};

function buildProjection(
  contract: SqlContract<SqlStorage>,
  tableName: string,
  selectedFields: readonly string[] | undefined,
  tableRef = tableName,
): ProjectionItem[] {
  const columns =
    selectedFields && selectedFields.length > 0
      ? [...selectedFields]
      : resolveTableColumns(contract, tableName);

  return columns.map((column) => ProjectionItem.of(column, ColumnRef.of(tableRef, column)));
}

function toOrderBy(
  tableName: string,
  orderBy: readonly OrderExpr[] | undefined,
): ReadonlyArray<OrderByItem> | undefined {
  if (!orderBy || orderBy.length === 0) {
    return undefined;
  }

  return orderBy.map(
    (entry) => new OrderByItem(ColumnRef.of(tableName, entry.column), entry.direction),
  );
}

function columnParam(
  contract: SqlContract<SqlStorage>,
  tableName: string,
  column: string,
  value: unknown,
): ParamRef {
  const codecId = resolveColumnCodecId(contract, tableName, column);
  return ParamRef.of(value, { name: column, ...(codecId ? { codecId } : {}) });
}

function createBoundaryExpr(
  contract: SqlContract<SqlStorage>,
  tableName: string,
  entry: CursorOrderEntry,
): AnyExpression {
  const comparator: BinaryOp = entry.direction === 'asc' ? 'gt' : 'lt';
  return new BinaryExpr(
    comparator,
    ColumnRef.of(tableName, entry.column),
    columnParam(contract, tableName, entry.column, entry.value),
  );
}

function buildLexicographicCursorWhere(
  contract: SqlContract<SqlStorage>,
  tableName: string,
  entries: readonly CursorOrderEntry[],
): AnyExpression {
  const branches = entries.map((entry, index): AnyExpression => {
    const branchExprs: AnyExpression[] = [];

    for (const prefixEntry of entries.slice(0, index)) {
      branchExprs.push(
        BinaryExpr.eq(
          ColumnRef.of(tableName, prefixEntry.column),
          columnParam(contract, tableName, prefixEntry.column, prefixEntry.value),
        ),
      );
    }

    branchExprs.push(createBoundaryExpr(contract, tableName, entry));
    if (branchExprs.length === 1) {
      return branchExprs[0] as AnyExpression;
    }

    return AndExpr.of(branchExprs);
  });

  if (branches.length === 1) {
    return branches[0] as AnyExpression;
  }

  return OrExpr.of(branches);
}

function buildCursorWhere(
  contract: SqlContract<SqlStorage>,
  tableName: string,
  orderBy: readonly OrderExpr[] | undefined,
  cursor: Readonly<Record<string, unknown>> | undefined,
): AnyExpression | undefined {
  if (!cursor || !orderBy || orderBy.length === 0) {
    return undefined;
  }

  const entries: CursorOrderEntry[] = [];
  for (const order of orderBy) {
    const value = cursor[order.column];
    if (value === undefined) {
      throw new Error(`Missing cursor value for orderBy column "${order.column}"`);
    }
    entries.push({
      ...order,
      value,
    });
  }

  const firstEntry = entries[0];
  if (entries.length === 1 && firstEntry !== undefined) {
    return createBoundaryExpr(contract, tableName, firstEntry);
  }

  return buildLexicographicCursorWhere(contract, tableName, entries);
}

function createTableRefRemapper(fromTable: string, toTable: string): AstRewriter {
  return {
    columnRef: (col) => (col.table === fromTable ? ColumnRef.of(toTable, col.column) : col),
    tableSource: (source) => {
      if (source.alias === fromTable) return TableSource.named(source.name, toTable);
      if (!source.alias && source.name === fromTable)
        return TableSource.named(source.name, toTable);
      return source;
    },
    eqColJoinOn: (on) =>
      EqColJoinOn.of(
        on.left.table === fromTable ? ColumnRef.of(toTable, on.left.column) : on.left,
        on.right.table === fromTable ? ColumnRef.of(toTable, on.right.column) : on.right,
      ),
  };
}

function buildStateWhere(
  contract: SqlContract<SqlStorage>,
  tableName: string,
  state: CollectionState,
  options?: {
    readonly filterTableName?: string;
  },
): AnyExpression | undefined {
  const filterTableName = options?.filterTableName;
  const cursorTableName = filterTableName ?? tableName;
  const cursorWhere = buildCursorWhere(contract, cursorTableName, state.orderBy, state.cursor);
  const remappedFilters =
    filterTableName && filterTableName !== tableName
      ? state.filters.map((filter) =>
          filter.rewrite(createTableRefRemapper(filterTableName, tableName)),
        )
      : state.filters;
  const remappedCursorWhere =
    cursorWhere && filterTableName && filterTableName !== tableName
      ? cursorWhere.rewrite(createTableRefRemapper(filterTableName, tableName))
      : cursorWhere;
  const filters = remappedCursorWhere ? [...remappedFilters, remappedCursorWhere] : remappedFilters;
  return combineWhereExprs(filters);
}

function buildIncludeOrderArtifacts(
  relationName: string,
  childTableRef: string,
  rowAlias: string,
  orderBy: readonly OrderExpr[] | undefined,
): {
  readonly childOrderBy: ReadonlyArray<OrderByItem> | undefined;
  readonly hiddenOrderProjection: ReadonlyArray<ProjectionItem>;
  readonly aggregateOrderBy: ReadonlyArray<OrderByItem> | undefined;
} {
  const childOrderBy = toOrderBy(childTableRef, orderBy);
  if (!childOrderBy || childOrderBy.length === 0) {
    return {
      childOrderBy: undefined,
      hiddenOrderProjection: [],
      aggregateOrderBy: undefined,
    };
  }

  const hiddenOrderProjection = childOrderBy.map((orderItem, index) =>
    ProjectionItem.of(`${relationName}__order_${index}`, orderItem.expr),
  );
  const aggregateOrderBy = hiddenOrderProjection.map((projection, index) => {
    const orderItem = childOrderBy[index];
    if (!orderItem) {
      throw new Error(`Missing include order metadata at index ${index}`);
    }
    return new OrderByItem(ColumnRef.of(rowAlias, projection.alias), orderItem.dir);
  });

  return {
    childOrderBy,
    hiddenOrderProjection,
    aggregateOrderBy,
  };
}

function buildIncludeChildRowsSelect(
  contract: SqlContract<SqlStorage>,
  parentTableName: string,
  include: IncludeExpr,
): {
  readonly childRows: SelectAst;
  readonly childProjection: ReadonlyArray<ProjectionItem>;
  readonly rowsAlias: string;
  readonly aggregateOrderBy: ReadonlyArray<OrderByItem> | undefined;
} {
  const childState = include.nested;
  const childTableAlias =
    include.relatedTableName === parentTableName ? `${include.relationName}__child` : undefined;
  const childTableRef = childTableAlias ?? include.relatedTableName;
  const rowsAlias = `${include.relationName}__rows`;
  const childProjection = buildProjection(
    contract,
    include.relatedTableName,
    childState.selectedFields,
    childTableRef,
  );
  const { childOrderBy, hiddenOrderProjection, aggregateOrderBy } = buildIncludeOrderArtifacts(
    include.relationName,
    childTableRef,
    rowsAlias,
    childState.orderBy,
  );
  const childWhere = buildStateWhere(contract, childTableRef, childState, {
    filterTableName: include.relatedTableName,
  });
  const joinExpr = BinaryExpr.eq(
    ColumnRef.of(childTableRef, include.fkColumn),
    ColumnRef.of(parentTableName, include.parentPkColumn),
  );
  const whereExpr = childWhere ? AndExpr.of([joinExpr, childWhere]) : joinExpr;

  let childRows = SelectAst.from(TableSource.named(include.relatedTableName, childTableAlias))
    .withProjection([...childProjection, ...hiddenOrderProjection])
    .withWhere(whereExpr);

  if (childOrderBy) {
    childRows = childRows.withOrderBy(childOrderBy);
  }
  if (childState.distinctOn && childState.distinctOn.length > 0) {
    childRows = childRows.withDistinctOn(
      childState.distinctOn.map((column) => ColumnRef.of(childTableRef, column)),
    );
  } else if (childState.distinct && childState.distinct.length > 0) {
    childRows = childRows.withDistinct(true);
  }
  if (childState.limit !== undefined) {
    childRows = childRows.withLimit(childState.limit);
  }
  if (childState.offset !== undefined) {
    childRows = childRows.withOffset(childState.offset);
  }

  return {
    childRows,
    childProjection,
    rowsAlias,
    aggregateOrderBy,
  };
}

function buildLateralIncludeArtifacts(
  contract: SqlContract<SqlStorage>,
  parentTableName: string,
  include: IncludeExpr,
): {
  readonly join: JoinAst;
  readonly projection: ProjectionItem;
} {
  const { childRows, childProjection, rowsAlias, aggregateOrderBy } = buildIncludeChildRowsSelect(
    contract,
    parentTableName,
    include,
  );
  const lateralAlias = `${include.relationName}_lateral`;
  const jsonObjectExpr = JsonObjectExpr.fromEntries(
    childProjection.map((item) =>
      JsonObjectExpr.entry(item.alias, ColumnRef.of(rowsAlias, item.alias)),
    ),
  );

  const aggregateQuery = SelectAst.from(DerivedTableSource.as(rowsAlias, childRows)).withProjection(
    [
      ProjectionItem.of(
        include.relationName,
        JsonArrayAggExpr.of(jsonObjectExpr, 'emptyArray', aggregateOrderBy),
      ),
    ],
  );

  return {
    join: JoinAst.left(DerivedTableSource.as(lateralAlias, aggregateQuery), AndExpr.true(), true),
    projection: ProjectionItem.of(
      include.relationName,
      ColumnRef.of(lateralAlias, include.relationName),
    ),
  };
}

function buildCorrelatedIncludeProjection(
  contract: SqlContract<SqlStorage>,
  parentTableName: string,
  include: IncludeExpr,
): {
  readonly projection: ProjectionItem;
} {
  const { childRows, childProjection, rowsAlias, aggregateOrderBy } = buildIncludeChildRowsSelect(
    contract,
    parentTableName,
    include,
  );
  const jsonObjectExpr = JsonObjectExpr.fromEntries(
    childProjection.map((item) =>
      JsonObjectExpr.entry(item.alias, ColumnRef.of(rowsAlias, item.alias)),
    ),
  );
  const aggregateQuery = SelectAst.from(DerivedTableSource.as(rowsAlias, childRows)).withProjection(
    [
      ProjectionItem.of(
        include.relationName,
        JsonArrayAggExpr.of(jsonObjectExpr, 'emptyArray', aggregateOrderBy),
      ),
    ],
  );

  return {
    projection: ProjectionItem.of(include.relationName, SubqueryExpr.of(aggregateQuery)),
  };
}

function buildSelectAst(
  contract: SqlContract<SqlStorage>,
  tableName: string,
  state: CollectionState,
  options: {
    readonly joins?: ReadonlyArray<JoinAst>;
    readonly includeProjection?: ReadonlyArray<ProjectionItem>;
    readonly where?: AnyExpression;
  } = {},
): SelectAst {
  const scalarProjection = buildProjection(contract, tableName, state.selectedFields);
  const projection = [...scalarProjection, ...(options.includeProjection ?? [])];
  const where = options.where ?? buildStateWhere(contract, tableName, state);
  const orderBy = toOrderBy(tableName, state.orderBy);

  let ast = SelectAst.from(TableSource.named(tableName)).withProjection(projection);
  if (where) {
    ast = ast.withWhere(where);
  }
  if (orderBy) {
    ast = ast.withOrderBy(orderBy);
  }
  if (state.selectedFields === undefined) {
    ast = ast.withSelectAllIntent({ table: tableName });
  }
  if (state.distinctOn && state.distinctOn.length > 0) {
    ast = ast.withDistinctOn(state.distinctOn.map((column) => ColumnRef.of(tableName, column)));
  } else if (state.distinct && state.distinct.length > 0) {
    ast = ast.withDistinct(true);
  }
  if (state.limit !== undefined) {
    ast = ast.withLimit(state.limit);
  }
  if (state.offset !== undefined) {
    ast = ast.withOffset(state.offset);
  }
  if (options.joins && options.joins.length > 0) {
    ast = ast.withJoins(options.joins);
  }

  return ast;
}

export function compileSelect(
  contract: SqlContract<SqlStorage>,
  tableName: string,
  state: CollectionState,
): SqlQueryPlan<Record<string, unknown>> {
  const ast = buildSelectAst(contract, tableName, {
    ...state,
    includes: [],
  });

  const { params, paramDescriptors } = deriveParamsFromAst(ast);
  return buildOrmQueryPlan(contract, ast, params, paramDescriptors);
}

export function compileRelationSelect(
  contract: SqlContract<SqlStorage>,
  relatedTableName: string,
  fkColumn: string,
  parentPks: readonly unknown[],
  nestedState: CollectionState,
): SqlQueryPlan<Record<string, unknown>> {
  const inFilter: AnyExpression = BinaryExpr.in(
    ColumnRef.of(relatedTableName, fkColumn),
    ListExpression.of(parentPks.map((pk) => columnParam(contract, relatedTableName, fkColumn, pk))),
  );

  return compileSelect(contract, relatedTableName, {
    ...nestedState,
    includes: [],
    limit: undefined,
    offset: undefined,
    filters: [inFilter, ...nestedState.filters],
  });
}

export function compileSelectWithIncludeStrategy(
  contract: SqlContract<SqlStorage>,
  tableName: string,
  state: CollectionState,
  strategy: 'lateral' | 'correlated',
): SqlQueryPlan<Record<string, unknown>> {
  if (
    state.includes.some((include) => include.scalar !== undefined || include.combine !== undefined)
  ) {
    throw new Error(
      'single-query include strategy does not support scalar include selectors or combine()',
    );
  }

  const includeJoins: JoinAst[] = [];
  const includeProjection: ProjectionItem[] = [];
  const topLevelWhere = buildStateWhere(contract, tableName, state);

  for (const include of state.includes) {
    if (strategy === 'lateral') {
      const artifact = buildLateralIncludeArtifacts(contract, tableName, include);
      includeJoins.push(artifact.join);
      includeProjection.push(artifact.projection);
      continue;
    }
    const artifact = buildCorrelatedIncludeProjection(contract, tableName, include);
    includeProjection.push(artifact.projection);
  }

  const ast = buildSelectAst(
    contract,
    tableName,
    {
      ...state,
      includes: [],
    },
    {
      joins: includeJoins,
      includeProjection,
      ...(topLevelWhere ? { where: topLevelWhere } : {}),
    },
  );

  const { params, paramDescriptors } = deriveParamsFromAst(ast);
  return buildOrmQueryPlan(contract, ast, params, paramDescriptors);
}