@tanstack/db 0.4.8 → 0.4.9

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

@@ -1,7 +1,12 @@
 import { D2, output } from "@tanstack/db-ivm"
 import { compileQuery } from "../compiler/index.js"
 import { buildQuery, getQueryIR } from "../builder/index.js"
+import { MissingAliasInputsError } from "../../errors.js"
+import { transactionScopedScheduler } from "../../scheduler.js"
+import { getActiveTransaction } from "../../transactions.js"
 import { CollectionSubscriber } from "./collection-subscriber.js"
+import { getCollectionBuilder } from "./collection-registry.js"
+import type { SchedulerContextId } from "../../scheduler.js"
 import type { CollectionSubscription } from "../../collection/subscription.js"
 import type { RootStreamBuilder } from "@tanstack/db-ivm"
 import type { OrderByOptimizationInfo } from "../compiler/order-by.js"
@@ -11,6 +16,7 @@ import type {
   KeyedStream,
   ResultStream,
   SyncConfig,
+  UtilsRecord,
 } from "../../types.js"
 import type { Context, GetResult } from "../builder/types.js"
 import type { BasicExpression, QueryIR } from "../ir.js"
@@ -23,6 +29,15 @@ import type {
 } from "./types.js"
 import type { AllCollectionEvents } from "../../collection/events.js"
 
+export type LiveQueryCollectionUtils = UtilsRecord & {
+  getRunCount: () => number
+  getBuilder: () => CollectionConfigBuilder<any, any>
+}
+
+type PendingGraphRun = {
+  loadCallbacks: Set<() => boolean>
+}
+
 // Global counter for auto-generated collection IDs
 let liveQueryCollectionCounter = 0
 
@@ -37,6 +52,9 @@ export class CollectionConfigBuilder<
   private readonly id: string
   readonly query: QueryIR
   private readonly collections: Record<string, Collection<any, any, any>>
+  private readonly collectionByAlias: Record<string, Collection<any, any, any>>
+  // Populated during compilation with all aliases (including subquery inner aliases)
+  private compiledAliasToCollectionId: Record<string, string> = {}
 
   // WeakMap to store the keys of the results
   // so that we can retrieve them in the getKey function
@@ -48,6 +66,14 @@ export class CollectionConfigBuilder<
   private readonly compare?: (val1: TResult, val2: TResult) => number
 
   private isGraphRunning = false
+  private runCount = 0
+
+  // Current sync session state (set when sync starts, cleared when it stops)
+  // Public for testing purposes (CollectionConfigBuilder is internal, not public API)
+  public currentSyncConfig:
+    | Parameters<SyncConfig<TResult>[`sync`]>[0]
+    | undefined
+  public currentSyncState: FullSyncState | undefined
 
   // Error state tracking
   private isInErrorState = false
@@ -55,19 +81,41 @@ export class CollectionConfigBuilder<
   // Reference to the live query collection for error state transitions
   private liveQueryCollection?: Collection<TResult, any, any>
 
+  private readonly aliasDependencies: Record<
+    string,
+    Array<CollectionConfigBuilder<any, any>>
+  > = {}
+
+  private readonly builderDependencies = new Set<
+    CollectionConfigBuilder<any, any>
+  >()
+
+  // Pending graph runs per scheduler context (e.g., per transaction)
+  // The builder manages its own state; the scheduler just orchestrates execution order
+  // Only stores callbacks - if sync ends, pending jobs gracefully no-op
+  private readonly pendingGraphRuns = new Map<
+    SchedulerContextId,
+    PendingGraphRun
+  >()
+
+  // Unsubscribe function for scheduler's onClear listener
+  // Registered when sync starts, unregistered when sync stops
+  // Prevents memory leaks by releasing the scheduler's reference to this builder
+  private unsubscribeFromSchedulerClears?: () => void
+
   private graphCache: D2 | undefined
   private inputsCache: Record<string, RootStreamBuilder<unknown>> | undefined
   private pipelineCache: ResultStream | undefined
-  public collectionWhereClausesCache:
+  public sourceWhereClausesCache:
     | Map<string, BasicExpression<boolean>>
     | undefined
 
-  // Map of collection ID to subscription
+  // Map of source alias to subscription
   readonly subscriptions: Record<string, CollectionSubscription> = {}
-  // Map of collection IDs to functions that load keys for that lazy collection
-  lazyCollectionsCallbacks: Record<string, LazyCollectionCallbacks> = {}
-  // Set of collection IDs that are lazy collections
-  readonly lazyCollections = new Set<string>()
+  // Map of source aliases to functions that load keys for that lazy source
+  lazySourcesCallbacks: Record<string, LazyCollectionCallbacks> = {}
+  // Set of source aliases that are lazy (don't load initial state)
+  readonly lazySources = new Set<string>()
   // Set of collection IDs that include an optimizable ORDER BY clause
   optimizableOrderByCollections: Record<string, OrderByOptimizationInfo> = {}
 
@@ -79,6 +127,19 @@ export class CollectionConfigBuilder<
 
     this.query = buildQueryFromConfig(config)
     this.collections = extractCollectionsFromQuery(this.query)
+    const collectionAliasesById = extractCollectionAliases(this.query)
+
+    // Build a reverse lookup map from alias to collection instance.
+    // This enables self-join support where the same collection can be referenced
+    // multiple times with different aliases (e.g., { employee: col, manager: col })
+    this.collectionByAlias = {}
+    for (const [collectionId, aliases] of collectionAliasesById.entries()) {
+      const collection = this.collections[collectionId]
+      if (!collection) continue
+      for (const alias of aliases) {
+        this.collectionByAlias[alias] = collection
+      }
+    }
 
     // Create compare function for ordering if the query has orderBy
     if (this.query.orderBy && this.query.orderBy.length > 0) {
@@ -90,7 +151,9 @@ export class CollectionConfigBuilder<
     this.compileBasePipeline()
   }
 
-  getConfig(): CollectionConfigSingleRowOption<TResult> {
+  getConfig(): CollectionConfigSingleRowOption<TResult> & {
+    utils: LiveQueryCollectionUtils
+  } {
     return {
       id: this.id,
       getKey:
@@ -105,7 +168,38 @@ export class CollectionConfigBuilder<
       onDelete: this.config.onDelete,
       startSync: this.config.startSync,
       singleResult: this.query.singleResult,
+      utils: {
+        getRunCount: this.getRunCount.bind(this),
+        getBuilder: () => this,
+      },
+    }
+  }
+
+  /**
+   * Resolves a collection alias to its collection ID.
+   *
+   * Uses a two-tier lookup strategy:
+   * 1. First checks compiled aliases (includes subquery inner aliases)
+   * 2. Falls back to declared aliases from the query's from/join clauses
+   *
+   * @param alias - The alias to resolve (e.g., "employee", "manager")
+   * @returns The collection ID that the alias references
+   * @throws {Error} If the alias is not found in either lookup
+   */
+  getCollectionIdForAlias(alias: string): string {
+    const compiled = this.compiledAliasToCollectionId[alias]
+    if (compiled) {
+      return compiled
+    }
+    const collection = this.collectionByAlias[alias]
+    if (collection) {
+      return collection.id
     }
+    throw new Error(`Unknown source alias "${alias}"`)
+  }
+
+  isLazyAlias(alias: string): boolean {
+    return this.lazySources.has(alias)
   }
 
   // The callback function is called after the graph has run.
@@ -115,12 +209,8 @@ export class CollectionConfigBuilder<
   // That can happen because even though we load N rows, the pipeline might filter some of these rows out
   // causing the orderBy operator to receive less than N rows or even no rows at all.
   // So this callback would notice that it doesn't have enough rows and load some more.
-  // The callback returns a boolean, when it's true it's done loading data.
-  maybeRunGraph(
-    config: SyncMethods<TResult>,
-    syncState: FullSyncState,
-    callback?: () => boolean
-  ) {
+  // The callback returns a boolean, when it's true it's done loading data and we can mark the collection as ready.
+  maybeRunGraph(callback?: () => boolean) {
     if (this.isGraphRunning) {
       // no nested runs of the graph
       // which is possible if the `callback`
@@ -128,10 +218,18 @@ export class CollectionConfigBuilder<
       return
     }
 
+    // Should only be called when sync is active
+    if (!this.currentSyncConfig || !this.currentSyncState) {
+      throw new Error(
+        `maybeRunGraph called without active sync session. This should not happen.`
+      )
+    }
+
     this.isGraphRunning = true
 
     try {
-      const { begin, commit } = config
+      const { begin, commit } = this.currentSyncConfig
+      const syncState = this.currentSyncState
 
       // Don't run if the live query is in an error state
       if (this.isInErrorState) {
@@ -152,7 +250,7 @@ export class CollectionConfigBuilder<
           commit()
           // After initial commit, check if we should mark ready
           // (in case all sources were already ready before we subscribed)
-          this.updateLiveQueryStatus(config)
+          this.updateLiveQueryStatus(this.currentSyncConfig)
         }
       }
     } finally {
@@ -160,6 +258,158 @@ export class CollectionConfigBuilder<
     }
   }
 
+  /**
+   * Schedules a graph run with the transaction-scoped scheduler.
+   * Ensures each builder runs at most once per transaction, with automatic dependency tracking
+   * to run parent queries before child queries. Outside a transaction, runs immediately.
+   *
+   * Multiple calls during a transaction are coalesced into a single execution.
+   * Dependencies are auto-discovered from subscribed live queries, or can be overridden.
+   * Load callbacks are combined when entries merge.
+   *
+   * Uses the current sync session's config and syncState from instance properties.
+   *
+   * @param callback - Optional callback to load more data if needed (returns true when done)
+   * @param options - Optional scheduling configuration
+   * @param options.contextId - Transaction ID to group work; defaults to active transaction
+   * @param options.jobId - Unique identifier for this job; defaults to this builder instance
+   * @param options.alias - Source alias that triggered this schedule; adds alias-specific dependencies
+   * @param options.dependencies - Explicit dependency list; overrides auto-discovered dependencies
+   */
+  scheduleGraphRun(
+    callback?: () => boolean,
+    options?: {
+      contextId?: SchedulerContextId
+      jobId?: unknown
+      alias?: string
+      dependencies?: Array<CollectionConfigBuilder<any, any>>
+    }
+  ) {
+    const contextId = options?.contextId ?? getActiveTransaction()?.id
+    // Use the builder instance as the job ID for deduplication. This is memory-safe
+    // because the scheduler's context Map is deleted after flushing (no long-term retention).
+    const jobId = options?.jobId ?? this
+    const dependentBuilders = (() => {
+      if (options?.dependencies) {
+        return options.dependencies
+      }
+
+      const deps = new Set(this.builderDependencies)
+      if (options?.alias) {
+        const aliasDeps = this.aliasDependencies[options.alias]
+        if (aliasDeps) {
+          for (const dep of aliasDeps) {
+            deps.add(dep)
+          }
+        }
+      }
+
+      deps.delete(this)
+
+      return Array.from(deps)
+    })()
+
+    // We intentionally scope deduplication to the builder instance. Each instance
+    // owns caches and compiled pipelines, so sharing work across instances that
+    // merely reuse the same string id would execute the wrong builder's graph.
+
+    if (!this.currentSyncConfig || !this.currentSyncState) {
+      throw new Error(
+        `scheduleGraphRun called without active sync session. This should not happen.`
+      )
+    }
+
+    // Manage our own state - get or create pending callbacks for this context
+    let pending = contextId ? this.pendingGraphRuns.get(contextId) : undefined
+    if (!pending) {
+      pending = {
+        loadCallbacks: new Set(),
+      }
+      if (contextId) {
+        this.pendingGraphRuns.set(contextId, pending)
+      }
+    }
+
+    // Add callback if provided (this is what accumulates between schedules)
+    if (callback) {
+      pending.loadCallbacks.add(callback)
+    }
+
+    // Schedule execution (scheduler just orchestrates order, we manage state)
+    // For immediate execution (no contextId), pass pending directly since it won't be in the map
+    const pendingToPass = contextId ? undefined : pending
+    transactionScopedScheduler.schedule({
+      contextId,
+      jobId,
+      dependencies: dependentBuilders,
+      run: () => this.executeGraphRun(contextId, pendingToPass),
+    })
+  }
+
+  /**
+   * Clears pending graph run state for a specific context.
+   * Called when the scheduler clears a context (e.g., transaction rollback/abort).
+   */
+  clearPendingGraphRun(contextId: SchedulerContextId): void {
+    this.pendingGraphRuns.delete(contextId)
+  }
+
+  /**
+   * Executes a pending graph run. Called by the scheduler when dependencies are satisfied.
+   * Clears the pending state BEFORE execution so that any re-schedules during the run
+   * create fresh state and don't interfere with the current execution.
+   * Uses instance sync state - if sync has ended, gracefully returns without executing.
+   *
+   * @param contextId - Optional context ID to look up pending state
+   * @param pendingParam - For immediate execution (no context), pending state is passed directly
+   */
+  private executeGraphRun(
+    contextId?: SchedulerContextId,
+    pendingParam?: PendingGraphRun
+  ): void {
+    // Get pending state: either from parameter (no context) or from map (with context)
+    // Remove from map BEFORE checking sync state to prevent leaking entries when sync ends
+    // before the transaction flushes (e.g., unsubscribe during in-flight transaction)
+    const pending =
+      pendingParam ??
+      (contextId ? this.pendingGraphRuns.get(contextId) : undefined)
+    if (contextId) {
+      this.pendingGraphRuns.delete(contextId)
+    }
+
+    // If no pending state, nothing to execute (context was cleared)
+    if (!pending) {
+      return
+    }
+
+    // If sync session has ended, don't execute (graph is finalized, subscriptions cleared)
+    if (!this.currentSyncConfig || !this.currentSyncState) {
+      return
+    }
+
+    this.incrementRunCount()
+
+    const combinedLoader = () => {
+      let allDone = true
+      let firstError: unknown
+      pending.loadCallbacks.forEach((loader) => {
+        try {
+          allDone = loader() && allDone
+        } catch (error) {
+          allDone = false
+          firstError ??= error
+        }
+      })
+      if (firstError) {
+        throw firstError
+      }
+      // Returning false signals that callers should schedule another pass.
+      return allDone
+    }
+
+    this.maybeRunGraph(combinedLoader)
+  }
+
   private getSyncConfig(): SyncConfig<TResult> {
     return {
       rowUpdateMode: `full`,
@@ -167,9 +417,19 @@ export class CollectionConfigBuilder<
     }
   }
 
+  incrementRunCount() {
+    this.runCount++
+  }
+
+  getRunCount() {
+    return this.runCount
+  }
+
   private syncFn(config: SyncMethods<TResult>) {
     // Store reference to the live query collection for error state transitions
     this.liveQueryCollection = config.collection
+    // Store config and syncState as instance properties for the duration of this sync session
+    this.currentSyncConfig = config
 
     const syncState: SyncState = {
       messagesCount: 0,
@@ -182,6 +442,15 @@ export class CollectionConfigBuilder<
       config,
       syncState
     )
+    this.currentSyncState = fullSyncState
+
+    // Listen for scheduler context clears to clean up our pending state
+    // Re-register on each sync start so the listener is active for the sync session's lifetime
+    this.unsubscribeFromSchedulerClears = transactionScopedScheduler.onClear(
+      (contextId) => {
+        this.clearPendingGraphRun(contextId)
+      }
+    )
 
     const loadMoreDataCallbacks = this.subscribeToAllCollections(
       config,
@@ -189,51 +458,81 @@ export class CollectionConfigBuilder<
     )
 
     // Initial run with callback to load more data if needed
-    this.maybeRunGraph(config, fullSyncState, loadMoreDataCallbacks)
+    this.scheduleGraphRun(loadMoreDataCallbacks)
 
     // Return the unsubscribe function
     return () => {
       syncState.unsubscribeCallbacks.forEach((unsubscribe) => unsubscribe())
 
+      // Clear current sync session state
+      this.currentSyncConfig = undefined
+      this.currentSyncState = undefined
+
+      // Clear all pending graph runs to prevent memory leaks from in-flight transactions
+      // that may flush after the sync session ends
+      this.pendingGraphRuns.clear()
+
       // Reset caches so a fresh graph/pipeline is compiled on next start
       // This avoids reusing a finalized D2 graph across GC restarts
       this.graphCache = undefined
       this.inputsCache = undefined
       this.pipelineCache = undefined
-      this.collectionWhereClausesCache = undefined
+      this.sourceWhereClausesCache = undefined
 
-      // Reset lazy collection state
-      this.lazyCollections.clear()
+      // Reset lazy source alias state
+      this.lazySources.clear()
       this.optimizableOrderByCollections = {}
-      this.lazyCollectionsCallbacks = {}
+      this.lazySourcesCallbacks = {}
+
+      // Clear subscription references to prevent memory leaks
+      // Note: Individual subscriptions are already unsubscribed via unsubscribeCallbacks
+      Object.keys(this.subscriptions).forEach(
+        (key) => delete this.subscriptions[key]
+      )
+      this.compiledAliasToCollectionId = {}
+
+      // Unregister from scheduler's onClear listener to prevent memory leaks
+      // The scheduler's listener Set would otherwise keep a strong reference to this builder
+      this.unsubscribeFromSchedulerClears?.()
+      this.unsubscribeFromSchedulerClears = undefined
     }
   }
 
+  /**
+   * Compiles the query pipeline with all declared aliases.
+   */
   private compileBasePipeline() {
     this.graphCache = new D2()
     this.inputsCache = Object.fromEntries(
-      Object.entries(this.collections).map(([key]) => [
-        key,
+      Object.keys(this.collectionByAlias).map((alias) => [
+        alias,
         this.graphCache!.newInput<any>(),
       ])
     )
 
-    // Compile the query and get both pipeline and collection WHERE clauses
-    const {
-      pipeline: pipelineCache,
-      collectionWhereClauses: collectionWhereClausesCache,
-    } = compileQuery(
+    const compilation = compileQuery(
       this.query,
       this.inputsCache as Record<string, KeyedStream>,
       this.collections,
       this.subscriptions,
-      this.lazyCollectionsCallbacks,
-      this.lazyCollections,
+      this.lazySourcesCallbacks,
+      this.lazySources,
       this.optimizableOrderByCollections
     )
 
-    this.pipelineCache = pipelineCache
-    this.collectionWhereClausesCache = collectionWhereClausesCache
+    this.pipelineCache = compilation.pipeline
+    this.sourceWhereClausesCache = compilation.sourceWhereClauses
+    this.compiledAliasToCollectionId = compilation.aliasToCollectionId
+
+    // Defensive check: verify all compiled aliases have corresponding inputs
+    // This should never happen since all aliases come from user declarations,
+    // but catch it early if the assumption is violated in the future.
+    const missingAliases = Object.keys(this.compiledAliasToCollectionId).filter(
+      (alias) => !Object.hasOwn(this.inputsCache!, alias)
+    )
+    if (missingAliases.length > 0) {
+      throw new MissingAliasInputsError(missingAliases)
+    }
   }
 
   private maybeCompileBasePipeline() {
@@ -400,44 +699,78 @@ export class CollectionConfigBuilder<
     )
   }
 
+  /**
+   * Creates per-alias subscriptions enabling self-join support.
+   * Each alias gets its own subscription with independent filters, even for the same collection.
+   * Example: `{ employee: col, manager: col }` creates two separate subscriptions.
+   */
   private subscribeToAllCollections(
     config: SyncMethods<TResult>,
     syncState: FullSyncState
   ) {
-    const loaders = Object.entries(this.collections).map(
-      ([collectionId, collection]) => {
-        const collectionSubscriber = new CollectionSubscriber(
-          collectionId,
-          collection,
-          config,
-          syncState,
-          this
-        )
-
-        const subscription = collectionSubscriber.subscribe()
-        this.subscriptions[collectionId] = subscription
-
-        // Subscribe to status changes for status flow
-        const statusUnsubscribe = collection.on(`status:change`, (event) => {
-          this.handleSourceStatusChange(config, collectionId, event)
-        })
-        syncState.unsubscribeCallbacks.add(statusUnsubscribe)
-
-        const loadMore = collectionSubscriber.loadMoreIfNeeded.bind(
-          collectionSubscriber,
-          subscription
-        )
-
-        return loadMore
+    // Use compiled aliases as the source of truth - these include all aliases from the query
+    // including those from subqueries, which may not be in collectionByAlias
+    const compiledAliases = Object.entries(this.compiledAliasToCollectionId)
+    if (compiledAliases.length === 0) {
+      throw new Error(
+        `Compiler returned no alias metadata for query '${this.id}'. This should not happen; please report.`
+      )
+    }
+
+    // Create a separate subscription for each alias, enabling self-joins where the same
+    // collection can be used multiple times with different filters and subscriptions
+    const loaders = compiledAliases.map(([alias, collectionId]) => {
+      // Try collectionByAlias first (for declared aliases), fall back to collections (for subquery aliases)
+      const collection =
+        this.collectionByAlias[alias] ?? this.collections[collectionId]!
+
+      const dependencyBuilder = getCollectionBuilder(collection)
+      if (dependencyBuilder && dependencyBuilder !== this) {
+        this.aliasDependencies[alias] = [dependencyBuilder]
+        this.builderDependencies.add(dependencyBuilder)
+      } else {
+        this.aliasDependencies[alias] = []
       }
-    )
 
+      // CollectionSubscriber handles the actual subscription to the source collection
+      // and feeds data into the D2 graph inputs for this specific alias
+      const collectionSubscriber = new CollectionSubscriber(
+        alias,
+        collectionId,
+        collection,
+        this
+      )
+
+      // Subscribe to status changes for status flow
+      const statusUnsubscribe = collection.on(`status:change`, (event) => {
+        this.handleSourceStatusChange(config, collectionId, event)
+      })
+      syncState.unsubscribeCallbacks.add(statusUnsubscribe)
+
+      const subscription = collectionSubscriber.subscribe()
+      // Store subscription by alias (not collection ID) to support lazy loading
+      // which needs to look up subscriptions by their query alias
+      this.subscriptions[alias] = subscription
+
+      // Create a callback for loading more data if needed (used by OrderBy optimization)
+      const loadMore = collectionSubscriber.loadMoreIfNeeded.bind(
+        collectionSubscriber,
+        subscription
+      )
+
+      return loadMore
+    })
+
+    // Combine all loaders into a single callback that initiates loading more data
+    // from any source that needs it. Returns true once all loaders have been called,
+    // but the actual async loading may still be in progress.
     const loadMoreDataCallback = () => {
       loaders.map((loader) => loader())
       return true
     }
 
-    // Mark the collections as subscribed in the sync state
+    // Mark as subscribed so the graph can start running
+    // (graph only runs when all collections are subscribed)
     syncState.subscribedToAllCollections = true
 
     // Initial status check after all subscriptions are set up
@@ -524,6 +857,65 @@ function extractCollectionsFromQuery(
   return collections
 }
 
+/**
+ * Extracts all aliases used for each collection across the entire query tree.
+ *
+ * Traverses the QueryIR recursively to build a map from collection ID to all aliases
+ * that reference that collection. This is essential for self-join support, where the
+ * same collection may be referenced multiple times with different aliases.
+ *
+ * For example, given a query like:
+ * ```ts
+ * q.from({ employee: employeesCollection })
+ *   .join({ manager: employeesCollection }, ({ employee, manager }) =>
+ *     eq(employee.managerId, manager.id)
+ *   )
+ * ```
+ *
+ * This function would return:
+ * ```
+ * Map { "employees" => Set { "employee", "manager" } }
+ * ```
+ *
+ * @param query - The query IR to extract aliases from
+ * @returns A map from collection ID to the set of all aliases referencing that collection
+ */
+function extractCollectionAliases(query: QueryIR): Map<string, Set<string>> {
+  const aliasesById = new Map<string, Set<string>>()
+
+  function recordAlias(source: any) {
+    if (!source) return
+
+    if (source.type === `collectionRef`) {
+      const { id } = source.collection
+      const existing = aliasesById.get(id)
+      if (existing) {
+        existing.add(source.alias)
+      } else {
+        aliasesById.set(id, new Set([source.alias]))
+      }
+    } else if (source.type === `queryRef`) {
+      traverse(source.query)
+    }
+  }
+
+  function traverse(q?: QueryIR) {
+    if (!q) return
+
+    recordAlias(q.from)
+
+    if (q.join) {
+      for (const joinClause of q.join) {
+        recordAlias(joinClause.from)
+      }
+    }
+  }
+
+  traverse(query)
+
+  return aliasesById
+}
+
 function accumulateChanges<T>(
   acc: Map<unknown, Changes<T>>,
   [[key, tupleData], multiplicity]: [