Refactor select improving spread (...obj) and enabling nested projection. (#389)

Co-authored-by: Kyle Mathews <mathews.kyle@gmail.com>
Co-authored-by: Claude <noreply@anthropic.com>

Author: 3 people

.changeset/quick-drinks-talk.md

@@ -0,0 +1,5 @@
+---
+"@tanstack/db": patch
+---
+
+Refactored `select` improving spread (`...obj`) support and enabling nested projection.

packages/db/src/query/builder/functions.ts

@@ -2,9 +2,12 @@ import { Aggregate, Func } from "../ir"
 import { toExpression } from "./ref-proxy.js"
 import type { BasicExpression } from "../ir"
 import type { RefProxy } from "./ref-proxy.js"
-import type { Ref } from "./types.js"
+import type { RefLeaf } from "./types.js"
 
-type StringRef = Ref<string> | Ref<string | null> | Ref<string | undefined>
+type StringRef =
+  | RefLeaf<string>
+  | RefLeaf<string | null>
+  | RefLeaf<string | undefined>
 type StringRefProxy =
   | RefProxy<string>
   | RefProxy<string | null>
@@ -23,7 +26,7 @@ type StringLike =
 
 type ComparisonOperand<T> =
   | RefProxy<T>
-  | Ref<T>
+  | RefLeaf<T>
   | T
   | BasicExpression<T>
   | undefined
@@ -35,13 +38,13 @@ type ComparisonOperandPrimitive<T extends string | number | boolean> =
   | null
 
 // Helper type for any expression-like value
-type ExpressionLike = BasicExpression | RefProxy<any> | Ref<any> | any
+type ExpressionLike = BasicExpression | RefProxy<any> | RefLeaf<any> | any
 
 // Helper type to extract the underlying type from various expression types
 type ExtractType<T> =
   T extends RefProxy<infer U>
     ? U
-    : T extends Ref<infer U>
+    : T extends RefLeaf<infer U>
       ? U
       : T extends BasicExpression<infer U>
         ? U

packages/db/src/query/builder/index.ts

@@ -1,13 +1,21 @@
 import { CollectionImpl } from "../../collection.js"
-import { CollectionRef, QueryRef } from "../ir.js"
+import {
+  Aggregate as AggregateExpr,
+  CollectionRef,
+  Func as FuncExpr,
+  PropRef,
+  QueryRef,
+  Value as ValueExpr,
+  isExpressionLike,
+} from "../ir.js"
 import {
   InvalidSourceError,
   JoinConditionMustBeEqualityError,
   OnlyOneSourceAllowedError,
   QueryMustHaveFromClauseError,
   SubQueryMustHaveFromClauseError,
 } from "../../errors.js"
-import { createRefProxy, isRefProxy, toExpression } from "./ref-proxy.js"
+import { createRefProxy, toExpression } from "./ref-proxy.js"
 import type { NamespacedRow } from "../../types.js"
 import type {
   Aggregate,
@@ -26,7 +34,7 @@ import type {
   MergeContextWithJoinType,
   OrderByCallback,
   OrderByOptions,
-  RefProxyForContext,
+  RefsForContext,
   ResultTypeFromSelect,
   SchemaFromSource,
   SelectObject,
@@ -152,7 +160,7 @@ export class BaseQueryBuilder<TContext extends Context = Context> {
     // Create a temporary context for the callback
     const currentAliases = this._getCurrentAliases()
     const newAliases = [...currentAliases, alias]
-    const refProxy = createRefProxy(newAliases) as RefProxyForContext<
+    const refProxy = createRefProxy(newAliases) as RefsForContext<
       MergeContextForJoinCallback<TContext, SchemaFromSource<TSource>>
     >
 
@@ -324,7 +332,7 @@ export class BaseQueryBuilder<TContext extends Context = Context> {
    */
   where(callback: WhereCallback<TContext>): QueryBuilder<TContext> {
     const aliases = this._getCurrentAliases()
-    const refProxy = createRefProxy(aliases) as RefProxyForContext<TContext>
+    const refProxy = createRefProxy(aliases) as RefsForContext<TContext>
     const expression = callback(refProxy)
 
     const existingWhere = this.query.where || []
@@ -365,7 +373,7 @@ export class BaseQueryBuilder<TContext extends Context = Context> {
    */
   having(callback: WhereCallback<TContext>): QueryBuilder<TContext> {
     const aliases = this._getCurrentAliases()
-    const refProxy = createRefProxy(aliases) as RefProxyForContext<TContext>
+    const refProxy = createRefProxy(aliases) as RefsForContext<TContext>
     const expression = callback(refProxy)
 
     const existingHaving = this.query.having || []
@@ -411,43 +419,16 @@ export class BaseQueryBuilder<TContext extends Context = Context> {
    * ```
    */
   select<TSelectObject extends SelectObject>(
-    callback: (refs: RefProxyForContext<TContext>) => TSelectObject
+    callback: (refs: RefsForContext<TContext>) => TSelectObject
   ): QueryBuilder<WithResult<TContext, ResultTypeFromSelect<TSelectObject>>> {
     const aliases = this._getCurrentAliases()
-    const refProxy = createRefProxy(aliases) as RefProxyForContext<TContext>
+    const refProxy = createRefProxy(aliases) as RefsForContext<TContext>
     const selectObject = callback(refProxy)
-
-    // Check if any tables were spread during the callback
-    const spreadSentinels = (refProxy as any).__spreadSentinels as Set<string>
-
-    // Convert the select object to use expressions, including spread sentinels
-    const select: Record<string, BasicExpression | Aggregate> = {}
-
-    // First, add spread sentinels for any tables that were spread
-    for (const spreadAlias of spreadSentinels) {
-      const sentinelKey = `__SPREAD_SENTINEL__${spreadAlias}`
-      select[sentinelKey] = toExpression(spreadAlias) // Use alias as a simple reference
-    }
-
-    // Then add the explicit select fields
-    for (const [key, value] of Object.entries(selectObject)) {
-      if (isRefProxy(value)) {
-        select[key] = toExpression(value)
-      } else if (
-        typeof value === `object` &&
-        value !== null &&
-        `type` in value &&
-        (value.type === `agg` || value.type === `func`)
-      ) {
-        select[key] = value as BasicExpression | Aggregate
-      } else {
-        select[key] = toExpression(value)
-      }
-    }
+    const select = buildNestedSelect(selectObject)
 
     return new BaseQueryBuilder({
       ...this.query,
-      select,
+      select: select,
       fnSelect: undefined, // remove the fnSelect clause if it exists
     }) as any
   }
@@ -483,7 +464,7 @@ export class BaseQueryBuilder<TContext extends Context = Context> {
     options: OrderByDirection | OrderByOptions = `asc`
   ): QueryBuilder<TContext> {
     const aliases = this._getCurrentAliases()
-    const refProxy = createRefProxy(aliases) as RefProxyForContext<TContext>
+    const refProxy = createRefProxy(aliases) as RefsForContext<TContext>
     const result = callback(refProxy)
 
     const opts: CompareOptions =
@@ -551,7 +532,7 @@ export class BaseQueryBuilder<TContext extends Context = Context> {
    */
   groupBy(callback: GroupByCallback<TContext>): QueryBuilder<TContext> {
     const aliases = this._getCurrentAliases()
-    const refProxy = createRefProxy(aliases) as RefProxyForContext<TContext>
+    const refProxy = createRefProxy(aliases) as RefsForContext<TContext>
     const result = callback(refProxy)
 
     const newExpressions = Array.isArray(result)
@@ -760,6 +741,43 @@ export class BaseQueryBuilder<TContext extends Context = Context> {
   }
 }
 
+// Helper to ensure we have a BasicExpression/Aggregate for a value
+function toExpr(value: any): BasicExpression | Aggregate {
+  if (value === undefined) return toExpression(null)
+  if (
+    value instanceof AggregateExpr ||
+    value instanceof FuncExpr ||
+    value instanceof PropRef ||
+    value instanceof ValueExpr
+  ) {
+    return value as BasicExpression | Aggregate
+  }
+  return toExpression(value)
+}
+
+function isPlainObject(value: any): value is Record<string, any> {
+  return (
+    value !== null &&
+    typeof value === `object` &&
+    !isExpressionLike(value) &&
+    !value.__refProxy
+  )
+}
+
+function buildNestedSelect(obj: any): any {
+  if (!isPlainObject(obj)) return toExpr(obj)
+  const out: Record<string, any> = {}
+  for (const [k, v] of Object.entries(obj)) {
+    if (typeof k === `string` && k.startsWith(`__SPREAD_SENTINEL__`)) {
+      // Preserve sentinel key and its value (value is unimportant at compile time)
+      out[k] = v
+      continue
+    }
+    out[k] = buildNestedSelect(v)
+  }
+  return out
+}
+
 // Internal function to build a query from a callback
 // used by liveQueryCollectionOptions.query
 export function buildQuery<TContext extends Context>(
@@ -799,4 +817,4 @@ export type ExtractContext<T> =
       : never
 
 // Export the types from types.ts for convenience
-export type { Context, Source, GetResult, Ref } from "./types.js"
+export type { Context, Source, GetResult, RefLeaf as Ref } from "./types.js"

packages/db/src/query/builder/ref-proxy.ts

@@ -1,6 +1,6 @@
 import { PropRef, Value } from "../ir.js"
 import type { BasicExpression } from "../ir.js"
-import type { Ref } from "./types.js"
+import type { RefLeaf } from "./types.js"
 
 export interface RefProxy<T = any> {
   /** @internal */
@@ -20,7 +20,7 @@ export type SingleRowRefProxy<T> =
     ? {
         [K in keyof T]: T[K] extends Record<string, any>
           ? SingleRowRefProxy<T[K]> & RefProxy<T[K]>
-          : Ref<T[K]>
+          : RefLeaf<T[K]>
       } & RefProxy<T>
     : RefProxy<T>
 
@@ -84,7 +84,7 @@ export function createRefProxy<T extends Record<string, any>>(
   aliases: Array<string>
 ): RefProxy<T> & T {
   const cache = new Map<string, any>()
-  const spreadSentinels = new Set<string>() // Track which aliases have been spread
+  let accessId = 0 // Monotonic counter to record evaluation order
 
   function createProxy(path: Array<string>): any {
     const pathKey = path.join(`.`)
@@ -110,10 +110,14 @@ export function createRefProxy<T extends Record<string, any>>(
       },
 
       ownKeys(target) {
-        // If this is a table-level proxy (path length 1), mark it as spread
-        if (path.length === 1) {
-          const aliasName = path[0]!
-          spreadSentinels.add(aliasName)
+        const id = ++accessId
+        const sentinelKey = `__SPREAD_SENTINEL__${path.join(`.`)}__${id}`
+        if (!Object.prototype.hasOwnProperty.call(target, sentinelKey)) {
+          Object.defineProperty(target, sentinelKey, {
+            enumerable: true,
+            configurable: true,
+            value: true,
+          })
         }
         return Reflect.ownKeys(target)
       },
@@ -136,7 +140,6 @@ export function createRefProxy<T extends Record<string, any>>(
       if (prop === `__refProxy`) return true
       if (prop === `__path`) return []
       if (prop === `__type`) return undefined // Type is only for TypeScript inference
-      if (prop === `__spreadSentinels`) return spreadSentinels // Expose spread sentinels
       if (typeof prop === `symbol`) return Reflect.get(target, prop, receiver)
 
       const propStr = String(prop)
@@ -148,28 +151,18 @@ export function createRefProxy<T extends Record<string, any>>(
     },
 
     has(target, prop) {
-      if (
-        prop === `__refProxy` ||
-        prop === `__path` ||
-        prop === `__type` ||
-        prop === `__spreadSentinels`
-      )
+      if (prop === `__refProxy` || prop === `__path` || prop === `__type`)
         return true
       if (typeof prop === `string` && aliases.includes(prop)) return true
       return Reflect.has(target, prop)
     },
 
     ownKeys(_target) {
-      return [...aliases, `__refProxy`, `__path`, `__type`, `__spreadSentinels`]
+      return [...aliases, `__refProxy`, `__path`, `__type`]
     },
 
     getOwnPropertyDescriptor(target, prop) {
-      if (
-        prop === `__refProxy` ||
-        prop === `__path` ||
-        prop === `__type` ||
-        prop === `__spreadSentinels`
-      ) {
+      if (prop === `__refProxy` || prop === `__path` || prop === `__type`) {
         return { enumerable: false, configurable: true }
       }
       if (typeof prop === `string` && aliases.includes(prop)) {