fix(db): preserve correlated include route identity

.changeset/fix-correlated-include-routes.md

@@ -0,0 +1,5 @@
+---
+'@tanstack/db': patch
+---
+
+Preserve parent-specific include results through nested includes, subqueries, unions, joins, ordering, projections, aggregates, and having clauses.

packages/db/package.json

@@ -21,7 +21,7 @@
     "dev": "vite build --watch",
     "lint": "eslint . --fix",
     "test": "vitest --run",
-    "test:oracles": "vitest --run tests/query/includes-oracle.property.test.ts tests/query/includes-collection-oracle.property.test.ts tests/query/includes-cross-formulation-oracle.property.test.ts tests/query/includes-temporal-oracle.test.ts tests/query/includes-optimistic-oracle.property.test.ts tests/query/includes-publication-oracle.test.ts tests/query/includes-query-shape-oracle.test.ts tests/query/includes-work-counter-oracle.test.ts"
+    "test:oracles": "vitest --run tests/query/includes-oracle.property.test.ts tests/query/includes-collection-oracle.property.test.ts tests/query/includes-cross-formulation-oracle.property.test.ts tests/query/includes-temporal-oracle.test.ts tests/query/includes-optimistic-oracle.property.test.ts tests/query/includes-publication-oracle.test.ts tests/query/includes-query-shape-oracle.test.ts tests/query/includes-work-counter-oracle.test.ts tests/query/includes-context-transport-oracle.test.ts"
   },
   "type": "module",
   "main": "dist/cjs/index.cjs",

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

@@ -1106,14 +1106,17 @@ function buildConditionalSelect(
 /**
  * Recursively collects all PropRef nodes from an expression tree.
  */
-function collectRefsFromExpression(expr: BasicExpression): Array<PropRef> {
+function collectRefsFromExpression(
+  expr: BasicExpression | Aggregate,
+): 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 ?? []) {
+    case `agg`:
+      for (const arg of expr.args) {
         refs.push(...collectRefsFromExpression(arg))
       }
       break
@@ -1123,6 +1126,164 @@ function collectRefsFromExpression(expr: BasicExpression): Array<PropRef> {
   return refs
 }
 
+function collectRefsFromSelectValue(value: unknown): Array<PropRef> {
+  if (
+    value instanceof PropRef ||
+    value instanceof FuncExpr ||
+    value instanceof AggregateExpr
+  ) {
+    return collectRefsFromExpression(value)
+  }
+  if (value instanceof ConditionalSelect) {
+    return [
+      ...value.branches.flatMap((branch) => [
+        ...collectRefsFromExpression(branch.condition),
+        ...collectRefsFromSelectValue(branch.value),
+      ]),
+      ...(value.defaultValue === undefined
+        ? []
+        : collectRefsFromSelectValue(value.defaultValue)),
+    ]
+  }
+  if (value instanceof IncludesSubquery) {
+    return [
+      value.correlationField,
+      ...(value.parentProjection ?? []),
+      ...collectExternalRefsFromQuery(value.query),
+    ]
+  }
+  if (!isPlainObject(value)) return []
+  return Object.values(value).flatMap(collectRefsFromSelectValue)
+}
+
+function collectExternalRefsFromQuery(query: QueryIR): Array<PropRef> {
+  const localAliases = new Set(collectQueryAliases(query))
+  const refs: Array<PropRef> = []
+  const addExpression = (expression: BasicExpression | Aggregate) => {
+    refs.push(...collectRefsFromExpression(expression))
+  }
+  const addWhere = (where: Where) => {
+    addExpression(
+      typeof where === `object` && `expression` in where
+        ? where.expression
+        : where,
+    )
+  }
+
+  for (const where of query.where ?? []) addWhere(where)
+  for (const join of query.join ?? []) {
+    addExpression(join.left)
+    addExpression(join.right)
+    if (join.from.type === `queryRef`) {
+      refs.push(...collectExternalRefsFromQuery(join.from.query))
+    }
+  }
+  for (const expression of query.groupBy ?? []) addExpression(expression)
+  for (const having of query.having ?? []) addWhere(having)
+  for (const { expression } of query.orderBy ?? []) addExpression(expression)
+  if (query.select) refs.push(...collectRefsFromSelectValue(query.select))
+
+  if (query.from.type === `queryRef`) {
+    refs.push(...collectExternalRefsFromQuery(query.from.query))
+  } else if (query.from.type === `unionFrom`) {
+    for (const source of query.from.sources) {
+      if (source.type === `queryRef`) {
+        refs.push(...collectExternalRefsFromQuery(source.query))
+      }
+    }
+  } else if (query.from.type === `unionAll`) {
+    for (const branch of query.from.queries) {
+      refs.push(...collectExternalRefsFromQuery(branch))
+    }
+  }
+
+  const seen = new Set<string>()
+  return refs.filter((ref) => {
+    const alias = ref.path.length > 1 ? ref.path[0] : undefined
+    const path = ref.path.join(`.`)
+    if (
+      alias == null ||
+      alias === `$selected` ||
+      localAliases.has(alias) ||
+      seen.has(path)
+    ) {
+      return false
+    }
+    seen.add(path)
+    return true
+  })
+}
+
+function collectParentRefsFromQuery(
+  query: QueryIR,
+  parentAliases: Array<string>,
+): Array<PropRef> {
+  const refs: Array<PropRef> = []
+  const addExpression = (expression: BasicExpression | Aggregate) => {
+    refs.push(...collectRefsFromExpression(expression))
+  }
+  const addWhere = (where: Where) => {
+    addExpression(
+      typeof where === `object` && `expression` in where
+        ? where.expression
+        : where,
+    )
+  }
+
+  for (const where of query.where ?? []) addWhere(where)
+  for (const join of query.join ?? []) {
+    addExpression(join.left)
+    addExpression(join.right)
+    if (join.from.type === `queryRef`) {
+      refs.push(...collectParentRefsFromQuery(join.from.query, parentAliases))
+    }
+  }
+  for (const expression of query.groupBy ?? []) addExpression(expression)
+  for (const having of query.having ?? []) addWhere(having)
+  for (const { expression } of query.orderBy ?? []) addExpression(expression)
+  if (query.select) {
+    refs.push(...collectRefsFromSelectValue(query.select))
+  }
+
+  if (query.from.type === `queryRef`) {
+    refs.push(...collectParentRefsFromQuery(query.from.query, parentAliases))
+  } else if (query.from.type === `unionFrom`) {
+    for (const source of query.from.sources) {
+      if (source.type === `queryRef`) {
+        refs.push(...collectParentRefsFromQuery(source.query, parentAliases))
+      }
+    }
+  } else if (query.from.type === `unionAll`) {
+    for (const branch of query.from.queries) {
+      refs.push(...collectParentRefsFromQuery(branch, parentAliases))
+    }
+  }
+
+  const seen = new Set<string>()
+  return refs.filter((ref) => {
+    const path = ref.path.join(`.`)
+    if (
+      ref.path[0] == null ||
+      !parentAliases.includes(ref.path[0]) ||
+      seen.has(path)
+    ) {
+      return false
+    }
+    seen.add(path)
+    return true
+  })
+}
+
+function collectExternalParentAliases(query: QueryIR): Array<string> {
+  return [
+    ...new Set(
+      collectExternalRefsFromQuery(query)
+        .map((ref) => ref.path[0])
+        .filter((alias): alias is string => alias !== undefined),
+    ),
+  ]
+}
+
 /**
  * Checks whether a WHERE clause references any parent alias.
  */
@@ -1151,6 +1312,9 @@ function buildIncludesSubquery(
 
   // Collect child's own aliases
   const childAliases = collectQueryAliases(childQuery)
+  const visibleParentAliases = [
+    ...new Set([...parentAliases, ...collectExternalParentAliases(childQuery)]),
+  ]
 
   // 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().
@@ -1176,7 +1340,7 @@ function buildIncludesSubquery(
         const result = extractCorrelation(
           expr.args[0]!,
           expr.args[1]!,
-          parentAliases,
+          visibleParentAliases,
           childAliases,
         )
         if (result) {
@@ -1203,7 +1367,7 @@ function buildIncludesSubquery(
             const result = extractCorrelation(
               arg.args[0]!,
               arg.args[1]!,
-              parentAliases,
+              visibleParentAliases,
               childAliases,
             )
             if (result) {
@@ -1264,32 +1428,21 @@ function buildIncludesSubquery(
   const pureChildWhere: Array<Where> = []
   const parentFilters: Array<Where> = []
   for (const w of modifiedWhere) {
-    if (referencesParent(w, parentAliases)) {
+    if (referencesParent(w, visibleParentAliases)) {
       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)
-        }
-      }
-    }
-  }
+  // Every parent input that can affect the child plan belongs to the route
+  // identity, not only the main equality key or residual filters.
+  const projectedParentRefs = collectParentRefsFromQuery(
+    { ...childQuery, where: modifiedWhere },
+    visibleParentAliases,
+  )
+  const parentProjection =
+    projectedParentRefs.length > 0 ? projectedParentRefs : undefined
 
   const modifiedQuery: QueryIR = {
     ...childQuery,

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

@@ -93,7 +93,7 @@ export type Source = {
   [alias: string]:
     | CollectionImpl<any, any>
     | CollectionOptionsIdentity<any, any, any, any, any>
-    | QueryBuilder<Context>
+    | QueryBuilder<any>
 }
 
 /**
@@ -168,7 +168,7 @@ export type ContextFromUnionSource<TSource extends Source> =
     : ContextFromSource<TSource>
 
 type ResultFromBranch<TBranch> =
-  TBranch extends QueryBuilder<infer TContext> ? GetResult<TContext> : never
+  TBranch extends QueryBuilder<infer TContext> ? GetRawResult<TContext> : never
 
 type UnionBranchResult<TBranches extends ReadonlyArray<QueryBuilder<any>>> =
   ResultFromBranch<TBranches[number]>