metal-orm 1.0.57 → 1.0.58

package/src/orm/lazy-batch.ts

@@ -1,11 +1,13 @@
 import { TableDef } from '../schema/table.js';
 import { BelongsToManyRelation, HasManyRelation, HasOneRelation, BelongsToRelation } from '../schema/relation.js';
 import { SelectQueryBuilder } from '../query-builder/select.js';
-import { inList, LiteralNode } from '../core/ast/expression.js';
+import { ExpressionNode, inList, LiteralNode } from '../core/ast/expression.js';
 import { EntityContext } from './entity-context.js';
 import type { QueryResult } from '../core/execution/db-executor.js';
 import { ColumnDef } from '../schema/column-types.js';
 import { findPrimaryKey } from '../query-builder/hydration-planner.js';
+import { RelationIncludeOptions } from '../query-builder/relation-types.js';
+import { buildDefaultPivotColumns } from '../query-builder/relation-utils.js';
 
 /**
  * An array of database rows, each represented as a record of string keys to unknown values.
@@ -17,16 +19,35 @@ type Rows = Record<string, unknown>[];
  */
 type EntityTracker = ReturnType<EntityContext['getEntitiesForTable']>[number];
 
-/**
- * Creates a record of all columns from the given table definition.
- * @param table - The table definition to select columns from.
- * @returns A record mapping column names to their definitions.
- */
-const selectAllColumns = (table: TableDef): Record<string, ColumnDef> =>
-  Object.entries(table.columns).reduce((acc, [name, def]) => {
-    acc[name] = def;
+const hasColumns = (columns?: readonly string[]): columns is readonly string[] =>
+  Boolean(columns && columns.length > 0);
+
+const buildColumnSelection = (
+  table: TableDef,
+  columns: string[],
+  missingMsg: (col: string) => string
+): Record<string, ColumnDef> => {
+  return columns.reduce((acc, column) => {
+    const def = table.columns[column];
+    if (!def) {
+      throw new Error(missingMsg(column));
+    }
+    acc[column] = def;
     return acc;
   }, {} as Record<string, ColumnDef>);
+};
+
+const filterRow = (row: Record<string, unknown>, columns: Set<string>): Record<string, unknown> => {
+  const filtered: Record<string, unknown> = {};
+  for (const column of columns) {
+    if (column in row) {
+      filtered[column] = row[column];
+    }
+  }
+  return filtered;
+};
+
+const filterRows = (rows: Rows, columns: Set<string>): Rows => rows.map(row => filterRow(row, columns));
 
 /**
  * Extracts rows from query results into a standardized format.
@@ -54,11 +75,11 @@ const rowsFromResults = (results: QueryResult[]): Rows => {
  * @param qb - The select query builder.
  * @returns A promise resolving to the rows from the query.
  */
-const executeQuery = async (ctx: EntityContext, qb: SelectQueryBuilder<unknown, TableDef>): Promise<Rows> => {
-  const compiled = ctx.dialect.compileSelect(qb.getAST());
-  const results = await ctx.executor.executeSql(compiled.sql, compiled.params);
-  return rowsFromResults(results);
-};
+const executeQuery = async (ctx: EntityContext, qb: SelectQueryBuilder<unknown, TableDef>): Promise<Rows> => {
+  const compiled = ctx.dialect.compileSelect(qb.getAST());
+  const results = await ctx.executor.executeSql(compiled.sql, compiled.params);
+  return rowsFromResults(results);
+};
 
 /**
  * Converts a value to a string key, handling null and undefined as empty string.
@@ -100,16 +121,21 @@ const buildInListValues = (keys: Set<unknown>): (string | number | LiteralNode)[
  * @param keys - The set of keys to match.
  * @returns A promise resolving to the matching rows.
  */
-const fetchRowsForKeys = async (
-  ctx: EntityContext,
-  table: TableDef,
-  column: ColumnDef,
-  keys: Set<unknown>
-): Promise<Rows> => {
-  const qb = new SelectQueryBuilder(table).select(selectAllColumns(table));
-  qb.where(inList(column, buildInListValues(keys)));
-  return executeQuery(ctx, qb);
-};
+const fetchRowsForKeys = async (
+  ctx: EntityContext,
+  table: TableDef,
+  column: ColumnDef,
+  keys: Set<unknown>,
+  selection: Record<string, ColumnDef>,
+  filter?: ExpressionNode
+): Promise<Rows> => {
+  let qb = new SelectQueryBuilder(table).select(selection);
+  qb = qb.where(inList(column, buildInListValues(keys)));
+  if (filter) {
+    qb = qb.where(filter);
+  }
+  return executeQuery(ctx, qb);
+};
 
 /**
  * Groups rows by the value of a key column, allowing multiple rows per key.
@@ -160,8 +186,9 @@ const groupRowsByUnique = (rows: Rows, keyColumn: string): Map<string, Record<st
 export const loadHasManyRelation = async (
   ctx: EntityContext,
   rootTable: TableDef,
-  _relationName: string,
-  relation: HasManyRelation
+  relationName: string,
+  relation: HasManyRelation,
+  options?: RelationIncludeOptions
 ): Promise<Map<string, Rows>> => {
   const localKey = relation.localKey || findPrimaryKey(rootTable);
   const roots = ctx.getEntitiesForTable(rootTable);
@@ -174,8 +201,33 @@ export const loadHasManyRelation = async (
   const fkColumn = relation.target.columns[relation.foreignKey];
   if (!fkColumn) return new Map();
 
-  const rows = await fetchRowsForKeys(ctx, relation.target, fkColumn, keys);
-  return groupRowsByMany(rows, relation.foreignKey);
+  const requestedColumns = hasColumns(options?.columns) ? [...options!.columns] : undefined;
+  const targetPrimaryKey = findPrimaryKey(relation.target);
+  const selectedColumns = requestedColumns ? [...requestedColumns] : Object.keys(relation.target.columns);
+  if (!selectedColumns.includes(targetPrimaryKey)) {
+    selectedColumns.push(targetPrimaryKey);
+  }
+
+  const queryColumns = new Set(selectedColumns);
+  queryColumns.add(relation.foreignKey);
+
+  const selection = buildColumnSelection(
+    relation.target,
+    Array.from(queryColumns),
+    column => `Column '${column}' not found on relation '${relationName}'`
+  );
+
+  const rows = await fetchRowsForKeys(ctx, relation.target, fkColumn, keys, selection, options?.filter);
+  const grouped = groupRowsByMany(rows, relation.foreignKey);
+
+  if (!requestedColumns) return grouped;
+
+  const visibleColumns = new Set(selectedColumns);
+  const filtered = new Map<string, Rows>();
+  for (const [key, bucket] of grouped.entries()) {
+    filtered.set(key, filterRows(bucket, visibleColumns));
+  }
+  return filtered;
 };
 
 /**
@@ -189,8 +241,9 @@ export const loadHasManyRelation = async (
 export const loadHasOneRelation = async (
   ctx: EntityContext,
   rootTable: TableDef,
-  _relationName: string,
-  relation: HasOneRelation
+  relationName: string,
+  relation: HasOneRelation,
+  options?: RelationIncludeOptions
 ): Promise<Map<string, Record<string, unknown>>> => {
   const localKey = relation.localKey || findPrimaryKey(rootTable);
   const roots = ctx.getEntitiesForTable(rootTable);
@@ -203,8 +256,33 @@ export const loadHasOneRelation = async (
   const fkColumn = relation.target.columns[relation.foreignKey];
   if (!fkColumn) return new Map();
 
-  const rows = await fetchRowsForKeys(ctx, relation.target, fkColumn, keys);
-  return groupRowsByUnique(rows, relation.foreignKey);
+  const requestedColumns = hasColumns(options?.columns) ? [...options!.columns] : undefined;
+  const targetPrimaryKey = findPrimaryKey(relation.target);
+  const selectedColumns = requestedColumns ? [...requestedColumns] : Object.keys(relation.target.columns);
+  if (!selectedColumns.includes(targetPrimaryKey)) {
+    selectedColumns.push(targetPrimaryKey);
+  }
+
+  const queryColumns = new Set(selectedColumns);
+  queryColumns.add(relation.foreignKey);
+
+  const selection = buildColumnSelection(
+    relation.target,
+    Array.from(queryColumns),
+    column => `Column '${column}' not found on relation '${relationName}'`
+  );
+
+  const rows = await fetchRowsForKeys(ctx, relation.target, fkColumn, keys, selection, options?.filter);
+  const grouped = groupRowsByUnique(rows, relation.foreignKey);
+
+  if (!requestedColumns) return grouped;
+
+  const visibleColumns = new Set(selectedColumns);
+  const filtered = new Map<string, Record<string, unknown>>();
+  for (const [key, row] of grouped.entries()) {
+    filtered.set(key, filterRow(row, visibleColumns));
+  }
+  return filtered;
 };
 
 /**
@@ -215,25 +293,91 @@ export const loadHasOneRelation = async (
  * @param relation - The belongs-to relation definition.
  * @returns A promise resolving to a map of foreign keys to single related rows.
  */
-export const loadBelongsToRelation = async (
-  ctx: EntityContext,
-  rootTable: TableDef,
-  _relationName: string,
-  relation: BelongsToRelation
-): Promise<Map<string, Record<string, unknown>>> => {
-  const roots = ctx.getEntitiesForTable(rootTable);
-  const foreignKeys = collectKeysFromRoots(roots, relation.foreignKey);
-
-  if (!foreignKeys.size) {
-    return new Map();
-  }
-
-  const targetKey = relation.localKey || findPrimaryKey(relation.target);
+export const loadBelongsToRelation = async (
+  ctx: EntityContext,
+  rootTable: TableDef,
+  relationName: string,
+  relation: BelongsToRelation,
+  options?: RelationIncludeOptions
+): Promise<Map<string, Record<string, unknown>>> => {
+  const roots = ctx.getEntitiesForTable(rootTable);
+
+  const getForeignKeys = (): Set<unknown> => collectKeysFromRoots(roots, relation.foreignKey);
+  let foreignKeys = getForeignKeys();
+
+  if (!foreignKeys.size) {
+    const pkName = findPrimaryKey(rootTable);
+    const pkColumn = rootTable.columns[pkName];
+    const fkColumn = rootTable.columns[relation.foreignKey];
+
+    if (pkColumn && fkColumn) {
+      const missingKeys = new Set<unknown>();
+      const entityByPk = new Map<unknown, Record<string, unknown>>();
+
+      for (const tracked of roots) {
+        const entity = tracked.entity as Record<string, unknown>;
+        const pkValue = entity[pkName];
+        if (pkValue === undefined || pkValue === null) continue;
+        const fkValue = entity[relation.foreignKey];
+        if (fkValue === undefined || fkValue === null) {
+          missingKeys.add(pkValue);
+          entityByPk.set(pkValue, entity);
+        }
+      }
+
+      if (missingKeys.size) {
+        const selection = buildColumnSelection(
+          rootTable,
+          [pkName, relation.foreignKey],
+          column => `Column '${column}' not found on table '${rootTable.name}'`
+        );
+        const keyRows = await fetchRowsForKeys(ctx, rootTable, pkColumn, missingKeys, selection);
+        for (const row of keyRows) {
+          const pkValue = row[pkName];
+          if (pkValue === undefined || pkValue === null) continue;
+          const entity = entityByPk.get(pkValue);
+          if (!entity) continue;
+          const fkValue = row[relation.foreignKey];
+          if (fkValue !== undefined && fkValue !== null) {
+            entity[relation.foreignKey] = fkValue;
+          }
+        }
+        foreignKeys = getForeignKeys();
+      }
+    }
+  }
+
+  if (!foreignKeys.size) {
+    return new Map();
+  }
+
+  const targetKey = relation.localKey || findPrimaryKey(relation.target);
   const pkColumn = relation.target.columns[targetKey];
   if (!pkColumn) return new Map();
 
-  const rows = await fetchRowsForKeys(ctx, relation.target, pkColumn, foreignKeys);
-  return groupRowsByUnique(rows, targetKey);
+  const requestedColumns = hasColumns(options?.columns) ? [...options!.columns] : undefined;
+  const selectedColumns = requestedColumns ? [...requestedColumns] : Object.keys(relation.target.columns);
+  if (!selectedColumns.includes(targetKey)) {
+    selectedColumns.push(targetKey);
+  }
+
+  const selection = buildColumnSelection(
+    relation.target,
+    selectedColumns,
+    column => `Column '${column}' not found on relation '${relationName}'`
+  );
+
+  const rows = await fetchRowsForKeys(ctx, relation.target, pkColumn, foreignKeys, selection, options?.filter);
+  const grouped = groupRowsByUnique(rows, targetKey);
+
+  if (!requestedColumns) return grouped;
+
+  const visibleColumns = new Set(selectedColumns);
+  const filtered = new Map<string, Record<string, unknown>>();
+  for (const [key, row] of grouped.entries()) {
+    filtered.set(key, filterRow(row, visibleColumns));
+  }
+  return filtered;
 };
 
 /**
@@ -247,8 +391,9 @@ export const loadBelongsToRelation = async (
 export const loadBelongsToManyRelation = async (
   ctx: EntityContext,
   rootTable: TableDef,
-  _relationName: string,
-  relation: BelongsToManyRelation
+  relationName: string,
+  relation: BelongsToManyRelation,
+  options?: RelationIncludeOptions
 ): Promise<Map<string, Rows>> => {
   const rootKey = relation.localKey || findPrimaryKey(rootTable);
   const roots = ctx.getEntitiesForTable(rootTable);
@@ -261,9 +406,32 @@ export const loadBelongsToManyRelation = async (
   const pivotColumn = relation.pivotTable.columns[relation.pivotForeignKeyToRoot];
   if (!pivotColumn) return new Map();
 
-  const pivotRows = await fetchRowsForKeys(ctx, relation.pivotTable, pivotColumn, rootIds);
+  const pivotColumnsRequested = hasColumns(options?.pivot?.columns) ? [...options!.pivot!.columns] : undefined;
+  const useIncludeDefaults = options !== undefined;
+  let pivotSelectedColumns: string[];
+  if (pivotColumnsRequested) {
+    pivotSelectedColumns = [...pivotColumnsRequested];
+  } else if (useIncludeDefaults) {
+    const pivotPk = relation.pivotPrimaryKey || findPrimaryKey(relation.pivotTable);
+    pivotSelectedColumns = relation.defaultPivotColumns ?? buildDefaultPivotColumns(relation, pivotPk);
+  } else {
+    pivotSelectedColumns = Object.keys(relation.pivotTable.columns);
+  }
+
+  const pivotQueryColumns = new Set(pivotSelectedColumns);
+  pivotQueryColumns.add(relation.pivotForeignKeyToRoot);
+  pivotQueryColumns.add(relation.pivotForeignKeyToTarget);
+
+  const pivotSelection = buildColumnSelection(
+    relation.pivotTable,
+    Array.from(pivotQueryColumns),
+    column => `Column '${column}' not found on pivot table '${relation.pivotTable.name}'`
+  );
+
+  const pivotRows = await fetchRowsForKeys(ctx, relation.pivotTable, pivotColumn, rootIds, pivotSelection);
   const rootLookup = new Map<string, { targetId: unknown; pivot: Record<string, unknown> }[]>();
   const targetIds = new Set<unknown>();
+  const pivotVisibleColumns = new Set(pivotSelectedColumns);
 
   for (const pivot of pivotRows) {
     const rootValue = pivot[relation.pivotForeignKeyToRoot];
@@ -274,7 +442,7 @@ export const loadBelongsToManyRelation = async (
     const bucket = rootLookup.get(toKey(rootValue)) ?? [];
     bucket.push({
       targetId: targetValue,
-      pivot: { ...pivot }
+      pivot: pivotVisibleColumns.size ? filterRow(pivot, pivotVisibleColumns) : {}
     });
     rootLookup.set(toKey(rootValue), bucket);
     targetIds.add(targetValue);
@@ -288,8 +456,23 @@ export const loadBelongsToManyRelation = async (
   const targetPkColumn = relation.target.columns[targetKey];
   if (!targetPkColumn) return new Map();
 
-  const targetRows = await fetchRowsForKeys(ctx, relation.target, targetPkColumn, targetIds);
+  const targetRequestedColumns = hasColumns(options?.columns) ? [...options!.columns] : undefined;
+  const targetSelectedColumns = targetRequestedColumns
+    ? [...targetRequestedColumns]
+    : Object.keys(relation.target.columns);
+  if (!targetSelectedColumns.includes(targetKey)) {
+    targetSelectedColumns.push(targetKey);
+  }
+
+  const targetSelection = buildColumnSelection(
+    relation.target,
+    targetSelectedColumns,
+    column => `Column '${column}' not found on relation '${relationName}'`
+  );
+
+  const targetRows = await fetchRowsForKeys(ctx, relation.target, targetPkColumn, targetIds, targetSelection, options?.filter);
   const targetMap = groupRowsByUnique(targetRows, targetKey);
+  const targetVisibleColumns = new Set(targetSelectedColumns);
   const result = new Map<string, Rows>();
 
   for (const [rootId, entries] of rootLookup.entries()) {
@@ -298,7 +481,7 @@ export const loadBelongsToManyRelation = async (
       const targetRow = targetMap.get(toKey(entry.targetId));
       if (!targetRow) continue;
       bucket.push({
-        ...targetRow,
+        ...(targetRequestedColumns ? filterRow(targetRow, targetVisibleColumns) : targetRow),
         _pivot: entry.pivot
       });
     }

package/src/orm/relations/belongs-to.ts

@@ -1,4 +1,4 @@
-import { BelongsToReference } from '../../schema/types.js';
+import { BelongsToReferenceApi } from '../../schema/types.js';
 import { EntityContext } from '../entity-context.js';
 import { RelationKey } from '../runtime-types.js';
 import { BelongsToRelation } from '../../schema/relation.js';
@@ -26,7 +26,7 @@ const hideInternal = (obj: object, keys: string[]): void => {
  *
  * @template TParent The type of the parent entity.
  */
-export class DefaultBelongsToReference<TParent> implements BelongsToReference<TParent> {
+export class DefaultBelongsToReference<TParent extends object> implements BelongsToReferenceApi<TParent> {
   private loaded = false;
   private current: TParent | null = null;
 

package/src/orm/relations/has-many.ts

@@ -69,15 +69,29 @@ export class DefaultHasManyCollection<TChild> implements HasManyCollection<TChil
     this.items = rows.map(row => this.createEntity(row));
     this.loaded = true;
     return this.items;
-  }
-
-  /**
-   * Gets the current items in the collection.
-   * @returns Array of child entities
-   */
-  getItems(): TChild[] {
-    return this.items;
-  }
+  }
+
+  /**
+   * Gets the current items in the collection.
+   * @returns Array of child entities
+   */
+  getItems(): TChild[] {
+    return this.items;
+  }
+ 
+  /**
+   * Array-compatible length for testing frameworks.
+   */
+  get length(): number {
+    return this.items.length;
+  }
+
+  /**
+   * Enables iteration over the collection like an array.
+   */
+  [Symbol.iterator](): Iterator<TChild> {
+    return this.items[Symbol.iterator]();
+  }
 
   /**
    * Adds a new child entity to the collection.

package/src/orm/relations/has-one.ts

@@ -1,4 +1,4 @@
-import { HasOneReference } from '../../schema/types.js';
+import { HasOneReferenceApi } from '../../schema/types.js';
 import { EntityContext } from '../entity-context.js';
 import { RelationKey } from '../runtime-types.js';
 import { HasOneRelation } from '../../schema/relation.js';
@@ -26,7 +26,7 @@ const hideInternal = (obj: object, keys: string[]): void => {
  *
  * @template TChild The type of the child entity.
  */
-export class DefaultHasOneReference<TChild> implements HasOneReference<TChild> {
+export class DefaultHasOneReference<TChild extends object> implements HasOneReferenceApi<TChild> {
   private loaded = false;
   private current: TChild | null = null;
 

package/src/orm/relations/many-to-many.ts

@@ -78,19 +78,33 @@ export class DefaultManyToManyCollection<TTarget> implements ManyToManyCollectio
     return this.items;
   }
 
-  /**
-   * Returns the currently loaded items.
-   * @returns Array of target entities.
-   */
-  getItems(): TTarget[] {
-    return this.items;
-  }
-
-  /**
-   * Attaches an entity to the collection.
-   * Registers an 'attach' change in the entity context.
-   * @param target Entity instance or its primary key value.
-   */
+  /**
+   * Returns the currently loaded items.
+   * @returns Array of target entities.
+   */
+  getItems(): TTarget[] {
+    return this.items;
+  }
+
+  /**
+   * Array-compatible length for testing frameworks.
+   */
+  get length(): number {
+    return this.items.length;
+  }
+
+  /**
+   * Enables iteration over the collection like an array.
+   */
+  [Symbol.iterator](): Iterator<TTarget> {
+    return this.items[Symbol.iterator]();
+  }
+
+  /**
+   * Attaches an entity to the collection.
+   * Registers an 'attach' change in the entity context.
+   * @param target Entity instance or its primary key value.
+   */
   attach(target: TTarget | number | string): void {
     const entity = this.ensureEntity(target);
     const id = this.extractId(entity);

package/src/orm/save-graph-types.ts

@@ -10,8 +10,8 @@ type AnyFn = (...args: unknown[]) => unknown;
 
 type RelationWrapper =
   | HasManyCollection<unknown>
-  | HasOneReference<unknown>
-  | BelongsToReference<unknown>
+  | HasOneReference
+  | BelongsToReference
   | ManyToManyCollection<unknown>;
 
 type FunctionKeys<T> = {

package/src/orm/save-graph.ts

@@ -184,15 +184,15 @@ const handleHasMany = async (
   }
 };
 
-const handleHasOne = async (
-  session: OrmSession,
-  root: AnyEntity,
-  relationName: string,
-  relation: HasOneRelation,
-  payload: unknown,
-  options: SaveGraphOptions
-): Promise<void> => {
-  const ref = root[relationName] as unknown as HasOneReference<unknown>;
+const handleHasOne = async (
+  session: OrmSession,
+  root: AnyEntity,
+  relationName: string,
+  relation: HasOneRelation,
+  payload: unknown,
+  options: SaveGraphOptions
+): Promise<void> => {
+  const ref = root[relationName] as unknown as HasOneReference<object>;
   if (payload === undefined) return;
   if (payload === null) {
     ref.set(null);
@@ -212,15 +212,15 @@ const handleHasOne = async (
   }
 };
 
-const handleBelongsTo = async (
-  session: OrmSession,
-  root: AnyEntity,
-  relationName: string,
-  relation: BelongsToRelation,
-  payload: unknown,
-  options: SaveGraphOptions
-): Promise<void> => {
-  const ref = root[relationName] as unknown as BelongsToReference<unknown>;
+const handleBelongsTo = async (
+  session: OrmSession,
+  root: AnyEntity,
+  relationName: string,
+  relation: BelongsToRelation,
+  payload: unknown,
+  options: SaveGraphOptions
+): Promise<void> => {
+  const ref = root[relationName] as unknown as BelongsToReference<object>;
   if (payload === undefined) return;
   if (payload === null) {
     ref.set(null);