metal-orm 1.0.62 → 1.0.64

package/src/query-builder/select.ts

@@ -2,18 +2,19 @@ import { TableDef } from '../schema/table.js';
 import { ColumnDef } from '../schema/column-types.js';
 import { OrderingTerm, SelectQueryNode, SetOperationKind } from '../core/ast/query.js';
 import { HydrationPlan } from '../core/hydration/types.js';
-import {
-  ColumnNode,
-  ExpressionNode,
-  FunctionNode,
-  BinaryExpressionNode,
-  CaseExpressionNode,
-  WindowFunctionNode,
-  and,
-  exists,
-  notExists,
-  OperandNode
-} from '../core/ast/expression.js';
+import {
+  ColumnNode,
+  ExpressionNode,
+  FunctionNode,
+  BinaryExpressionNode,
+  CaseExpressionNode,
+  WindowFunctionNode,
+  and,
+  exists,
+  notExists,
+  OperandNode
+} from '../core/ast/expression.js';
+import type { TypedExpression } from '../core/ast/expression.js';
 import { CompiledQuery, Dialect } from '../core/dialect/abstract.js';
 import { DialectKey, resolveDialectInput } from '../core/dialect/dialect-factory.js';
 
@@ -31,7 +32,8 @@ import { ColumnSelector } from './column-selector.js';
 import { RelationIncludeOptions, RelationTargetColumns, TypedRelationIncludeOptions } from './relation-types.js';
 import { RelationKinds } from '../schema/relation.js';
 import { JOIN_KINDS, JoinKind, ORDER_DIRECTIONS, OrderDirection } from '../core/sql/sql.js';
-import { EntityInstance, RelationMap } from '../schema/types.js';
+import { EntityInstance, RelationMap } from '../schema/types.js';
+import type { ColumnToTs, InferRow } from '../schema/types.js';
 import { OrmSession } from '../orm/orm-session.ts';
 import { ExecutionContext } from '../orm/execution-context.js';
 import { HydrationContext } from '../orm/hydration-context.js';
@@ -55,8 +57,27 @@ import { SelectCTEFacet } from './select/cte-facet.js';
 import { SelectSetOpFacet } from './select/setop-facet.js';
 import { SelectRelationFacet } from './select/relation-facet.js';
 
-type ColumnSelectionValue = ColumnDef | FunctionNode | CaseExpressionNode | WindowFunctionNode;
-
+type ColumnSelectionValue =
+  | ColumnDef
+  | FunctionNode
+  | CaseExpressionNode
+  | WindowFunctionNode
+  | TypedExpression<unknown>;
+
+type SelectionValueType<TValue> =
+  TValue extends TypedExpression<infer TRuntime> ? TRuntime :
+  TValue extends ColumnDef ? ColumnToTs<TValue> :
+  unknown;
+
+type SelectionResult<TSelection extends Record<string, ColumnSelectionValue>> = {
+  [K in keyof TSelection]: SelectionValueType<TSelection[K]>;
+};
+
+type SelectionFromKeys<
+  TTable extends TableDef,
+  K extends keyof TTable['columns'] & string
+> = Pick<InferRow<TTable>, K>;
+
 type DeepSelectEntry<TTable extends TableDef> = {
   type: 'root';
   columns: (keyof TTable['columns'] & string)[];
@@ -133,21 +154,21 @@ export class SelectQueryBuilder<T = EntityInstance<TableDef>, TTable extends Tab
    * @param lazyRelations - Updated lazy relations set
    * @returns New SelectQueryBuilder instance
    */
-  private clone(
-    context: SelectQueryBuilderContext = this.context,
-    lazyRelations = new Set(this.lazyRelations),
-    lazyRelationOptions = new Map(this.lazyRelationOptions)
-  ): SelectQueryBuilder<T, TTable> {
-    return new SelectQueryBuilder(
-      this.env.table as TTable,
-      context.state,
-      context.hydration,
-      this.env.deps,
-      lazyRelations,
-      lazyRelationOptions,
-      this.entityConstructor
-    );
-  }
+  private clone<TNext = T>(
+    context: SelectQueryBuilderContext = this.context,
+    lazyRelations = new Set(this.lazyRelations),
+    lazyRelationOptions = new Map(this.lazyRelationOptions)
+  ): SelectQueryBuilder<TNext, TTable> {
+    return new SelectQueryBuilder(
+      this.env.table as TTable,
+      context.state,
+      context.hydration,
+      this.env.deps,
+      lazyRelations,
+      lazyRelationOptions,
+      this.entityConstructor
+    ) as SelectQueryBuilder<TNext, TTable>;
+  }
 
   /**
    * Applies an alias to the root FROM table.
@@ -213,32 +234,41 @@ export class SelectQueryBuilder<T = EntityInstance<TableDef>, TTable extends Tab
    *   fullName: concat(userTable.columns.firstName, ' ', userTable.columns.lastName)
    * });
    */
-  select<K extends keyof TTable['columns'] & string>(
-    ...args: K[]
-  ): SelectQueryBuilder<T, TTable>;
-  select(columns: Record<string, ColumnSelectionValue>): SelectQueryBuilder<T, TTable>;
-  select<K extends keyof TTable['columns'] & string>(
-    ...args: K[] | [Record<string, ColumnSelectionValue>]
-  ): SelectQueryBuilder<T, TTable> {
+  select<K extends keyof TTable['columns'] & string>(
+    ...args: K[]
+  ): SelectQueryBuilder<T & SelectionFromKeys<TTable, K>, TTable>;
+  select<TSelection extends Record<string, ColumnSelectionValue>>(
+    columns: TSelection
+  ): SelectQueryBuilder<T & SelectionResult<TSelection>, TTable>;
+  select<
+    K extends keyof TTable['columns'] & string,
+    TSelection extends Record<string, ColumnSelectionValue>
+  >(
+    ...args: K[] | [TSelection]
+  ): SelectQueryBuilder<T, TTable> {
     // If first arg is an object (not a string), treat as projection map
     if (args.length === 1 && typeof args[0] === 'object' && args[0] !== null && typeof args[0] !== 'string') {
-      const columns = args[0] as Record<string, ColumnSelectionValue>;
-      return this.clone(this.projectionFacet.select(this.context, columns));
-    }
+      const columns = args[0] as TSelection;
+      return this.clone<T & SelectionResult<TSelection>>(
+        this.projectionFacet.select(this.context, columns)
+      );
+    }
 
     // Otherwise, treat as column names
     const cols = args as K[];
     const selection: Record<string, ColumnDef> = {};
-    for (const key of cols) {
-      const col = this.env.table.columns[key];
-      if (!col) {
-        throw new Error(`Column '${key}' not found on table '${this.env.table.name}'`);
-      }
-      selection[key] = col;
-    }
-
-    return this.clone(this.projectionFacet.select(this.context, selection));
-  }
+    for (const key of cols) {
+      const col = this.env.table.columns[key];
+      if (!col) {
+        throw new Error(`Column '${key}' not found on table '${this.env.table.name}'`);
+      }
+      selection[key] = col;
+    }
+
+    return this.clone<T & SelectionFromKeys<TTable, K>>(
+      this.projectionFacet.select(this.context, selection)
+    );
+  }
 
   /**
    * Selects raw column expressions
@@ -354,10 +384,15 @@ export class SelectQueryBuilder<T = EntityInstance<TableDef>, TTable extends Tab
    * qb.select('id', 'name')
    *   .selectSubquery('postCount', postCount);
    */
-  selectSubquery<TSub extends TableDef>(alias: string, sub: SelectQueryBuilder<unknown, TSub> | SelectQueryNode): SelectQueryBuilder<T, TTable> {
-    const query = resolveSelectQuery(sub);
-    return this.clone(this.projectionFacet.selectSubquery(this.context, alias, query));
-  }
+  selectSubquery<TValue = unknown, K extends string = string, TSub extends TableDef = TableDef>(
+    alias: K,
+    sub: SelectQueryBuilder<unknown, TSub> | SelectQueryNode
+  ): SelectQueryBuilder<T & Record<K, TValue>, TTable> {
+    const query = resolveSelectQuery(sub);
+    return this.clone<T & Record<K, TValue>>(
+      this.projectionFacet.selectSubquery(this.context, alias, query)
+    );
+  }
 
   /**
    * Adds a JOIN against a derived table (subquery with alias)
@@ -675,7 +710,8 @@ export class SelectQueryBuilder<T = EntityInstance<TableDef>, TTable extends Tab
    */
   async executePlain(ctx: OrmSession): Promise<EntityInstance<TTable>[]> {
     const builder = this.ensureDefaultSelection();
-    return executeHydratedPlain(ctx, builder) as EntityInstance<TTable>[];
+    const rows = await executeHydratedPlain(ctx, builder);
+    return rows as EntityInstance<TTable>[];
   }
 
   /**