@tanstack/db 0.4.7 → 0.4.9

package/src/query/live/collection-registry.ts

@@ -0,0 +1,47 @@
+import type { Collection } from "../../collection/index.js"
+import type { CollectionConfigBuilder } from "./collection-config-builder.js"
+
+const collectionBuilderRegistry = new WeakMap<
+  Collection<any, any, any>,
+  CollectionConfigBuilder<any, any>
+>()
+
+/**
+ * Retrieves the builder attached to a config object via its utils.getBuilder() method.
+ *
+ * @param config - The collection config object
+ * @returns The attached builder, or `undefined` if none exists
+ */
+export function getBuilderFromConfig(
+  config: object
+): CollectionConfigBuilder<any, any> | undefined {
+  return (config as any).utils?.getBuilder?.()
+}
+
+/**
+ * Registers a builder for a collection in the global registry.
+ * Used to detect when a live query depends on another live query,
+ * enabling the scheduler to ensure parent queries run first.
+ *
+ * @param collection - The collection to register the builder for
+ * @param builder - The builder that produces this collection
+ */
+export function registerCollectionBuilder(
+  collection: Collection<any, any, any>,
+  builder: CollectionConfigBuilder<any, any>
+): void {
+  collectionBuilderRegistry.set(collection, builder)
+}
+
+/**
+ * Retrieves the builder registered for a collection.
+ * Used to discover dependencies when a live query subscribes to another live query.
+ *
+ * @param collection - The collection to look up
+ * @returns The registered builder, or `undefined` if none exists
+ */
+export function getCollectionBuilder(
+  collection: Collection<any, any, any>
+): CollectionConfigBuilder<any, any> | undefined {
+  return collectionBuilderRegistry.get(collection)
+}

package/src/query/live/collection-subscriber.ts

@@ -3,15 +3,20 @@ import {
   convertOrderByToBasicExpression,
   convertToBasicExpression,
 } from "../compiler/expressions.js"
-import type { FullSyncState } from "./types.js"
+import { WhereClauseConversionError } from "../../errors.js"
 import type { MultiSetArray, RootStreamBuilder } from "@tanstack/db-ivm"
 import type { Collection } from "../../collection/index.js"
-import type { ChangeMessage, SyncConfig } from "../../types.js"
+import type { ChangeMessage } from "../../types.js"
 import type { Context, GetResult } from "../builder/types.js"
 import type { BasicExpression } from "../ir.js"
+import type { OrderByOptimizationInfo } from "../compiler/order-by.js"
 import type { CollectionConfigBuilder } from "./collection-config-builder.js"
 import type { CollectionSubscription } from "../../collection/subscription.js"
 
+const loadMoreCallbackSymbol = Symbol.for(
+  `@tanstack/db.collection-config-builder`
+)
+
 export class CollectionSubscriber<
   TContext extends Context,
   TResult extends object = GetResult<TContext>,
@@ -19,61 +24,42 @@ export class CollectionSubscriber<
   // Keep track of the biggest value we've sent so far (needed for orderBy optimization)
   private biggest: any = undefined
 
-  private collectionAlias: string
-
   constructor(
+    private alias: string,
     private collectionId: string,
     private collection: Collection,
-    private config: Parameters<SyncConfig<TResult>[`sync`]>[0],
-    private syncState: FullSyncState,
     private collectionConfigBuilder: CollectionConfigBuilder<TContext, TResult>
-  ) {
-    this.collectionAlias = findCollectionAlias(
-      this.collectionId,
-      this.collectionConfigBuilder.query
-    )!
-  }
+  ) {}
 
   subscribe(): CollectionSubscription {
-    const whereClause = this.getWhereClauseFromAlias(this.collectionAlias)
+    const whereClause = this.getWhereClauseForAlias()
 
     if (whereClause) {
-      // Convert WHERE clause to BasicExpression format for collection subscription
-      const whereExpression = convertToBasicExpression(
-        whereClause,
-        this.collectionAlias
-      )
+      const whereExpression = convertToBasicExpression(whereClause, this.alias)
 
       if (whereExpression) {
-        // Use index optimization for this collection
         return this.subscribeToChanges(whereExpression)
-      } else {
-        // This should not happen - if we have a whereClause but can't create whereExpression,
-        // it indicates a bug in our optimization logic
-        throw new Error(
-          `Failed to convert WHERE clause to collection filter for collection '${this.collectionId}'. ` +
-            `This indicates a bug in the query optimization logic.`
-        )
       }
-    } else {
-      // No WHERE clause for this collection, use regular subscription
-      return this.subscribeToChanges()
+
+      throw new WhereClauseConversionError(this.collectionId, this.alias)
     }
+
+    return this.subscribeToChanges()
   }
 
   private subscribeToChanges(whereExpression?: BasicExpression<boolean>) {
     let subscription: CollectionSubscription
-    if (
-      Object.hasOwn(
-        this.collectionConfigBuilder.optimizableOrderByCollections,
-        this.collectionId
+    const orderByInfo = this.getOrderByInfo()
+    if (orderByInfo) {
+      subscription = this.subscribeToOrderedChanges(
+        whereExpression,
+        orderByInfo
       )
-    ) {
-      subscription = this.subscribeToOrderedChanges(whereExpression)
     } else {
-      // If the collection is lazy then we should not include the initial state
-      const includeInitialState =
-        !this.collectionConfigBuilder.lazyCollections.has(this.collectionId)
+      // If the source alias is lazy then we should not include the initial state
+      const includeInitialState = !this.collectionConfigBuilder.isLazyAlias(
+        this.alias
+      )
 
       subscription = this.subscribeToMatchingChanges(
         whereExpression,
@@ -83,7 +69,11 @@ export class CollectionSubscriber<
     const unsubscribe = () => {
       subscription.unsubscribe()
     }
-    this.syncState.unsubscribeCallbacks.add(unsubscribe)
+    // currentSyncState is always defined when subscribe() is called
+    // (called during sync session setup)
+    this.collectionConfigBuilder.currentSyncState!.unsubscribeCallbacks.add(
+      unsubscribe
+    )
     return subscription
   }
 
@@ -91,7 +81,10 @@ export class CollectionSubscriber<
     changes: Iterable<ChangeMessage<any, string | number>>,
     callback?: () => boolean
   ) {
-    const input = this.syncState.inputs[this.collectionId]!
+    // currentSyncState and input are always defined when this method is called
+    // (only called from active subscriptions during a sync session)
+    const input =
+      this.collectionConfigBuilder.currentSyncState!.inputs[this.alias]!
     const sentChanges = sendChangesToInput(
       input,
       changes,
@@ -103,14 +96,12 @@ export class CollectionSubscriber<
     // otherwise we end up in an infinite loop trying to load more data
     const dataLoader = sentChanges > 0 ? callback : undefined
 
-    // We need to call `maybeRunGraph` even if there's no data to load
+    // We need to schedule a graph run even if there's no data to load
     // because we need to mark the collection as ready if it's not already
-    // and that's only done in `maybeRunGraph`
-    this.collectionConfigBuilder.maybeRunGraph(
-      this.config,
-      this.syncState,
-      dataLoader
-    )
+    // and that's only done in `scheduleGraphRun`
+    this.collectionConfigBuilder.scheduleGraphRun(dataLoader, {
+      alias: this.alias,
+    })
   }
 
   private subscribeToMatchingChanges(
@@ -132,12 +123,11 @@ export class CollectionSubscriber<
   }
 
   private subscribeToOrderedChanges(
-    whereExpression: BasicExpression<boolean> | undefined
+    whereExpression: BasicExpression<boolean> | undefined,
+    orderByInfo: OrderByOptimizationInfo
   ) {
     const { orderBy, offset, limit, comparator, dataNeeded, index } =
-      this.collectionConfigBuilder.optimizableOrderByCollections[
-        this.collectionId
-      ]!
+      orderByInfo
 
     const sendChangesInRange = (
       changes: Iterable<ChangeMessage<any, string | number>>
@@ -147,7 +137,7 @@ export class CollectionSubscriber<
       // because they can't affect the topK (and if later we need more data, we will dynamically load more data)
       const splittedChanges = splitUpdates(changes)
       let filteredChanges = splittedChanges
-      if (dataNeeded!() === 0) {
+      if (dataNeeded && dataNeeded() === 0) {
         // If the topK is full [..., maxSentValue] then we do not need to send changes > maxSentValue
         // because they can never make it into the topK.
         // However, if the topK isn't full yet, we need to also send changes > maxSentValue
@@ -173,7 +163,7 @@ export class CollectionSubscriber<
     // Normalize the orderBy clauses such that the references are relative to the collection
     const normalizedOrderBy = convertOrderByToBasicExpression(
       orderBy,
-      this.collectionAlias
+      this.alias
     )
 
     // Load the first `offset + limit` values from the index
@@ -190,10 +180,7 @@ export class CollectionSubscriber<
   // after each iteration of the query pipeline
   // to ensure that the orderBy operator has enough data to work with
   loadMoreIfNeeded(subscription: CollectionSubscription) {
-    const orderByInfo =
-      this.collectionConfigBuilder.optimizableOrderByCollections[
-        this.collectionId
-      ]
+    const orderByInfo = this.getOrderByInfo()
 
     if (!orderByInfo) {
       // This query has no orderBy operator
@@ -224,24 +211,40 @@ export class CollectionSubscriber<
     changes: Iterable<ChangeMessage<any, string | number>>,
     subscription: CollectionSubscription
   ) {
-    const { comparator } =
-      this.collectionConfigBuilder.optimizableOrderByCollections[
-        this.collectionId
-      ]!
-    const trackedChanges = this.trackSentValues(changes, comparator)
+    const orderByInfo = this.getOrderByInfo()
+    if (!orderByInfo) {
+      this.sendChangesToPipeline(changes)
+      return
+    }
+
+    const trackedChanges = this.trackSentValues(changes, orderByInfo.comparator)
+
+    // Cache the loadMoreIfNeeded callback on the subscription using a symbol property.
+    // This ensures we pass the same function instance to the scheduler each time,
+    // allowing it to deduplicate callbacks when multiple changes arrive during a transaction.
+    type SubscriptionWithLoader = CollectionSubscription & {
+      [loadMoreCallbackSymbol]?: () => boolean
+    }
+
+    const subscriptionWithLoader = subscription as SubscriptionWithLoader
+
+    subscriptionWithLoader[loadMoreCallbackSymbol] ??=
+      this.loadMoreIfNeeded.bind(this, subscription)
+
     this.sendChangesToPipeline(
       trackedChanges,
-      this.loadMoreIfNeeded.bind(this, subscription)
+      subscriptionWithLoader[loadMoreCallbackSymbol]
     )
   }
 
   // Loads the next `n` items from the collection
   // starting from the biggest item it has sent
   private loadNextItems(n: number, subscription: CollectionSubscription) {
-    const { orderBy, valueExtractorForRawRow } =
-      this.collectionConfigBuilder.optimizableOrderByCollections[
-        this.collectionId
-      ]!
+    const orderByInfo = this.getOrderByInfo()
+    if (!orderByInfo) {
+      return
+    }
+    const { orderBy, valueExtractorForRawRow } = orderByInfo
     const biggestSentRow = this.biggest
     const biggestSentValue = biggestSentRow
       ? valueExtractorForRawRow(biggestSentRow)
@@ -250,7 +253,7 @@ export class CollectionSubscriber<
     // Normalize the orderBy clauses such that the references are relative to the collection
     const normalizedOrderBy = convertOrderByToBasicExpression(
       orderBy,
-      this.collectionAlias
+      this.alias
     )
 
     // Take the `n` items after the biggest sent value
@@ -261,13 +264,22 @@ export class CollectionSubscriber<
     })
   }
 
-  private getWhereClauseFromAlias(
-    collectionAlias: string | undefined
-  ): BasicExpression<boolean> | undefined {
-    const collectionWhereClausesCache =
-      this.collectionConfigBuilder.collectionWhereClausesCache
-    if (collectionAlias && collectionWhereClausesCache) {
-      return collectionWhereClausesCache.get(collectionAlias)
+  private getWhereClauseForAlias(): BasicExpression<boolean> | undefined {
+    const sourceWhereClausesCache =
+      this.collectionConfigBuilder.sourceWhereClausesCache
+    if (!sourceWhereClausesCache) {
+      return undefined
+    }
+    return sourceWhereClausesCache.get(this.alias)
+  }
+
+  private getOrderByInfo(): OrderByOptimizationInfo | undefined {
+    const info =
+      this.collectionConfigBuilder.optimizableOrderByCollections[
+        this.collectionId
+      ]
+    if (info && info.alias === this.alias) {
+      return info
     }
     return undefined
   }
@@ -288,36 +300,6 @@ export class CollectionSubscriber<
   }
 }
 
-/**
- * Finds the alias for a collection ID in the query
- */
-function findCollectionAlias(
-  collectionId: string,
-  query: any
-): string | undefined {
-  // Check FROM clause
-  if (
-    query.from?.type === `collectionRef` &&
-    query.from.collection?.id === collectionId
-  ) {
-    return query.from.alias
-  }
-
-  // Check JOIN clauses
-  if (query.join) {
-    for (const joinClause of query.join) {
-      if (
-        joinClause.from?.type === `collectionRef` &&
-        joinClause.from.collection?.id === collectionId
-      ) {
-        return joinClause.from.alias
-      }
-    }
-  }
-
-  return undefined
-}
-
 /**
  * Helper function to send changes to a D2 input stream
  */

package/src/query/live-query-collection.ts

@@ -1,5 +1,10 @@
 import { createCollection } from "../collection/index.js"
 import { CollectionConfigBuilder } from "./live/collection-config-builder.js"
+import {
+  getBuilderFromConfig,
+  registerCollectionBuilder,
+} from "./live/collection-registry.js"
+import type { LiveQueryCollectionUtils } from "./live/collection-config-builder.js"
 import type { LiveQueryCollectionConfig } from "./live/types.js"
 import type { InitialQueryBuilder, QueryBuilder } from "./builder/index.js"
 import type { Collection } from "../collection/index.js"
@@ -55,7 +60,9 @@ export function liveQueryCollectionOptions<
   TResult extends object = GetResult<TContext>,
 >(
   config: LiveQueryCollectionConfig<TContext, TResult>
-): CollectionConfigForContext<TContext, TResult> {
+): CollectionConfigForContext<TContext, TResult> & {
+  utils: LiveQueryCollectionUtils
+} {
   const collectionConfigBuilder = new CollectionConfigBuilder<
     TContext,
     TResult
@@ -63,7 +70,7 @@ export function liveQueryCollectionOptions<
   return collectionConfigBuilder.getConfig() as CollectionConfigForContext<
     TContext,
     TResult
-  >
+  > & { utils: LiveQueryCollectionUtils }
 }
 
 /**
@@ -106,7 +113,9 @@ export function createLiveQueryCollection<
   TResult extends object = GetResult<TContext>,
 >(
   query: (q: InitialQueryBuilder) => QueryBuilder<TContext>
-): CollectionForContext<TContext, TResult>
+): CollectionForContext<TContext, TResult> & {
+  utils: LiveQueryCollectionUtils
+}
 
 // Overload 2: Accept full config object with optional utilities
 export function createLiveQueryCollection<
@@ -115,7 +124,9 @@ export function createLiveQueryCollection<
   TUtils extends UtilsRecord = {},
 >(
   config: LiveQueryCollectionConfig<TContext, TResult> & { utils?: TUtils }
-): CollectionForContext<TContext, TResult>
+): CollectionForContext<TContext, TResult> & {
+  utils: LiveQueryCollectionUtils & TUtils
+}
 
 // Implementation
 export function createLiveQueryCollection<
@@ -126,7 +137,9 @@ export function createLiveQueryCollection<
   configOrQuery:
     | (LiveQueryCollectionConfig<TContext, TResult> & { utils?: TUtils })
     | ((q: InitialQueryBuilder) => QueryBuilder<TContext>)
-): CollectionForContext<TContext, TResult> {
+): CollectionForContext<TContext, TResult> & {
+  utils: LiveQueryCollectionUtils & TUtils
+} {
   // Determine if the argument is a function (query) or a config object
   if (typeof configOrQuery === `function`) {
     // Simple query function case
@@ -139,7 +152,7 @@ export function createLiveQueryCollection<
     return bridgeToCreateCollection(options) as CollectionForContext<
       TContext,
       TResult
-    >
+    > & { utils: LiveQueryCollectionUtils & TUtils }
   } else {
     // Config object case
     const config = configOrQuery as LiveQueryCollectionConfig<
@@ -147,10 +160,16 @@ export function createLiveQueryCollection<
       TResult
     > & { utils?: TUtils }
     const options = liveQueryCollectionOptions<TContext, TResult>(config)
-    return bridgeToCreateCollection({
-      ...options,
-      utils: config.utils,
-    }) as CollectionForContext<TContext, TResult>
+
+    // Merge custom utils if provided, preserving the getBuilder() method for dependency tracking
+    if (config.utils) {
+      options.utils = { ...options.utils, ...config.utils }
+    }
+
+    return bridgeToCreateCollection(options) as CollectionForContext<
+      TContext,
+      TResult
+    > & { utils: LiveQueryCollectionUtils & TUtils }
   }
 }
 
@@ -162,12 +181,18 @@ function bridgeToCreateCollection<
   TResult extends object,
   TUtils extends UtilsRecord = {},
 >(
-  options: CollectionConfig<TResult> & { utils?: TUtils }
+  options: CollectionConfig<TResult> & { utils: TUtils }
 ): Collection<TResult, string | number, TUtils> {
-  // This is the only place we need a type assertion, hidden from user API
-  return createCollection(options as any) as unknown as Collection<
+  const collection = createCollection(options as any) as unknown as Collection<
     TResult,
     string | number,
-    TUtils
+    LiveQueryCollectionUtils
   >
+
+  const builder = getBuilderFromConfig(options)
+  if (builder) {
+    registerCollectionBuilder(collection, builder)
+  }
+
+  return collection as unknown as Collection<TResult, string | number, TUtils>
 }

package/src/query/optimizer.ts

@@ -162,8 +162,8 @@ export interface GroupedWhereClauses {
 export interface OptimizationResult {
   /** The optimized query with WHERE clauses potentially moved to subqueries */
   optimizedQuery: QueryIR
-  /** Map of collection aliases to their extracted WHERE clauses for index optimization */
-  collectionWhereClauses: Map<string, BasicExpression<boolean>>
+  /** Map of source aliases to their extracted WHERE clauses for index optimization */
+  sourceWhereClauses: Map<string, BasicExpression<boolean>>
 }
 
 /**
@@ -184,14 +184,14 @@ export interface OptimizationResult {
  *   where: [eq(u.dept_id, 1), gt(p.views, 100)]
  * }
  *
- * const { optimizedQuery, collectionWhereClauses } = optimizeQuery(originalQuery)
+ * const { optimizedQuery, sourceWhereClauses } = optimizeQuery(originalQuery)
  * // Result: Single-source clauses moved to deepest possible subqueries
- * // collectionWhereClauses: Map { 'u' => eq(u.dept_id, 1), 'p' => gt(p.views, 100) }
+ * // sourceWhereClauses: Map { 'u' => eq(u.dept_id, 1), 'p' => gt(p.views, 100) }
  * ```
  */
 export function optimizeQuery(query: QueryIR): OptimizationResult {
-  // First, extract collection WHERE clauses before optimization
-  const collectionWhereClauses = extractCollectionWhereClauses(query)
+  // First, extract source WHERE clauses before optimization
+  const sourceWhereClauses = extractSourceWhereClauses(query)
 
   // Apply multi-level predicate pushdown with iterative convergence
   let optimized = query
@@ -214,7 +214,7 @@ export function optimizeQuery(query: QueryIR): OptimizationResult {
 
   return {
     optimizedQuery: cleaned,
-    collectionWhereClauses,
+    sourceWhereClauses,
   }
 }
 
@@ -224,16 +224,16 @@ export function optimizeQuery(query: QueryIR): OptimizationResult {
  * to specific collections, but only for simple queries without joins.
  *
  * @param query - The original QueryIR to analyze
- * @returns Map of collection aliases to their WHERE clauses
+ * @returns Map of source aliases to their WHERE clauses
  */
-function extractCollectionWhereClauses(
+function extractSourceWhereClauses(
   query: QueryIR
 ): Map<string, BasicExpression<boolean>> {
-  const collectionWhereClauses = new Map<string, BasicExpression<boolean>>()
+  const sourceWhereClauses = new Map<string, BasicExpression<boolean>>()
 
   // Only analyze queries that have WHERE clauses
   if (!query.where || query.where.length === 0) {
-    return collectionWhereClauses
+    return sourceWhereClauses
   }
 
   // Split all AND clauses at the root level for granular analysis
@@ -254,12 +254,12 @@ function extractCollectionWhereClauses(
     if (isCollectionReference(query, sourceAlias)) {
       // Check if the WHERE clause can be converted to collection-compatible format
       if (isConvertibleToCollectionFilter(whereClause)) {
-        collectionWhereClauses.set(sourceAlias, whereClause)
+        sourceWhereClauses.set(sourceAlias, whereClause)
       }
     }
   }
 
-  return collectionWhereClauses
+  return sourceWhereClauses
 }
 
 /**
@@ -782,8 +782,6 @@ function optimizeFromWithTracking(
     return new QueryRefClass(subQuery, from.alias)
   }
 
-  // Must be queryRef due to type system
-
   // SAFETY CHECK: Only check safety when pushing WHERE clauses into existing subqueries
   // We need to be careful about pushing WHERE clauses into subqueries that already have
   // aggregates, HAVING, or ORDER BY + LIMIT since that could change their semantics
@@ -793,6 +791,12 @@ function optimizeFromWithTracking(
     return new QueryRefClass(deepCopyQuery(from.query), from.alias)
   }
 
+  // Skip pushdown when a clause references a field that only exists via a renamed
+  // projection inside the subquery; leaving it outside preserves the alias mapping.
+  if (referencesAliasWithRemappedSelect(from.query, whereClause, from.alias)) {
+    return new QueryRefClass(deepCopyQuery(from.query), from.alias)
+  }
+
   // Add the WHERE clause to the existing subquery
   // Create a deep copy to ensure immutability
   const existingWhere = from.query.where || []
@@ -943,6 +947,72 @@ function whereReferencesComputedSelectFields(
   return false
 }
 
+/**
+ * Detects whether a WHERE clause references the subquery alias through fields that
+ * are re-exposed under different names (renamed SELECT projections or fnSelect output).
+ * In those cases we keep the clause at the outer level to avoid alias remapping bugs.
+ * TODO: in future we should handle this by rewriting the clause to use the subquery's
+ * internal field references, but it likely needs a wider refactor to do cleanly.
+ */
+function referencesAliasWithRemappedSelect(
+  subquery: QueryIR,
+  whereClause: BasicExpression<boolean>,
+  outerAlias: string
+): boolean {
+  const refs = collectRefs(whereClause)
+  // Only care about clauses that actually reference the outer alias.
+  if (refs.every((ref) => ref.path[0] !== outerAlias)) {
+    return false
+  }
+
+  // fnSelect always rewrites the row shape, so alias-safe pushdown is impossible.
+  if (subquery.fnSelect) {
+    return true
+  }
+
+  const select = subquery.select
+  // Without an explicit SELECT the clause still refers to the original collection.
+  if (!select) {
+    return false
+  }
+
+  for (const ref of refs) {
+    const path = ref.path
+    // Need at least alias + field to matter.
+    if (path.length < 2) continue
+    if (path[0] !== outerAlias) continue
+
+    const projected = select[path[1]!]
+    // Unselected fields can't be remapped, so skip - only care about fields in the SELECT.
+    if (!projected) continue
+
+    // Non-PropRef projections are computed values; cannot push down.
+    if (!(projected instanceof PropRef)) {
+      return true
+    }
+
+    // If the projection is just the alias (whole row) without a specific field,
+    // we can't verify whether the field we're referencing is being preserved or remapped.
+    if (projected.path.length < 2) {
+      return true
+    }
+
+    const [innerAlias, innerField] = projected.path
+
+    // Safe only when the projection points straight back to the same alias or the
+    // underlying source alias and preserves the field name.
+    if (innerAlias !== outerAlias && innerAlias !== subquery.from.alias) {
+      return true
+    }
+
+    if (innerField !== path[1]) {
+      return true
+    }
+  }
+
+  return false
+}
+
 /**
  * Helper function to combine multiple expressions with AND.
  *