metal-orm 1.0.11 → 1.0.13

package/src/orm/unit-of-work.ts

@@ -1,11 +1,11 @@
-import { eq } from '../core/ast/expression.js';
-import type { Dialect, CompiledQuery } from '../core/dialect/abstract.js';
+import { ColumnNode, eq } from '../core/ast/expression.js';
+import type { Dialect, CompiledQuery } from '../core/dialect/abstract.js';
 import { InsertQueryBuilder } from '../query-builder/insert.js';
 import { UpdateQueryBuilder } from '../query-builder/update.js';
 import { DeleteQueryBuilder } from '../query-builder/delete.js';
 import { findPrimaryKey } from '../query-builder/hydration-planner.js';
 import type { TableDef, TableHooks } from '../schema/table.js';
-import type { DbExecutor } from './db-executor.js';
+import type { DbExecutor, QueryResult } from './db-executor.js';
 import { IdentityMap } from './identity-map.js';
 import { EntityStatus } from './runtime-types.js';
 import type { TrackedEntity } from './runtime-types.js';
@@ -120,10 +120,14 @@ export class UnitOfWork {
   private async flushInsert(tracked: TrackedEntity): Promise<void> {
     await this.runHook(tracked.table.hooks?.beforeInsert, tracked);
 
-    const payload = this.extractColumns(tracked.table, tracked.entity);
-    const builder = new InsertQueryBuilder(tracked.table).values(payload);
-    const compiled = builder.compile(this.dialect);
-    await this.executeCompiled(compiled);
+    const payload = this.extractColumns(tracked.table, tracked.entity);
+    let builder = new InsertQueryBuilder(tracked.table).values(payload);
+    if (this.dialect.supportsReturning()) {
+      builder = builder.returning(...this.getReturningColumns(tracked.table));
+    }
+    const compiled = builder.compile(this.dialect);
+    const results = await this.executeCompiled(compiled);
+    this.applyReturningResults(tracked, results);
 
     tracked.status = EntityStatus.Managed;
     tracked.original = this.createSnapshot(tracked.table, tracked.entity);
@@ -146,12 +150,17 @@ export class UnitOfWork {
     const pkColumn = tracked.table.columns[findPrimaryKey(tracked.table)];
     if (!pkColumn) return;
 
-    const builder = new UpdateQueryBuilder(tracked.table)
-      .set(changes)
-      .where(eq(pkColumn, tracked.pk));
-
-    const compiled = builder.compile(this.dialect);
-    await this.executeCompiled(compiled);
+    let builder = new UpdateQueryBuilder(tracked.table)
+      .set(changes)
+      .where(eq(pkColumn, tracked.pk));
+
+    if (this.dialect.supportsReturning()) {
+      builder = builder.returning(...this.getReturningColumns(tracked.table));
+    }
+
+    const compiled = builder.compile(this.dialect);
+    const results = await this.executeCompiled(compiled);
+    this.applyReturningResults(tracked, results);
 
     tracked.status = EntityStatus.Managed;
     tracked.original = this.createSnapshot(tracked.table, tracked.entity);
@@ -206,16 +215,44 @@ export class UnitOfWork {
     return payload;
   }
 
-  private async executeCompiled(compiled: CompiledQuery): Promise<void> {
-    await this.executor.executeSql(compiled.sql, compiled.params);
-  }
-
-  private registerIdentity(tracked: TrackedEntity): void {
-    if (tracked.pk == null) return;
-    this.identityMap.register(tracked);
-  }
-
-  private createSnapshot(table: TableDef, entity: any): Record<string, any> {
+  private async executeCompiled(compiled: CompiledQuery): Promise<QueryResult[]> {
+    return this.executor.executeSql(compiled.sql, compiled.params);
+  }
+
+  private getReturningColumns(table: TableDef): ColumnNode[] {
+    return Object.values(table.columns).map(column => ({
+      type: 'Column',
+      table: table.name,
+      name: column.name,
+      alias: column.name
+    }));
+  }
+
+  private applyReturningResults(tracked: TrackedEntity, results: QueryResult[]): void {
+    if (!this.dialect.supportsReturning()) return;
+    const first = results[0];
+    if (!first || first.values.length === 0) return;
+
+    const row = first.values[0];
+    for (let i = 0; i < first.columns.length; i++) {
+      const columnName = this.normalizeColumnName(first.columns[i]);
+      if (!(columnName in tracked.table.columns)) continue;
+      tracked.entity[columnName] = row[i];
+    }
+  }
+
+  private normalizeColumnName(column: string): string {
+    const parts = column.split('.');
+    const candidate = parts[parts.length - 1];
+    return candidate.replace(/^["`[\]]+|["`[\]]+$/g, '');
+  }
+
+  private registerIdentity(tracked: TrackedEntity): void {
+    if (tracked.pk == null) return;
+    this.identityMap.register(tracked);
+  }
+
+  private createSnapshot(table: TableDef, entity: any): Record<string, any> {
     const snapshot: Record<string, any> = {};
     for (const column of Object.keys(table.columns)) {
       snapshot[column] = entity[column];

package/src/query-builder/hydration-manager.ts

@@ -1,8 +1,11 @@
-import { TableDef } from '../schema/table.js';
-import { RelationDef } from '../schema/relation.js';
-import { SelectQueryNode, HydrationPlan } from '../core/ast/query.js';
-import { HydrationPlanner } from './hydration-planner.js';
-import { SelectQueryState, ProjectionNode } from './select-query-state.js';
+import { TableDef } from '../schema/table.js';
+import { RelationDef, RelationKinds } from '../schema/relation.js';
+import { CommonTableExpressionNode, HydrationPlan, OrderByNode, SelectQueryNode } from '../core/ast/query.js';
+import { HydrationPlanner } from './hydration-planner.js';
+import { ProjectionNode, SelectQueryState } from './select-query-state.js';
+import { ColumnNode, eq } from '../core/ast/expression.js';
+import { createJoinNode } from '../core/ast/join-node.js';
+import { JOIN_KINDS } from '../core/sql/sql.js';
 
 /**
  * Manages hydration planning for query results
@@ -60,28 +63,229 @@ export class HydrationManager {
     return this.clone(next);
   }
 
-  /**
-   * Applies hydration plan to the AST
-   * @param ast - Query AST to modify
-   * @returns AST with hydration metadata
-   */
-  applyToAst(ast: SelectQueryNode): SelectQueryNode {
-    const plan = this.planner.getPlan();
-    if (!plan) return ast;
-    return {
-      ...ast,
-      meta: {
-        ...(ast.meta || {}),
-        hydration: plan
-      }
-    };
-  }
+  /**
+   * Applies hydration plan to the AST
+   * @param ast - Query AST to modify
+   * @returns AST with hydration metadata
+   */
+  applyToAst(ast: SelectQueryNode): SelectQueryNode {
+    // Hydration is not applied to compound set queries since row identity is ambiguous.
+    if (ast.setOps && ast.setOps.length > 0) {
+      return ast;
+    }
+
+    const plan = this.planner.getPlan();
+    if (!plan) return ast;
+
+    const needsPaginationGuard = this.requiresParentPagination(ast, plan);
+    const rewritten = needsPaginationGuard ? this.wrapForParentPagination(ast, plan) : ast;
+    return this.attachHydrationMeta(rewritten, plan);
+  }
 
   /**
    * Gets the current hydration plan
    * @returns Hydration plan or undefined if none exists
    */
-  getPlan(): HydrationPlan | undefined {
-    return this.planner.getPlan();
-  }
-}
+  getPlan(): HydrationPlan | undefined {
+    return this.planner.getPlan();
+  }
+
+  /**
+   * Attaches hydration metadata to a query AST node.
+   */
+  private attachHydrationMeta(ast: SelectQueryNode, plan: HydrationPlan): SelectQueryNode {
+    return {
+      ...ast,
+      meta: {
+        ...(ast.meta || {}),
+        hydration: plan
+      }
+    };
+  }
+
+  /**
+   * Determines whether the query needs pagination rewriting to keep LIMIT/OFFSET
+   * applied to parent rows when eager-loading multiplicative relations.
+   */
+  private requiresParentPagination(ast: SelectQueryNode, plan: HydrationPlan): boolean {
+    const hasPagination = ast.limit !== undefined || ast.offset !== undefined;
+    return hasPagination && this.hasMultiplyingRelations(plan);
+  }
+
+  private hasMultiplyingRelations(plan: HydrationPlan): boolean {
+    return plan.relations.some(
+      rel => rel.type === RelationKinds.HasMany || rel.type === RelationKinds.BelongsToMany
+    );
+  }
+
+  /**
+   * Rewrites the query using CTEs so LIMIT/OFFSET target distinct parent rows
+   * instead of the joined result set.
+   *
+   * The strategy:
+   * - Hoist the original query (minus limit/offset) into a base CTE.
+   * - Select distinct parent ids from that base CTE with the original ordering and pagination.
+   * - Join the base CTE against the paged ids to retrieve the joined rows for just that page.
+   */
+  private wrapForParentPagination(ast: SelectQueryNode, plan: HydrationPlan): SelectQueryNode {
+    const projectionNames = this.getProjectionNames(ast.columns);
+    if (!projectionNames) {
+      return ast;
+    }
+
+    const projectionAliases = this.buildProjectionAliasMap(ast.columns);
+    const projectionSet = new Set(projectionNames);
+    const rootPkAlias = projectionAliases.get(`${plan.rootTable}.${plan.rootPrimaryKey}`) ?? plan.rootPrimaryKey;
+
+    const baseCteName = this.nextCteName(ast.ctes, '__metal_pagination_base');
+    const baseQuery: SelectQueryNode = {
+      ...ast,
+      ctes: undefined,
+      limit: undefined,
+      offset: undefined,
+      orderBy: undefined,
+      meta: undefined
+    };
+
+    const baseCte: CommonTableExpressionNode = {
+      type: 'CommonTableExpression',
+      name: baseCteName,
+      query: baseQuery,
+      recursive: false
+    };
+
+    const orderBy = this.mapOrderBy(ast.orderBy, plan, projectionAliases, baseCteName, projectionSet);
+    // When an order-by uses child-table columns we cannot safely rewrite pagination,
+    // so preserve the original query to avoid changing semantics.
+    if (orderBy === null) {
+      return ast;
+    }
+
+    const pageCteName = this.nextCteName([...(ast.ctes ?? []), baseCte], '__metal_pagination_page');
+    const pagingColumns = this.buildPagingColumns(rootPkAlias, orderBy, baseCteName);
+
+    const pageCte: CommonTableExpressionNode = {
+      type: 'CommonTableExpression',
+      name: pageCteName,
+      query: {
+        type: 'SelectQuery',
+        from: { type: 'Table', name: baseCteName },
+        columns: pagingColumns,
+        joins: [],
+        distinct: [{ type: 'Column', table: baseCteName, name: rootPkAlias }],
+        orderBy,
+        limit: ast.limit,
+        offset: ast.offset
+      },
+      recursive: false
+    };
+
+    const joinCondition = eq(
+      { type: 'Column', table: baseCteName, name: rootPkAlias },
+      { type: 'Column', table: pageCteName, name: rootPkAlias }
+    );
+
+    const outerColumns: ColumnNode[] = projectionNames.map(name => ({
+      type: 'Column',
+      table: baseCteName,
+      name,
+      alias: name
+    }));
+
+    return {
+      type: 'SelectQuery',
+      from: { type: 'Table', name: baseCteName },
+      columns: outerColumns,
+      joins: [createJoinNode(JOIN_KINDS.INNER, pageCteName, joinCondition)],
+      orderBy,
+      ctes: [...(ast.ctes ?? []), baseCte, pageCte]
+    };
+  }
+
+  private nextCteName(existing: CommonTableExpressionNode[] | undefined, baseName: string): string {
+    const names = new Set((existing ?? []).map(cte => cte.name));
+    let candidate = baseName;
+    let suffix = 1;
+
+    while (names.has(candidate)) {
+      suffix += 1;
+      candidate = `${baseName}_${suffix}`;
+    }
+
+    return candidate;
+  }
+
+  private getProjectionNames(columns: ProjectionNode[]): string[] | undefined {
+    const names: string[] = [];
+    for (const col of columns) {
+      const alias = (col as any).alias ?? (col as any).name;
+      if (!alias) return undefined;
+      names.push(alias);
+    }
+    return names;
+  }
+
+  private buildProjectionAliasMap(columns: ProjectionNode[]): Map<string, string> {
+    const map = new Map<string, string>();
+    for (const col of columns) {
+      if ((col as ColumnNode).type !== 'Column') continue;
+      const node = col as ColumnNode;
+      const key = `${node.table}.${node.name}`;
+      map.set(key, node.alias ?? node.name);
+    }
+    return map;
+  }
+
+  private mapOrderBy(
+    orderBy: OrderByNode[] | undefined,
+    plan: HydrationPlan,
+    projectionAliases: Map<string, string>,
+    baseAlias: string,
+    availableColumns: Set<string>
+  ): OrderByNode[] | undefined | null {
+    if (!orderBy || orderBy.length === 0) {
+      return undefined;
+    }
+
+    const mapped: OrderByNode[] = [];
+
+    for (const ob of orderBy) {
+      // Only rewrite when ordering by root columns; child columns would reintroduce the pagination bug.
+      if (ob.column.table !== plan.rootTable) {
+        return null;
+      }
+
+      const alias = projectionAliases.get(`${ob.column.table}.${ob.column.name}`) ?? ob.column.name;
+      if (!availableColumns.has(alias)) {
+        return null;
+      }
+
+      mapped.push({
+        type: 'OrderBy',
+        column: { type: 'Column', table: baseAlias, name: alias },
+        direction: ob.direction
+      });
+    }
+
+    return mapped;
+  }
+
+  private buildPagingColumns(primaryKey: string, orderBy: OrderByNode[] | undefined, tableAlias: string): ColumnNode[] {
+    const columns: ColumnNode[] = [{ type: 'Column', table: tableAlias, name: primaryKey, alias: primaryKey }];
+
+    if (!orderBy) return columns;
+
+    for (const ob of orderBy) {
+      if (!columns.some(col => col.name === ob.column.name)) {
+        columns.push({
+          type: 'Column',
+          table: tableAlias,
+          name: ob.column.name,
+          alias: ob.column.name
+        });
+      }
+    }
+
+    return columns;
+  }
+}

package/src/query-builder/query-ast-service.ts

@@ -1,6 +1,6 @@
 import { TableDef } from '../schema/table.js';
 import { ColumnDef } from '../schema/column.js';
-import { SelectQueryNode, CommonTableExpressionNode } from '../core/ast/query.js';
+import { SelectQueryNode, CommonTableExpressionNode, SetOperationKind, SetOperationNode } from '../core/ast/query.js';
 import { buildColumnNode } from '../core/ast/builders.js';
 import {
   ColumnNode,
@@ -95,17 +95,32 @@ export class QueryAstService {
    * @param recursive - Whether the CTE is recursive
    * @returns Updated query state with CTE
    */
-  withCte(name: string, query: SelectQueryNode, columns?: string[], recursive = false): SelectQueryState {
-    const cte: CommonTableExpressionNode = {
-      type: 'CommonTableExpression',
-      name,
-      query,
-      columns,
-      recursive
-    };
-
-    return this.state.withCte(cte);
-  }
+  withCte(name: string, query: SelectQueryNode, columns?: string[], recursive = false): SelectQueryState {
+    const cte: CommonTableExpressionNode = {
+      type: 'CommonTableExpression',
+      name,
+      query,
+      columns,
+      recursive
+    };
+
+    return this.state.withCte(cte);
+  }
+
+  /**
+   * Adds a set operation (UNION/UNION ALL/INTERSECT/EXCEPT) to the query
+   * @param operator - Set operator
+   * @param query - Right-hand side query
+   * @returns Updated query state with set operation
+   */
+  withSetOperation(operator: SetOperationKind, query: SelectQueryNode): SelectQueryState {
+    const op: SetOperationNode = {
+      type: 'SetOperation',
+      operator,
+      query
+    };
+    return this.state.withSetOperation(op);
+  }
 
   /**
    * Selects a subquery as a column

package/src/query-builder/select-query-state.ts

@@ -1,5 +1,5 @@
 import { TableDef } from '../schema/table.js';
-import { SelectQueryNode, CommonTableExpressionNode, OrderByNode } from '../core/ast/query.js';
+import { SelectQueryNode, CommonTableExpressionNode, OrderByNode, SetOperationNode } from '../core/ast/query.js';
 import {
   ColumnNode,
   ExpressionNode,
@@ -166,14 +166,26 @@ export class SelectQueryState {
   }
 
   /**
-   * Adds a Common Table Expression (CTE) to the query
-   * @param cte - CTE node to add
-   * @returns New SelectQueryState with CTE
-   */
-  withCte(cte: CommonTableExpressionNode): SelectQueryState {
-    return this.clone({
-      ...this.ast,
-      ctes: [...(this.ast.ctes ?? []), cte]
-    });
-  }
-}
+   * Adds a Common Table Expression (CTE) to the query
+   * @param cte - CTE node to add
+   * @returns New SelectQueryState with CTE
+   */
+  withCte(cte: CommonTableExpressionNode): SelectQueryState {
+    return this.clone({
+      ...this.ast,
+      ctes: [...(this.ast.ctes ?? []), cte]
+    });
+  }
+
+  /**
+   * Adds a set operation (UNION/INTERSECT/EXCEPT) to the query
+   * @param op - Set operation node to add
+   * @returns New SelectQueryState with set operation
+   */
+  withSetOperation(op: SetOperationNode): SelectQueryState {
+    return this.clone({
+      ...this.ast,
+      setOps: [...(this.ast.setOps ?? []), op]
+    });
+  }
+}

package/src/query-builder/select.ts

@@ -1,6 +1,6 @@
 import { TableDef } from '../schema/table.js';
 import { ColumnDef } from '../schema/column.js';
-import { SelectQueryNode, HydrationPlan } from '../core/ast/query.js';
+import { SelectQueryNode, HydrationPlan, SetOperationKind } from '../core/ast/query.js';
 import {
   ColumnNode,
   ExpressionNode,
@@ -96,15 +96,23 @@ export class SelectQueryBuilder<T = any, TTable extends TableDef = TableDef> {
     return { state: nextState, hydration: context.hydration };
   }
 
-  private applyJoin(
-    context: SelectQueryBuilderContext,
-    table: TableDef,
-    condition: BinaryExpressionNode,
-    kind: JoinKind
-  ): SelectQueryBuilderContext {
-    const joinNode = createJoinNode(kind, table.name, condition);
-    return this.applyAst(context, service => service.withJoin(joinNode));
-  }
+  private applyJoin(
+    context: SelectQueryBuilderContext,
+    table: TableDef,
+    condition: BinaryExpressionNode,
+    kind: JoinKind
+  ): SelectQueryBuilderContext {
+    const joinNode = createJoinNode(kind, table.name, condition);
+    return this.applyAst(context, service => service.withJoin(joinNode));
+  }
+
+  private applySetOperation(
+    operator: SetOperationKind,
+    query: SelectQueryBuilder<any, TableDef<any>> | SelectQueryNode
+  ): SelectQueryBuilderContext {
+    const subAst = this.resolveQueryNode(query);
+    return this.applyAst(this.context, service => service.withSetOperation(operator, subAst));
+  }
 
   /**
    * Selects specific columns for the query
@@ -315,10 +323,46 @@ export class SelectQueryBuilder<T = any, TTable extends TableDef = TableDef> {
    * @param n - Number of rows to skip
    * @returns New query builder instance with the OFFSET clause
    */
-  offset(n: number): SelectQueryBuilder<T, TTable> {
-    const nextContext = this.applyAst(this.context, service => service.withOffset(n));
-    return this.clone(nextContext);
-  }
+  offset(n: number): SelectQueryBuilder<T, TTable> {
+    const nextContext = this.applyAst(this.context, service => service.withOffset(n));
+    return this.clone(nextContext);
+  }
+
+  /**
+   * Combines this query with another using UNION
+   * @param query - Query to union with
+   * @returns New query builder instance with the set operation
+   */
+  union(query: SelectQueryBuilder<any, TableDef<any>> | SelectQueryNode): SelectQueryBuilder<T, TTable> {
+    return this.clone(this.applySetOperation('UNION', query));
+  }
+
+  /**
+   * Combines this query with another using UNION ALL
+   * @param query - Query to union with
+   * @returns New query builder instance with the set operation
+   */
+  unionAll(query: SelectQueryBuilder<any, TableDef<any>> | SelectQueryNode): SelectQueryBuilder<T, TTable> {
+    return this.clone(this.applySetOperation('UNION ALL', query));
+  }
+
+  /**
+   * Combines this query with another using INTERSECT
+   * @param query - Query to intersect with
+   * @returns New query builder instance with the set operation
+   */
+  intersect(query: SelectQueryBuilder<any, TableDef<any>> | SelectQueryNode): SelectQueryBuilder<T, TTable> {
+    return this.clone(this.applySetOperation('INTERSECT', query));
+  }
+
+  /**
+   * Combines this query with another using EXCEPT
+   * @param query - Query to subtract
+   * @returns New query builder instance with the set operation
+   */
+  except(query: SelectQueryBuilder<any, TableDef<any>> | SelectQueryNode): SelectQueryBuilder<T, TTable> {
+    return this.clone(this.applySetOperation('EXCEPT', query));
+  }
 
   /**
    * Adds a WHERE EXISTS condition to the query