@tanstack/db 0.4.8 → 0.4.10

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,71 +24,106 @@ 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
+  // Track deferred promises for subscription loading states
+  private subscriptionLoadingPromises = new Map<
+    CollectionSubscription,
+    { resolve: () => void }
+  >()
 
   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,
         includeInitialState
       )
     }
+
+    // Subscribe to subscription status changes to propagate loading state
+    const statusUnsubscribe = subscription.on(`status:change`, (event) => {
+      // TODO: For now we are setting this loading state whenever the subscription
+      // status changes to 'loadingSubset'. But we have discussed it only happening
+      // when the the live query has it's offset/limit changed, and that triggers the
+      // subscription to request a snapshot. This will require more work to implement,
+      // and builds on https://github.com/TanStack/db/pull/663 which this PR
+      // does not yet depend on.
+      if (event.status === `loadingSubset`) {
+        // Guard against duplicate transitions
+        if (!this.subscriptionLoadingPromises.has(subscription)) {
+          let resolve: () => void
+          const promise = new Promise<void>((res) => {
+            resolve = res
+          })
+
+          this.subscriptionLoadingPromises.set(subscription, {
+            resolve: resolve!,
+          })
+          this.collectionConfigBuilder.liveQueryCollection!._sync.trackLoadPromise(
+            promise
+          )
+        }
+      } else {
+        // status is 'ready'
+        const deferred = this.subscriptionLoadingPromises.get(subscription)
+        if (deferred) {
+          // Clear the map entry FIRST (before resolving)
+          this.subscriptionLoadingPromises.delete(subscription)
+          deferred.resolve()
+        }
+      }
+    })
+
     const unsubscribe = () => {
+      // If subscription has a pending promise, resolve it before unsubscribing
+      const deferred = this.subscriptionLoadingPromises.get(subscription)
+      if (deferred) {
+        // Clear the map entry FIRST (before resolving)
+        this.subscriptionLoadingPromises.delete(subscription)
+        deferred.resolve()
+      }
+
+      statusUnsubscribe()
       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 +131,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 +146,12 @@ export class CollectionSubscriber<
     // otherwise we end up in an infinite loop trying to load more data
     const dataLoader = sentChanges > 0 ? callback : undefined
 
-    // Always call maybeRunGraph to process changes eagerly.
-    // The graph will run unless the live query is in an error state.
-    // Status management is handled separately via status:change event listeners.
-    this.collectionConfigBuilder.maybeRunGraph(
-      this.config,
-      this.syncState,
-      dataLoader
-    )
+    // 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 `scheduleGraphRun`
+    this.collectionConfigBuilder.scheduleGraphRun(dataLoader, {
+      alias: this.alias,
+    })
   }
 
   private subscribeToMatchingChanges(
@@ -132,12 +173,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 +187,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 +213,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 +230,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 +261,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 +303,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 +314,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 +350,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"
@@ -15,16 +20,20 @@ import type { Context, GetResult } from "./builder/types.js"
 type CollectionConfigForContext<
   TContext extends Context,
   TResult extends object,
+  TUtils extends UtilsRecord = {},
 > = TContext extends SingleResult
-  ? CollectionConfigSingleRowOption<TResult> & SingleResult
-  : CollectionConfigSingleRowOption<TResult> & NonSingleResult
+  ? CollectionConfigSingleRowOption<TResult, string | number, never, TUtils> &
+      SingleResult
+  : CollectionConfigSingleRowOption<TResult, string | number, never, TUtils> &
+      NonSingleResult
 
 type CollectionForContext<
   TContext extends Context,
   TResult extends object,
+  TUtils extends UtilsRecord = {},
 > = TContext extends SingleResult
-  ? Collection<TResult> & SingleResult
-  : Collection<TResult> & NonSingleResult
+  ? Collection<TResult, string | number, TUtils> & SingleResult
+  : Collection<TResult, string | number, TUtils> & NonSingleResult
 
 /**
  * Creates live query collection options for use with createCollection
@@ -55,7 +64,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 +74,7 @@ export function liveQueryCollectionOptions<
   return collectionConfigBuilder.getConfig() as CollectionConfigForContext<
     TContext,
     TResult
-  >
+  > & { utils: LiveQueryCollectionUtils }
 }
 
 /**
@@ -106,7 +117,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 +128,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 +141,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 +156,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 +164,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 +185,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>
 }