@tanstack/db 0.5.33 → 0.6.1

package/src/query/builder/index.ts

@@ -3,6 +3,8 @@ import {
   Aggregate as AggregateExpr,
   CollectionRef,
   Func as FuncExpr,
+  INCLUDES_SCALAR_FIELD,
+  IncludesSubquery,
   PropRef,
   QueryRef,
   Value as ValueExpr,
@@ -23,14 +25,17 @@ import {
   isRefProxy,
   toExpression,
 } from './ref-proxy.js'
+import { ConcatToArrayWrapper, ToArrayWrapper } from './functions.js'
 import type { NamespacedRow, SingleResult } from '../../types.js'
 import type {
   Aggregate,
   BasicExpression,
+  IncludesMaterialization,
   JoinClause,
   OrderBy,
   OrderByDirection,
   QueryIR,
+  Where,
 } from '../ir.js'
 import type {
   CompareOptions,
@@ -41,10 +46,13 @@ import type {
   JoinOnCallback,
   MergeContextForJoinCallback,
   MergeContextWithJoinType,
+  NonScalarSelectObject,
   OrderByCallback,
   OrderByOptions,
   RefsForContext,
   ResultTypeFromSelect,
+  ResultTypeFromSelectValue,
+  ScalarSelectValue,
   SchemaFromSource,
   SelectObject,
   Source,
@@ -486,12 +494,30 @@ export class BaseQueryBuilder<TContext extends Context = Context> {
    * ```
    */
   select<TSelectObject extends SelectObject>(
-    callback: (refs: RefsForContext<TContext>) => TSelectObject,
-  ): QueryBuilder<WithResult<TContext, ResultTypeFromSelect<TSelectObject>>> {
+    callback: (
+      refs: RefsForContext<TContext>,
+    ) => NonScalarSelectObject<TSelectObject>,
+  ): QueryBuilder<WithResult<TContext, ResultTypeFromSelect<TSelectObject>>>
+  select<TSelectValue extends ScalarSelectValue>(
+    callback: (refs: RefsForContext<TContext>) => TSelectValue,
+  ): QueryBuilder<WithResult<TContext, ResultTypeFromSelectValue<TSelectValue>>>
+  select(
+    callback: (
+      refs: RefsForContext<TContext>,
+    ) => SelectObject | ScalarSelectValue,
+  ) {
     const aliases = this._getCurrentAliases()
     const refProxy = createRefProxy(aliases) as RefsForContext<TContext>
-    const selectObject = callback(refProxy)
-    const select = buildNestedSelect(selectObject)
+    let selectObject = callback(refProxy)
+
+    // Returning a top-level alias directly is equivalent to spreading it.
+    // Leaf refs like `row.name` must remain scalar selections.
+    if (isRefProxy(selectObject) && selectObject.__path.length === 1) {
+      const sentinelKey = `__SPREAD_SENTINEL__${selectObject.__path[0]}__0`
+      selectObject = { [sentinelKey]: true }
+    }
+
+    const select = buildNestedSelect(selectObject, aliases)
 
     return new BaseQueryBuilder({
       ...this.query,
@@ -676,7 +702,7 @@ export class BaseQueryBuilder<TContext extends Context = Context> {
    * // Get countries our users are from
    * query
    *   .from({ users: usersCollection })
-   *   .select(({users}) => users.country)
+   *   .select(({users}) => ({ country: users.country }))
    *   .distinct()
    * ```
    */
@@ -706,7 +732,7 @@ export class BaseQueryBuilder<TContext extends Context = Context> {
       // TODO: enforcing return only one result with also a default orderBy if none is specified
       // limit: 1,
       singleResult: true,
-    })
+    }) as any
   }
 
   // Helper methods
@@ -769,7 +795,7 @@ export class BaseQueryBuilder<TContext extends Context = Context> {
           ...builder.query,
           select: undefined, // remove the select clause if it exists
           fnSelect: callback,
-        })
+        }) as any
       },
       /**
        * Filter rows using a function that operates on each row
@@ -867,7 +893,7 @@ function isPlainObject(value: any): value is Record<string, any> {
   )
 }
 
-function buildNestedSelect(obj: any): any {
+function buildNestedSelect(obj: any, parentAliases: Array<string> = []): any {
   if (!isPlainObject(obj)) return toExpr(obj)
   const out: Record<string, any> = {}
   for (const [k, v] of Object.entries(obj)) {
@@ -876,11 +902,305 @@ function buildNestedSelect(obj: any): any {
       out[k] = v
       continue
     }
-    out[k] = buildNestedSelect(v)
+    if (v instanceof BaseQueryBuilder) {
+      out[k] = buildIncludesSubquery(v, k, parentAliases, `collection`)
+      continue
+    }
+    if (v instanceof ToArrayWrapper) {
+      if (!(v.query instanceof BaseQueryBuilder)) {
+        throw new Error(`toArray() must wrap a subquery builder`)
+      }
+      out[k] = buildIncludesSubquery(v.query, k, parentAliases, `array`)
+      continue
+    }
+    if (v instanceof ConcatToArrayWrapper) {
+      if (!(v.query instanceof BaseQueryBuilder)) {
+        throw new Error(`concat(toArray(...)) must wrap a subquery builder`)
+      }
+      out[k] = buildIncludesSubquery(v.query, k, parentAliases, `concat`)
+      continue
+    }
+    out[k] = buildNestedSelect(v, parentAliases)
   }
   return out
 }
 
+/**
+ * Recursively collects all PropRef nodes from an expression tree.
+ */
+function collectRefsFromExpression(expr: BasicExpression): Array<PropRef> {
+  const refs: Array<PropRef> = []
+  switch (expr.type) {
+    case `ref`:
+      refs.push(expr)
+      break
+    case `func`:
+      for (const arg of (expr as any).args ?? []) {
+        refs.push(...collectRefsFromExpression(arg))
+      }
+      break
+    default:
+      break
+  }
+  return refs
+}
+
+/**
+ * Checks whether a WHERE clause references any parent alias.
+ */
+function referencesParent(where: Where, parentAliases: Array<string>): boolean {
+  const expr =
+    typeof where === `object` && `expression` in where
+      ? where.expression
+      : where
+  return collectRefsFromExpression(expr).some(
+    (ref) => ref.path[0] != null && parentAliases.includes(ref.path[0]),
+  )
+}
+
+/**
+ * Builds an IncludesSubquery IR node from a child query builder.
+ * Extracts the correlation condition from the child's WHERE clauses by finding
+ * an eq() predicate that references both a parent alias and a child alias.
+ */
+function buildIncludesSubquery(
+  childBuilder: BaseQueryBuilder,
+  fieldName: string,
+  parentAliases: Array<string>,
+  materialization: IncludesMaterialization,
+): IncludesSubquery {
+  const childQuery = childBuilder._getQuery()
+
+  // Collect child's own aliases
+  const childAliases: Array<string> = [childQuery.from.alias]
+  if (childQuery.join) {
+    for (const j of childQuery.join) {
+      childAliases.push(j.from.alias)
+    }
+  }
+
+  // Walk child's WHERE clauses to find the correlation condition.
+  // The correlation eq() may be a standalone WHERE or nested inside a top-level and().
+  let parentRef: PropRef | undefined
+  let childRef: PropRef | undefined
+  let correlationWhereIndex = -1
+  let correlationAndArgIndex = -1 // >= 0 when found inside an and()
+
+  if (childQuery.where) {
+    for (let i = 0; i < childQuery.where.length; i++) {
+      const where = childQuery.where[i]!
+      const expr =
+        typeof where === `object` && `expression` in where
+          ? where.expression
+          : where
+
+      // Try standalone eq()
+      if (
+        expr.type === `func` &&
+        expr.name === `eq` &&
+        expr.args.length === 2
+      ) {
+        const result = extractCorrelation(
+          expr.args[0]!,
+          expr.args[1]!,
+          parentAliases,
+          childAliases,
+        )
+        if (result) {
+          parentRef = result.parentRef
+          childRef = result.childRef
+          correlationWhereIndex = i
+          break
+        }
+      }
+
+      // Try inside top-level and()
+      if (
+        expr.type === `func` &&
+        expr.name === `and` &&
+        expr.args.length >= 2
+      ) {
+        for (let j = 0; j < expr.args.length; j++) {
+          const arg = expr.args[j]!
+          if (
+            arg.type === `func` &&
+            arg.name === `eq` &&
+            arg.args.length === 2
+          ) {
+            const result = extractCorrelation(
+              arg.args[0]!,
+              arg.args[1]!,
+              parentAliases,
+              childAliases,
+            )
+            if (result) {
+              parentRef = result.parentRef
+              childRef = result.childRef
+              correlationWhereIndex = i
+              correlationAndArgIndex = j
+              break
+            }
+          }
+        }
+        if (parentRef) break
+      }
+    }
+  }
+
+  if (!parentRef || !childRef || correlationWhereIndex === -1) {
+    throw new Error(
+      `Includes subquery for "${fieldName}" must have a WHERE clause with an eq() condition ` +
+        `that correlates a parent field with a child field. ` +
+        `Example: .where(({child}) => eq(child.parentId, parent.id))`,
+    )
+  }
+
+  // Remove the correlation eq() from the child query's WHERE clauses.
+  // If it was inside an and(), remove just that arg (collapsing the and() if needed).
+  const modifiedWhere = [...childQuery.where!]
+  if (correlationAndArgIndex >= 0) {
+    const where = modifiedWhere[correlationWhereIndex]!
+    const expr =
+      typeof where === `object` && `expression` in where
+        ? where.expression
+        : where
+    const remainingArgs = (expr as any).args.filter(
+      (_: any, idx: number) => idx !== correlationAndArgIndex,
+    )
+    if (remainingArgs.length === 1) {
+      // Collapse and() with single remaining arg to just that expression
+      const isResidual =
+        typeof where === `object` && `expression` in where && where.residual
+      modifiedWhere[correlationWhereIndex] = isResidual
+        ? { expression: remainingArgs[0], residual: true }
+        : remainingArgs[0]
+    } else {
+      // Rebuild and() without the extracted arg
+      const newAnd = new FuncExpr(`and`, remainingArgs)
+      const isResidual =
+        typeof where === `object` && `expression` in where && where.residual
+      modifiedWhere[correlationWhereIndex] = isResidual
+        ? { expression: newAnd, residual: true }
+        : newAnd
+    }
+  } else {
+    modifiedWhere.splice(correlationWhereIndex, 1)
+  }
+
+  // Separate remaining WHEREs into pure-child vs parent-referencing
+  const pureChildWhere: Array<Where> = []
+  const parentFilters: Array<Where> = []
+  for (const w of modifiedWhere) {
+    if (referencesParent(w, parentAliases)) {
+      parentFilters.push(w)
+    } else {
+      pureChildWhere.push(w)
+    }
+  }
+
+  // Collect distinct parent PropRefs from parent-referencing filters
+  let parentProjection: Array<PropRef> | undefined
+  if (parentFilters.length > 0) {
+    const seen = new Set<string>()
+    parentProjection = []
+    for (const w of parentFilters) {
+      const expr = typeof w === `object` && `expression` in w ? w.expression : w
+      for (const ref of collectRefsFromExpression(expr)) {
+        if (
+          ref.path[0] != null &&
+          parentAliases.includes(ref.path[0]) &&
+          !seen.has(ref.path.join(`.`))
+        ) {
+          seen.add(ref.path.join(`.`))
+          parentProjection.push(ref)
+        }
+      }
+    }
+  }
+
+  const modifiedQuery: QueryIR = {
+    ...childQuery,
+    where: pureChildWhere.length > 0 ? pureChildWhere : undefined,
+  }
+
+  const rawChildSelect = modifiedQuery.select as any
+  const hasObjectSelect =
+    rawChildSelect === undefined || isPlainObject(rawChildSelect)
+  let includesQuery = modifiedQuery
+  let scalarField: string | undefined
+
+  if (materialization === `concat`) {
+    if (rawChildSelect === undefined || hasObjectSelect) {
+      throw new Error(
+        `concat(toArray(...)) for "${fieldName}" requires the subquery to select a scalar value`,
+      )
+    }
+  }
+
+  if (!hasObjectSelect) {
+    if (materialization === `collection`) {
+      throw new Error(
+        `Includes subquery for "${fieldName}" must select an object when materializing as a Collection`,
+      )
+    }
+
+    scalarField = INCLUDES_SCALAR_FIELD
+    includesQuery = {
+      ...modifiedQuery,
+      select: {
+        [scalarField]: rawChildSelect,
+      },
+    }
+  }
+
+  return new IncludesSubquery(
+    includesQuery,
+    parentRef,
+    childRef,
+    fieldName,
+    parentFilters.length > 0 ? parentFilters : undefined,
+    parentProjection,
+    materialization,
+    scalarField,
+  )
+}
+
+/**
+ * Checks if two eq() arguments form a parent-child correlation.
+ * Returns the parent and child PropRefs if found, undefined otherwise.
+ */
+function extractCorrelation(
+  argA: BasicExpression,
+  argB: BasicExpression,
+  parentAliases: Array<string>,
+  childAliases: Array<string>,
+): { parentRef: PropRef; childRef: PropRef } | undefined {
+  if (argA.type === `ref` && argB.type === `ref`) {
+    const aAlias = argA.path[0]
+    const bAlias = argB.path[0]
+
+    if (
+      aAlias &&
+      bAlias &&
+      parentAliases.includes(aAlias) &&
+      childAliases.includes(bAlias)
+    ) {
+      return { parentRef: argA, childRef: argB }
+    }
+
+    if (
+      aAlias &&
+      bAlias &&
+      parentAliases.includes(bAlias) &&
+      childAliases.includes(aAlias)
+    ) {
+      return { parentRef: argB, childRef: argA }
+    }
+  }
+
+  return undefined
+}
+
 // Internal function to build a query from a callback
 // used by liveQueryCollectionOptions.query
 export function buildQuery<TContext extends Context>(

package/src/query/builder/ref-proxy.ts

@@ -1,6 +1,7 @@
 import { PropRef, Value } from '../ir.js'
 import type { BasicExpression } from '../ir.js'
 import type { RefLeaf } from './types.js'
+import type { VirtualRowProps } from '../../virtual-props.js'
 
 export interface RefProxy<T = any> {
   /** @internal */
@@ -11,18 +12,35 @@ export interface RefProxy<T = any> {
   readonly __type: T
 }
 
+/**
+ * Virtual properties available on all row ref proxies.
+ * These allow querying on sync status, origin, key, and collection ID.
+ */
+export type VirtualPropsRefProxy<
+  TKey extends string | number = string | number,
+> = {
+  readonly [K in keyof VirtualRowProps<TKey>]: RefLeaf<VirtualRowProps<TKey>[K]>
+}
+
 /**
  * Type for creating a RefProxy for a single row/type without namespacing
  * Used in collection indexes and where clauses
+ *
+ * Includes virtual properties ($synced, $origin, $key, $collectionId) for
+ * querying on sync status and row metadata.
  */
-export type SingleRowRefProxy<T> =
+export type SingleRowRefProxy<
+  T,
+  TKey extends string | number = string | number,
+> =
   T extends Record<string, any>
     ? {
         [K in keyof T]: T[K] extends Record<string, any>
-          ? SingleRowRefProxy<T[K]> & RefProxy<T[K]>
+          ? SingleRowRefProxy<T[K], TKey> & RefProxy<T[K]>
           : RefLeaf<T[K]>
-      } & RefProxy<T>
-    : RefProxy<T>
+      } & RefProxy<T> &
+        VirtualPropsRefProxy<TKey>
+    : RefProxy<T> & VirtualPropsRefProxy<TKey>
 
 /**
  * Creates a proxy object that records property access paths for a single row