@tanstack/db/dist/esm/query/live/collection-config-builder.js

A reactive client store for building super fast apps on sync

import { D2, output } from "@tanstack/db-ivm";
import { compileQuery } from "../compiler/index.js";
import { SetWindowRequiresOrderByError, MissingAliasInputsError } from "../../errors.js";
import { getActivePublicationContext, transactionScopedScheduler } from "../../scheduler.js";
import { getActiveTransaction } from "../../transactions.js";
import { deepEquals } from "../../utils.js";
import { CollectionSubscriber } from "./collection-subscriber.js";
import { getCollectionBuilder } from "./collection-registry.js";
import { LIVE_QUERY_INTERNAL } from "./internal.js";
import { materializeCompilation } from "./materialized-pipeline.js";
import { BucketFacadeAdapter } from "./bucket-facade-adapter.js";
import { buildQueryFromConfig, extractCollectionsFromQuery, extractCollectionFromSource } from "./utils.js";
import { collectCollectionSources } from "../ir.js";
let liveQueryCollectionCounter = 0;
class CollectionConfigBuilder {
  constructor(config) {
    this.config = config;
    this.compiledAliasToCollectionId = {};
    this.resultKeys = /* @__PURE__ */ new WeakMap();
    this.orderByIndices = /* @__PURE__ */ new WeakMap();
    this.isGraphRunning = false;
    this.runCount = 0;
    this.isInErrorState = false;
    this.sourceDependencies = {};
    this.builderDependencies = /* @__PURE__ */ new Set();
    this.pendingGraphRuns = /* @__PURE__ */ new Map();
    this.subscriptions = {};
    this.lazySourcesCallbacks = {};
    this.lazySources = /* @__PURE__ */ new Set();
    this.activeDemands = /* @__PURE__ */ new Map();
    this.demandGenerations = /* @__PURE__ */ new Map();
    this.optimizableOrderByCollections = {};
    this.id = config.id || `live-query-${++liveQueryCollectionCounter}`;
    this.query = buildQueryFromConfig({
      query: config.query,
      requireObjectResult: true
    });
    this.initialWindow = this.query.orderBy?.length ? {
      offset: this.query.offset ?? 0,
      limit: this.query.limit ?? Infinity
    } : void 0;
    this.collections = extractCollectionsFromQuery(this.query);
    this.collectionSources = collectCollectionSources(this.query);
    this.collectionByAlias = Object.fromEntries(
      this.collectionSources.map(({ alias, collection }) => [
        alias,
        collection
      ])
    );
    if (this.query.orderBy && this.query.orderBy.length > 0) {
      this.compare = createOrderByComparator(this.orderByIndices);
    }
    this.compareOptions = this.config.defaultStringCollation ?? extractCollectionFromSource(this.query).compareOptions;
    this.compileBasePipeline();
  }
  /**
   * Recursively checks if a query or any of its subqueries contains joins
   */
  hasJoins(query) {
    if (query.join && query.join.length > 0) {
      return true;
    }
    if (query.from.type === `queryRef`) {
      if (this.hasJoins(query.from.query)) {
        return true;
      }
    } else if (query.from.type === `unionFrom`) {
      for (const source of query.from.sources) {
        if (source.type === `queryRef` && this.hasJoins(source.query)) {
          return true;
        }
      }
    } else if (query.from.type === `unionAll`) {
      for (const branch of query.from.queries) {
        if (this.hasJoins(branch)) {
          return true;
        }
      }
    }
    return false;
  }
  getConfig() {
    return {
      id: this.id,
      getKey: this.config.getKey || ((item) => this.resultKeys.get(item) ?? item.$key),
      sync: this.getSyncConfig(),
      compare: this.compare,
      defaultStringCollation: this.compareOptions,
      gcTime: this.config.gcTime ?? 5e3,
      // 5 seconds by default for live queries
      schema: this.config.schema,
      onInsert: this.config.onInsert,
      onUpdate: this.config.onUpdate,
      onDelete: this.config.onDelete,
      startSync: this.config.startSync,
      singleResult: this.query.singleResult,
      utils: {
        getRunCount: this.getRunCount.bind(this),
        setWindow: this.setWindow.bind(this),
        getWindow: this.getWindow.bind(this),
        [LIVE_QUERY_INTERNAL]: {
          getBuilder: () => this,
          hasCustomGetKey: !!this.config.getKey,
          hasJoins: this.hasJoins(this.query),
          hasDistinct: !!this.query.distinct
        }
      }
    };
  }
  setWindow(options) {
    if (!this.windowFn) {
      throw new SetWindowRequiresOrderByError();
    }
    const previousWindow = this.currentWindow ?? this.initialWindow;
    try {
      this.windowFn(options);
      this.maybeRunGraphFn?.();
      this.currentWindow = options;
    } catch (error) {
      if (previousWindow) {
        try {
          this.windowFn(previousWindow);
          this.maybeRunGraphFn?.();
        } catch {
        }
      }
      throw error;
    }
    return this.liveQueryCollection?._sync.waitForCurrentLoadSubset() ?? true;
  }
  getWindow() {
    const window = this.currentWindow ?? this.initialWindow;
    if (!this.windowFn || !window) {
      return void 0;
    }
    return {
      offset: window.offset ?? 0,
      limit: window.limit ?? 0
    };
  }
  /**
   * 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) {
    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}"`);
  }
  isLazySource(sourceId) {
    return this.lazySources.has(sourceId);
  }
  beginDemand(planId) {
    const generation = (this.demandGenerations.get(planId) ?? 0) + 1;
    this.demandGenerations.set(planId, generation);
    this.activeDemands.set(planId, { generation, settled: false });
    return generation;
  }
  settleDemand(planId, generation) {
    const demand = this.activeDemands.get(planId);
    if (!demand || demand.generation !== generation || demand.settled) return;
    demand.settled = true;
    this.maybeRunGraphFn?.();
  }
  failDemand(planId, generation, error) {
    const demand = this.activeDemands.get(planId);
    if (!demand || demand.generation !== generation) return;
    const message = error instanceof Error ? error.message : String(error);
    this.transitionToError(`Subset demand '${planId}' failed: ${message}`);
  }
  retireDemand(planId) {
    this.activeDemands.delete(planId);
  }
  // The callback function is called after the graph has run.
  // This gives the callback a chance to load more data if needed,
  // that's used to optimize orderBy operators that set a limit,
  // in order to load some more data if we still don't have enough rows after the pipeline has run.
  // 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 and we can mark the collection as ready.
  maybeRunGraph(callback) {
    if (this.isGraphRunning) {
      return;
    }
    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 } = this.currentSyncConfig;
      const syncState = this.currentSyncState;
      if (this.isInErrorState) {
        return;
      }
      if (syncState.subscribedToAllCollections) {
        let callbackCalled = false;
        while (syncState.graph.pendingWork()) {
          syncState.graph.run();
          callback?.();
          callbackCalled = true;
        }
        syncState.flushPendingChanges?.();
        if (!callbackCalled) {
          callback?.();
        }
        if (syncState.messagesCount === 0) {
          begin();
          commit();
        }
        this.updateLiveQueryStatus(this.currentSyncConfig);
      }
    } finally {
      this.isGraphRunning = false;
    }
  }
  /**
   * 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.sourceId - Source that triggered this schedule; adds its dependencies
   * @param options.dependencies - Explicit dependency list; overrides auto-discovered dependencies
   */
  scheduleGraphRun(callback, options) {
    const contextId = options?.contextId ?? getActiveTransaction()?.id ?? getActivePublicationContext();
    const jobId = options?.jobId ?? this;
    const dependentBuilders = (() => {
      if (options?.dependencies) {
        return options.dependencies;
      }
      const deps = new Set(this.builderDependencies);
      if (options?.sourceId) {
        const sourceDeps = this.sourceDependencies[options.sourceId];
        if (sourceDeps) {
          for (const dep of sourceDeps) {
            deps.add(dep);
          }
        }
      }
      deps.delete(this);
      return Array.from(deps);
    })();
    if (contextId) {
      for (const dep of dependentBuilders) {
        if (typeof dep.scheduleGraphRun === `function`) {
          dep.scheduleGraphRun(void 0, { contextId });
        }
      }
    }
    if (!this.currentSyncConfig || !this.currentSyncState) {
      throw new Error(
        `scheduleGraphRun called without active sync session. This should not happen.`
      );
    }
    let pending = contextId ? this.pendingGraphRuns.get(contextId) : void 0;
    if (!pending) {
      pending = {
        loadCallbacks: /* @__PURE__ */ new Set()
      };
      if (contextId) {
        this.pendingGraphRuns.set(contextId, pending);
      }
    }
    if (callback) {
      pending.loadCallbacks.add(callback);
    }
    const pendingToPass = contextId ? void 0 : 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) {
    this.pendingGraphRuns.delete(contextId);
  }
  /**
   * Returns true if this builder has a pending graph run for the given context.
   */
  hasPendingGraphRun(contextId) {
    return this.pendingGraphRuns.has(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
   */
  executeGraphRun(contextId, pendingParam) {
    const pending = pendingParam ?? (contextId ? this.pendingGraphRuns.get(contextId) : void 0);
    if (contextId) {
      this.pendingGraphRuns.delete(contextId);
    }
    if (!pending) {
      return;
    }
    if (!this.currentSyncConfig || !this.currentSyncState) {
      return;
    }
    this.incrementRunCount();
    const combinedLoader = () => {
      let allDone = true;
      let firstError;
      pending.loadCallbacks.forEach((loader) => {
        try {
          allDone = loader() && allDone;
        } catch (error) {
          allDone = false;
          firstError ??= error;
        }
      });
      if (firstError) {
        throw firstError;
      }
      return allDone;
    };
    this.maybeRunGraph(combinedLoader);
  }
  getSyncConfig() {
    return {
      rowUpdateMode: `full`,
      sync: this.syncFn.bind(this)
    };
  }
  incrementRunCount() {
    this.runCount++;
  }
  getRunCount() {
    return this.runCount;
  }
  syncFn(config) {
    this.liveQueryCollection = config.collection;
    this.isInErrorState = false;
    this.currentSyncConfig = config;
    const syncState = {
      messagesCount: 0,
      subscribedToAllCollections: false,
      unsubscribeCallbacks: /* @__PURE__ */ new Set()
    };
    const fullSyncState = this.extendPipelineWithChangeProcessing(
      config,
      syncState
    );
    this.currentSyncState = fullSyncState;
    this.unsubscribeFromSchedulerClears = transactionScopedScheduler.onClear(
      (contextId) => {
        this.clearPendingGraphRun(contextId);
      }
    );
    const loadingSubsetUnsubscribe = config.collection.on(
      `loadingSubset:change`,
      (event) => {
        if (!event.isLoadingSubset) {
          this.updateLiveQueryStatus(config);
        }
      }
    );
    syncState.unsubscribeCallbacks.add(loadingSubsetUnsubscribe);
    const loadSubsetDataCallbacks = this.subscribeToAllCollections(
      config,
      fullSyncState
    );
    this.maybeRunGraphFn = () => this.scheduleGraphRun(loadSubsetDataCallbacks);
    this.scheduleGraphRun(loadSubsetDataCallbacks);
    return () => {
      syncState.unsubscribeCallbacks.forEach((unsubscribe) => unsubscribe());
      this.currentSyncConfig = void 0;
      this.currentSyncState = void 0;
      this.maybeRunGraphFn = void 0;
      this.currentWindow = void 0;
      this.pendingGraphRuns.clear();
      this.graphCache = void 0;
      this.inputsCache = void 0;
      this.pipelineCache = void 0;
      this.sourceWhereClausesCache = void 0;
      this.bucketFacadesCache = void 0;
      this.lazySources.clear();
      this.demandGenerations.clear();
      this.activeDemands.clear();
      this.optimizableOrderByCollections = {};
      this.lazySourcesCallbacks = {};
      Object.keys(this.subscriptions).forEach(
        (key) => delete this.subscriptions[key]
      );
      this.compiledAliasToCollectionId = {};
      this.unsubscribeFromSchedulerClears?.();
      this.unsubscribeFromSchedulerClears = void 0;
    };
  }
  /**
   * Compiles the query pipeline with all declared aliases.
   */
  compileBasePipeline() {
    this.graphCache = new D2();
    this.inputsCache = Object.fromEntries(
      this.collectionSources.map((source) => [
        source.sourceId,
        this.graphCache.newInput()
      ])
    );
    const compilation = compileQuery(
      this.query,
      this.inputsCache,
      this.collections,
      this.subscriptions,
      this.lazySourcesCallbacks,
      this.lazySources,
      this.optimizableOrderByCollections,
      (windowFn) => {
        this.windowFn = windowFn;
        if (this.currentWindow) {
          windowFn(this.currentWindow);
        }
      }
    );
    const materialized = materializeCompilation(
      compilation,
      this.config.getKey,
      this.hasJoins(this.query)
    );
    this.pipelineCache = materialized.pipeline;
    this.sourceWhereClausesCache = compilation.sourceWhereClauses;
    this.compiledAliasToCollectionId = compilation.aliasToCollectionId;
    this.bucketFacadesCache = materialized.facades;
    const missingSources = this.collectionSources.map((source) => source.sourceId).filter((sourceId) => !Object.hasOwn(this.inputsCache, sourceId));
    if (missingSources.length > 0) {
      throw new MissingAliasInputsError(missingSources);
    }
  }
  maybeCompileBasePipeline() {
    if (!this.graphCache || !this.inputsCache || !this.pipelineCache) {
      this.compileBasePipeline();
    }
    return {
      graph: this.graphCache,
      inputs: this.inputsCache,
      pipeline: this.pipelineCache
    };
  }
  extendPipelineWithChangeProcessing(config, syncState) {
    const { begin, commit } = config;
    const { graph, inputs, pipeline } = this.maybeCompileBasePipeline();
    let pendingChanges = /* @__PURE__ */ new Map();
    pipeline.pipe(
      output((data) => {
        const messages = data.getInner();
        syncState.messagesCount += messages.length;
        messages.reduce(accumulateChanges, pendingChanges);
      })
    );
    const bucketFacades = new BucketFacadeAdapter(
      this.id,
      this.bucketFacadesCache ?? [],
      (count) => {
        syncState.messagesCount += count;
      }
    );
    syncState.unsubscribeCallbacks.add(() => bucketFacades.cleanup());
    syncState.flushPendingChanges = () => {
      const hasParentChanges = pendingChanges.size > 0;
      const hasChildChanges = bucketFacades.hasPendingChanges();
      if (!hasParentChanges && !hasChildChanges) {
        return;
      }
      let facadePublication;
      let rootPublication;
      try {
        facadePublication = bucketFacades.flush();
        rootPublication = hasParentChanges ? config.collection._deferPublication() : void 0;
        const changesToApply = new Map(
          [...pendingChanges].map(([key, changes]) => {
            const resolved = {
              ...changes,
              value: bucketFacades.resolve(changes.value)
            };
            if (changes.previousValue !== void 0) {
              resolved.previousValue = bucketFacades.resolve(
                changes.previousValue
              );
            }
            return [key, resolved];
          })
        );
        if (hasParentChanges) {
          begin();
          changesToApply.forEach(this.applyChanges.bind(this, config));
          if (hasOrderOnlyMove(changesToApply)) {
            markLayoutChange(config.collection);
          }
          commit();
        }
      } catch (error) {
        pendingChanges = /* @__PURE__ */ new Map();
        rootPublication?.discard();
        facadePublication?.rollback();
        throw error;
      }
      pendingChanges = /* @__PURE__ */ new Map();
      let publicationError;
      for (const publish of [
        rootPublication?.publish,
        facadePublication.publish
      ]) {
        if (!publish) continue;
        try {
          publish();
        } catch (error) {
          publicationError ??= error;
        }
      }
      if (publicationError !== void 0) throw publicationError;
    };
    graph.finalize();
    syncState.graph = graph;
    syncState.inputs = inputs;
    syncState.pipeline = pipeline;
    return syncState;
  }
  applyChanges(config, changes, key) {
    const { write, collection } = config;
    const { deletes, inserts, value, orderByIndex } = changes;
    this.resultKeys.set(value, key);
    if (orderByIndex !== void 0) {
      this.orderByIndices.set(value, orderByIndex);
    }
    if (inserts && deletes === 0) {
      write({
        value,
        type: `insert`
      });
    } else if (
      // Insert & update(s) (updates are a delete & insert)
      inserts > deletes || // Just update(s) but the item is already in the collection (so
      // was inserted previously).
      inserts === deletes && collection.has(collection.getKeyFromItem(value))
    ) {
      write({
        value,
        type: `update`
      });
    } else if (deletes > 0) {
      write({
        value,
        type: `delete`
      });
    } else {
      throw new Error(
        `Could not apply changes: ${JSON.stringify(changes)}. This should never happen.`
      );
    }
  }
  /**
   * Handle status changes from source collections
   */
  handleSourceStatusChange(config, collectionId, event) {
    const { status } = event;
    if (status === `error`) {
      this.transitionToError(
        `Source collection '${collectionId}' entered error state`
      );
      return;
    }
    if (status === `cleaned-up`) {
      this.transitionToError(
        `Source collection '${collectionId}' was manually cleaned up while live query '${this.id}' depends on it. Live queries prevent automatic GC, so this was likely a manual cleanup() call.`
      );
      return;
    }
    this.updateLiveQueryStatus(config);
  }
  /**
   * Update the live query status based on source collection statuses
   */
  updateLiveQueryStatus(config) {
    const { markReady } = config;
    if (this.isInErrorState) {
      return;
    }
    const subscribedToAll = this.currentSyncState?.subscribedToAllCollections;
    const allReady = this.allRequiredSourcesReady();
    const allDemandsSettled = [...this.activeDemands.values()].every(
      (demand) => demand.settled
    );
    const isLoading = this.liveQueryCollection?.isLoadingSubset;
    if (subscribedToAll && allReady && allDemandsSettled && !isLoading) {
      markReady();
    }
  }
  /**
   * Transition the live query to error state
   */
  transitionToError(message) {
    this.isInErrorState = true;
    console.error(`[Live Query Error] ${message}`);
    this.liveQueryCollection?._lifecycle.setStatus(`error`);
  }
  allRequiredSourcesReady() {
    return this.collectionSources.every(
      (source) => (
        // Only on-demand sources settle through route demand. Eager
        // loadSubset calls return immediately, so they must reach ready.
        this.lazySources.has(source.sourceId) && source.collection.config.syncMode === `on-demand` || source.collection.isReady()
      )
    );
  }
  /**
   * Creates one subscription per lexical collection source.
   * Each source gets independent filters, even when aliases or collections repeat.
   * Example: `{ employee: col, manager: col }` creates two separate subscriptions.
   */
  subscribeToAllCollections(config, syncState) {
    if (this.collectionSources.length === 0) {
      throw new Error(
        `Query '${this.id}' has no collection sources. This should not happen; please report.`
      );
    }
    const loaders = this.collectionSources.map((source) => {
      const { sourceId, alias, collection } = source;
      const collectionId = collection.id;
      const dependencyBuilder = getCollectionBuilder(collection);
      if (dependencyBuilder && dependencyBuilder !== this) {
        this.sourceDependencies[sourceId] = [dependencyBuilder];
        this.builderDependencies.add(dependencyBuilder);
      } else {
        this.sourceDependencies[sourceId] = [];
      }
      const collectionSubscriber = new CollectionSubscriber(
        sourceId,
        alias,
        collection,
        this
      );
      const statusUnsubscribe = collection.on(`status:change`, (event) => {
        this.handleSourceStatusChange(config, collectionId, event);
      });
      syncState.unsubscribeCallbacks.add(statusUnsubscribe);
      const subscription = collectionSubscriber.subscribe();
      this.subscriptions[sourceId] = subscription;
      const lazyCallbacks = this.lazySourcesCallbacks[sourceId];
      if (lazyCallbacks) {
        lazyCallbacks.setDemand = (plan, keys) => collectionSubscriber.setDemand(subscription, plan, keys);
        for (const plan of lazyCallbacks.plans ?? []) {
          if (plan.initialKeys.size > 0) {
            lazyCallbacks.setDemand(plan, plan.initialKeys);
          }
        }
      }
      const loadMore = collectionSubscriber.loadMoreIfNeeded.bind(
        collectionSubscriber,
        subscription
      );
      return loadMore;
    });
    const loadSubsetDataCallbacks = () => {
      loaders.map((loader) => loader());
      return true;
    };
    syncState.subscribedToAllCollections = true;
    return loadSubsetDataCallbacks;
  }
}
function createOrderByComparator(orderByIndices) {
  return (val1, val2) => {
    const index1 = orderByIndices.get(val1);
    const index2 = orderByIndices.get(val2);
    if (index1 && index2) {
      if (index1 < index2) {
        return -1;
      } else if (index1 > index2) {
        return 1;
      } else {
        return 0;
      }
    }
    return 0;
  };
}
function accumulateChanges(acc, [[key, tupleData], multiplicity]) {
  const [value, orderByIndex] = tupleData;
  const changes = acc.get(key) || {
    deletes: 0,
    inserts: 0,
    value,
    orderByIndex
  };
  if (multiplicity < 0) {
    changes.deletes += Math.abs(multiplicity);
    changes.previousValue = value;
    changes.previousOrderByIndex = orderByIndex;
  } else if (multiplicity > 0) {
    changes.inserts += multiplicity;
    changes.value = value;
    if (orderByIndex !== void 0) {
      changes.orderByIndex = orderByIndex;
    }
  }
  acc.set(key, changes);
  return acc;
}
function hasOrderOnlyMove(changesToApply) {
  for (const changes of changesToApply.values()) {
    const isUpdate = changes.inserts > 0 && changes.deletes > 0;
    if (isUpdate && changes.previousValue !== void 0 && deepEquals(changes.previousValue, changes.value) && changes.orderByIndex !== changes.previousOrderByIndex) {
      return true;
    }
  }
  return false;
}
function markLayoutChange(collection) {
  collection._markLayoutChange();
}
export {
  CollectionConfigBuilder
};
//# sourceMappingURL=collection-config-builder.js.map