@tanstack/db 0.5.32 → 0.6.0

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

@@ -1,9 +1,15 @@
-import { MultiSet, serializeValue } from '@tanstack/db-ivm'
 import {
   normalizeExpressionPaths,
   normalizeOrderByPaths,
 } from '../compiler/expressions.js'
-import type { MultiSetArray, RootStreamBuilder } from '@tanstack/db-ivm'
+import {
+  computeOrderedLoadCursor,
+  computeSubscriptionOrderByHints,
+  filterDuplicateInserts,
+  sendChangesToInput,
+  splitUpdates,
+  trackBiggestSentValue,
+} from './utils.js'
 import type { Collection } from '../../collection/index.js'
 import type {
   ChangeMessage,
@@ -147,25 +153,11 @@ export class CollectionSubscriber<
     changes: Iterable<ChangeMessage<any, string | number>>,
     callback?: () => boolean,
   ) {
-    // Filter changes to prevent duplicate inserts to D2 pipeline.
-    // This ensures D2 multiplicity stays at 1 for visible items, so deletes
-    // properly reduce multiplicity to 0 (triggering DELETE output).
     const changesArray = Array.isArray(changes) ? changes : [...changes]
-    const filteredChanges: Array<ChangeMessage<any, string | number>> = []
-    for (const change of changesArray) {
-      if (change.type === `insert`) {
-        if (this.sentToD2Keys.has(change.key)) {
-          // Skip duplicate insert - already sent to D2
-          continue
-        }
-        this.sentToD2Keys.add(change.key)
-      } else if (change.type === `delete`) {
-        // Remove from tracking so future re-inserts are allowed
-        this.sentToD2Keys.delete(change.key)
-      }
-      // Updates are handled as delete+insert by splitUpdates, so no special handling needed
-      filteredChanges.push(change)
-    }
+    const filteredChanges = filterDuplicateInserts(
+      changesArray,
+      this.sentToD2Keys,
+    )
 
     // currentSyncState and input are always defined when this method is called
     // (only called from active subscriptions during a sync session)
@@ -202,27 +194,10 @@ export class CollectionSubscriber<
     }
 
     // Get the query's orderBy and limit to pass to loadSubset.
-    // Only include orderBy when it is scoped to this alias and uses simple refs,
-    // to avoid leaking cross-collection paths into backend-specific compilers.
-    const { orderBy, limit, offset } = this.collectionConfigBuilder.query
-    const effectiveLimit =
-      limit !== undefined && offset !== undefined ? limit + offset : limit
-    const normalizedOrderBy = orderBy
-      ? normalizeOrderByPaths(orderBy, this.alias)
-      : undefined
-    const canPassOrderBy =
-      normalizedOrderBy?.every((clause) => {
-        const exp = clause.expression
-        if (exp.type !== `ref`) {
-          return false
-        }
-        const path = exp.path
-        return Array.isArray(path) && path.length === 1
-      }) ?? false
-    const orderByForSubscription = canPassOrderBy
-      ? normalizedOrderBy
-      : undefined
-    const limitForSubscription = canPassOrderBy ? effectiveLimit : undefined
+    const hints = computeSubscriptionOrderByHints(
+      this.collectionConfigBuilder.query,
+      this.alias,
+    )
 
     // Track loading via the loadSubset promise directly.
     // requestSnapshot uses trackLoadSubsetPromise: false (needed for truncate handling),
@@ -241,8 +216,8 @@ export class CollectionSubscriber<
       ...(includeInitialState && { includeInitialState }),
       whereExpression,
       onStatusChange,
-      orderBy: orderByForSubscription,
-      limit: limitForSubscription,
+      orderBy: hints.orderBy,
+      limit: hints.limit,
       onLoadSubsetResult,
     })
 
@@ -415,52 +390,28 @@ export class CollectionSubscriber<
     if (!orderByInfo) {
       return
     }
-    const { orderBy, valueExtractorForRawRow, offset } = orderByInfo
-    const biggestSentRow = this.biggest
-
-    // Extract all orderBy column values from the biggest sent row
-    // For single-column: returns single value, for multi-column: returns array
-    const extractedValues = biggestSentRow
-      ? valueExtractorForRawRow(biggestSentRow)
-      : undefined
-
-    // Normalize to array format for minValues
-    let minValues: Array<unknown> | undefined
-    if (extractedValues !== undefined) {
-      minValues = Array.isArray(extractedValues)
-        ? extractedValues
-        : [extractedValues]
-    }
-
-    const loadRequestKey = this.getLoadRequestKey({
-      minValues,
-      offset,
-      limit: n,
-    })
 
-    // Skip if we already requested a load for this cursor+window.
-    // This prevents infinite loops from cached data re-writes while still allowing
-    // window moves (offset/limit changes) to trigger new requests.
-    if (this.lastLoadRequestKey === loadRequestKey) {
-      return
-    }
+    const cursor = computeOrderedLoadCursor(
+      orderByInfo,
+      this.biggest,
+      this.lastLoadRequestKey,
+      this.alias,
+      n,
+    )
+    if (!cursor) return // Duplicate request — skip
 
-    // Normalize the orderBy clauses such that the references are relative to the collection
-    const normalizedOrderBy = normalizeOrderByPaths(orderBy, this.alias)
+    this.lastLoadRequestKey = cursor.loadRequestKey
 
     // Take the `n` items after the biggest sent value
-    // Pass the current window offset to ensure proper deduplication
+    // Omit offset so requestLimitedSnapshot can advance based on
+    // the number of rows already loaded (supports offset-based backends).
     subscription.requestLimitedSnapshot({
-      orderBy: normalizedOrderBy,
+      orderBy: cursor.normalizedOrderBy,
       limit: n,
-      minValues,
-      // Omit offset so requestLimitedSnapshot can advance the offset based on
-      // the number of rows already loaded (supports offset-based backends).
+      minValues: cursor.minValues,
       trackLoadSubsetPromise: false,
       onLoadSubsetResult: this.orderedLoadSubsetResult,
     })
-
-    this.lastLoadRequestKey = loadRequestKey
   }
 
   private getWhereClauseForAlias(): BasicExpression<boolean> | undefined {
@@ -487,24 +438,15 @@ export class CollectionSubscriber<
     changes: Array<ChangeMessage<any, string | number>>,
     comparator: (a: any, b: any) => number,
   ): void {
-    for (const change of changes) {
-      if (change.type === `delete`) {
-        continue
-      }
-
-      const isNewKey = !this.sentToD2Keys.has(change.key)
-
-      // Only track inserts/updates for cursor positioning, not deletes
-      if (!this.biggest) {
-        this.biggest = change.value
-        this.lastLoadRequestKey = undefined
-      } else if (comparator(this.biggest, change.value) < 0) {
-        this.biggest = change.value
-        this.lastLoadRequestKey = undefined
-      } else if (isNewKey) {
-        // New key with same orderBy value - allow another load if needed
-        this.lastLoadRequestKey = undefined
-      }
+    const result = trackBiggestSentValue(
+      changes,
+      this.biggest,
+      this.sentToD2Keys,
+      comparator,
+    )
+    this.biggest = result.biggest
+    if (result.shouldResetLoadKey) {
+      this.lastLoadRequestKey = undefined
     }
   }
 
@@ -525,62 +467,4 @@ export class CollectionSubscriber<
       promise,
     )
   }
-
-  private getLoadRequestKey(options: {
-    minValues: Array<unknown> | undefined
-    offset: number
-    limit: number
-  }): string {
-    return serializeValue({
-      minValues: options.minValues ?? null,
-      offset: options.offset,
-      limit: options.limit,
-    })
-  }
-}
-
-/**
- * Helper function to send changes to a D2 input stream
- */
-function sendChangesToInput(
-  input: RootStreamBuilder<unknown>,
-  changes: Iterable<ChangeMessage>,
-  getKey: (item: ChangeMessage[`value`]) => any,
-): number {
-  const multiSetArray: MultiSetArray<unknown> = []
-  for (const change of changes) {
-    const key = getKey(change.value)
-    if (change.type === `insert`) {
-      multiSetArray.push([[key, change.value], 1])
-    } else if (change.type === `update`) {
-      multiSetArray.push([[key, change.previousValue], -1])
-      multiSetArray.push([[key, change.value], 1])
-    } else {
-      // change.type === `delete`
-      multiSetArray.push([[key, change.value], -1])
-    }
-  }
-
-  if (multiSetArray.length !== 0) {
-    input.sendData(new MultiSet(multiSetArray))
-  }
-
-  return multiSetArray.length
-}
-
-/** Splits updates into a delete of the old value and an insert of the new value */
-function* splitUpdates<
-  T extends object = Record<string, unknown>,
-  TKey extends string | number = string | number,
->(
-  changes: Iterable<ChangeMessage<T, TKey>>,
-): Generator<ChangeMessage<T, TKey>> {
-  for (const change of changes) {
-    if (change.type === `update`) {
-      yield { type: `delete`, key: change.key, value: change.previousValue! }
-      yield { type: `insert`, key: change.key, value: change.value }
-    } else {
-      yield change
-    }
-  }
 }

package/src/query/live/types.ts

@@ -5,7 +5,11 @@ import type {
   StringCollationConfig,
 } from '../../types.js'
 import type { InitialQueryBuilder, QueryBuilder } from '../builder/index.js'
-import type { Context, GetResult } from '../builder/types.js'
+import type {
+  Context,
+  RootObjectResultConstraint,
+  RootQueryResult,
+} from '../builder/types.js'
 
 export type Changes<T> = {
   deletes: number
@@ -54,7 +58,7 @@ export type FullSyncState = Required<Omit<SyncState, `flushPendingChanges`>> &
  */
 export interface LiveQueryCollectionConfig<
   TContext extends Context,
-  TResult extends object = GetResult<TContext> & object,
+  TResult extends object = RootQueryResult<TContext>,
 > {
   /**
    * Unique identifier for the collection
@@ -66,8 +70,10 @@ export interface LiveQueryCollectionConfig<
    * Query builder function that defines the live query
    */
   query:
-    | ((q: InitialQueryBuilder) => QueryBuilder<TContext>)
-    | QueryBuilder<TContext>
+    | ((
+        q: InitialQueryBuilder,
+      ) => QueryBuilder<TContext> & RootObjectResultConstraint<TContext>)
+    | (QueryBuilder<TContext> & RootObjectResultConstraint<TContext>)
 
   /**
    * Function to extract the key from result items

package/src/query/live/utils.ts

@@ -0,0 +1,417 @@
+import { MultiSet, serializeValue } from '@tanstack/db-ivm'
+import { UnsupportedRootScalarSelectError } from '../../errors.js'
+import { normalizeOrderByPaths } from '../compiler/expressions.js'
+import { buildQuery, getQueryIR } from '../builder/index.js'
+import { IncludesSubquery } from '../ir.js'
+import type { MultiSetArray, RootStreamBuilder } from '@tanstack/db-ivm'
+import type { Collection } from '../../collection/index.js'
+import type { ChangeMessage } from '../../types.js'
+import type { InitialQueryBuilder, QueryBuilder } from '../builder/index.js'
+import type { Context } from '../builder/types.js'
+import type { OrderBy, QueryIR } from '../ir.js'
+import type { OrderByOptimizationInfo } from '../compiler/order-by.js'
+
+/**
+ * Helper function to extract collections from a compiled query.
+ * Traverses the query IR to find all collection references.
+ * Maps collections by their ID (not alias) as expected by the compiler.
+ */
+export function extractCollectionsFromQuery(
+  query: any,
+): Record<string, Collection<any, any, any>> {
+  const collections: Record<string, any> = {}
+
+  // Helper function to recursively extract collections from a query or source
+  function extractFromSource(source: any) {
+    if (source.type === `collectionRef`) {
+      collections[source.collection.id] = source.collection
+    } else if (source.type === `queryRef`) {
+      // Recursively extract from subquery
+      extractFromQuery(source.query)
+    }
+  }
+
+  // Helper function to recursively extract collections from a query
+  function extractFromQuery(q: any) {
+    // Extract from FROM clause
+    if (q.from) {
+      extractFromSource(q.from)
+    }
+
+    // Extract from JOIN clauses
+    if (q.join && Array.isArray(q.join)) {
+      for (const joinClause of q.join) {
+        if (joinClause.from) {
+          extractFromSource(joinClause.from)
+        }
+      }
+    }
+
+    // Extract from SELECT (for IncludesSubquery)
+    if (q.select) {
+      extractFromSelect(q.select)
+    }
+  }
+
+  function extractFromSelect(select: any) {
+    for (const [key, value] of Object.entries(select)) {
+      if (typeof key === `string` && key.startsWith(`__SPREAD_SENTINEL__`)) {
+        continue
+      }
+      if (value instanceof IncludesSubquery) {
+        extractFromQuery(value.query)
+      } else if (isNestedSelectObject(value)) {
+        extractFromSelect(value)
+      }
+    }
+  }
+
+  // Start extraction from the root query
+  extractFromQuery(query)
+
+  return collections
+}
+
+/**
+ * Helper function to extract the collection that is referenced in the query's FROM clause.
+ * The FROM clause may refer directly to a collection or indirectly to a subquery.
+ */
+export function extractCollectionFromSource(
+  query: any,
+): Collection<any, any, any> {
+  const from = query.from
+
+  if (from.type === `collectionRef`) {
+    return from.collection
+  } else if (from.type === `queryRef`) {
+    // Recursively extract from subquery
+    return extractCollectionFromSource(from.query)
+  }
+
+  throw new Error(
+    `Failed to extract collection. Invalid FROM clause: ${JSON.stringify(query)}`,
+  )
+}
+
+/**
+ * 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
+ */
+export 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 traverseSelect(select: any) {
+    for (const [key, value] of Object.entries(select)) {
+      if (typeof key === `string` && key.startsWith(`__SPREAD_SENTINEL__`)) {
+        continue
+      }
+      if (value instanceof IncludesSubquery) {
+        traverse(value.query)
+      } else if (isNestedSelectObject(value)) {
+        traverseSelect(value)
+      }
+    }
+  }
+
+  function traverse(q?: QueryIR) {
+    if (!q) return
+
+    recordAlias(q.from)
+
+    if (q.join) {
+      for (const joinClause of q.join) {
+        recordAlias(joinClause.from)
+      }
+    }
+
+    if (q.select) {
+      traverseSelect(q.select)
+    }
+  }
+
+  traverse(query)
+
+  return aliasesById
+}
+
+/**
+ * Check if a value is a nested select object (plain object, not an expression)
+ */
+function isNestedSelectObject(obj: any): boolean {
+  if (obj === null || typeof obj !== `object`) return false
+  if (obj instanceof IncludesSubquery) return false
+  // Expression-like objects have a .type property
+  if (`type` in obj && typeof obj.type === `string`) return false
+  // Ref proxies from spread operations
+  if (obj.__refProxy) return false
+  return true
+}
+
+/**
+ * Builds a query IR from a config object that contains either a query builder
+ * function or a QueryBuilder instance.
+ */
+export function buildQueryFromConfig<TContext extends Context>(config: {
+  query:
+    | ((q: InitialQueryBuilder) => QueryBuilder<TContext>)
+    | QueryBuilder<TContext>
+  requireObjectResult?: boolean
+}): QueryIR {
+  // Build the query using the provided query builder function or instance
+  const query =
+    typeof config.query === `function`
+      ? buildQuery<TContext>(config.query)
+      : getQueryIR(config.query)
+
+  if (
+    config.requireObjectResult &&
+    query.select &&
+    !isNestedSelectObject(query.select)
+  ) {
+    throw new UnsupportedRootScalarSelectError()
+  }
+
+  return query
+}
+
+/**
+ * Helper function to send changes to a D2 input stream.
+ * Converts ChangeMessages to D2 MultiSet data and sends to the input.
+ *
+ * @returns The number of multiset entries sent
+ */
+export function sendChangesToInput(
+  input: RootStreamBuilder<unknown>,
+  changes: Iterable<ChangeMessage>,
+  getKey: (item: ChangeMessage[`value`]) => any,
+): number {
+  const multiSetArray: MultiSetArray<unknown> = []
+  for (const change of changes) {
+    const key = getKey(change.value)
+    if (change.type === `insert`) {
+      multiSetArray.push([[key, change.value], 1])
+    } else if (change.type === `update`) {
+      multiSetArray.push([[key, change.previousValue], -1])
+      multiSetArray.push([[key, change.value], 1])
+    } else {
+      // change.type === `delete`
+      multiSetArray.push([[key, change.value], -1])
+    }
+  }
+
+  if (multiSetArray.length !== 0) {
+    input.sendData(new MultiSet(multiSetArray))
+  }
+
+  return multiSetArray.length
+}
+
+/** Splits updates into a delete of the old value and an insert of the new value */
+export function* splitUpdates<
+  T extends object = Record<string, unknown>,
+  TKey extends string | number = string | number,
+>(
+  changes: Iterable<ChangeMessage<T, TKey>>,
+): Generator<ChangeMessage<T, TKey>> {
+  for (const change of changes) {
+    if (change.type === `update`) {
+      yield { type: `delete`, key: change.key, value: change.previousValue! }
+      yield { type: `insert`, key: change.key, value: change.value }
+    } else {
+      yield change
+    }
+  }
+}
+
+/**
+ * Filter changes to prevent duplicate inserts to a D2 pipeline.
+ * Maintains D2 multiplicity at 1 for visible items so that deletes
+ * properly reduce multiplicity to 0.
+ *
+ * Mutates `sentKeys` in place: adds keys on insert, removes on delete.
+ */
+export function filterDuplicateInserts(
+  changes: Array<ChangeMessage<any, string | number>>,
+  sentKeys: Set<string | number>,
+): Array<ChangeMessage<any, string | number>> {
+  const filtered: Array<ChangeMessage<any, string | number>> = []
+  for (const change of changes) {
+    if (change.type === `insert`) {
+      if (sentKeys.has(change.key)) {
+        continue // Skip duplicate
+      }
+      sentKeys.add(change.key)
+    } else if (change.type === `delete`) {
+      sentKeys.delete(change.key)
+    }
+    filtered.push(change)
+  }
+  return filtered
+}
+
+/**
+ * Track the biggest value seen in a stream of changes, used for cursor-based
+ * pagination in ordered subscriptions. Returns whether the load request key
+ * should be reset (allowing another load).
+ *
+ * @param changes   - changes to process (deletes are skipped)
+ * @param current   - the current biggest value (or undefined if none)
+ * @param sentKeys  - set of keys already sent to D2 (for new-key detection)
+ * @param comparator - orderBy comparator
+ * @returns `{ biggest, shouldResetLoadKey }` — the new biggest value and
+ *          whether the caller should clear its last-load-request-key
+ */
+export function trackBiggestSentValue(
+  changes: Array<ChangeMessage<any, string | number>>,
+  current: unknown | undefined,
+  sentKeys: Set<string | number>,
+  comparator: (a: any, b: any) => number,
+): { biggest: unknown; shouldResetLoadKey: boolean } {
+  let biggest = current
+  let shouldResetLoadKey = false
+
+  for (const change of changes) {
+    if (change.type === `delete`) continue
+
+    const isNewKey = !sentKeys.has(change.key)
+
+    if (biggest === undefined) {
+      biggest = change.value
+      shouldResetLoadKey = true
+    } else if (comparator(biggest, change.value) < 0) {
+      biggest = change.value
+      shouldResetLoadKey = true
+    } else if (isNewKey) {
+      // New key at same sort position — allow another load if needed
+      shouldResetLoadKey = true
+    }
+  }
+
+  return { biggest, shouldResetLoadKey }
+}
+
+/**
+ * Compute orderBy/limit subscription hints for an alias.
+ * Returns normalised orderBy and effective limit suitable for passing to
+ * `subscribeChanges`, or `undefined` values when the query's orderBy cannot
+ * be scoped to the given alias (e.g. cross-collection refs or aggregates).
+ */
+export function computeSubscriptionOrderByHints(
+  query: { orderBy?: OrderBy; limit?: number; offset?: number },
+  alias: string,
+): { orderBy: OrderBy | undefined; limit: number | undefined } {
+  const { orderBy, limit, offset } = query
+  const effectiveLimit =
+    limit !== undefined && offset !== undefined ? limit + offset : limit
+
+  const normalizedOrderBy = orderBy
+    ? normalizeOrderByPaths(orderBy, alias)
+    : undefined
+
+  // Only pass orderBy when it is scoped to this alias and uses simple refs,
+  // to avoid leaking cross-collection paths into backend-specific compilers.
+  const canPassOrderBy =
+    normalizedOrderBy?.every((clause) => {
+      const exp = clause.expression
+      if (exp.type !== `ref`) return false
+      const path = exp.path
+      return Array.isArray(path) && path.length === 1
+    }) ?? false
+
+  return {
+    orderBy: canPassOrderBy ? normalizedOrderBy : undefined,
+    limit: canPassOrderBy ? effectiveLimit : undefined,
+  }
+}
+
+/**
+ * Compute the cursor for loading the next batch of ordered data.
+ * Extracts values from the biggest sent row and builds the `minValues`
+ * array and a deduplication key.
+ *
+ * @returns `undefined` if the load should be skipped (duplicate request),
+ *          otherwise `{ minValues, normalizedOrderBy, loadRequestKey }`.
+ */
+export function computeOrderedLoadCursor(
+  orderByInfo: Pick<
+    OrderByOptimizationInfo,
+    'orderBy' | 'valueExtractorForRawRow' | 'offset'
+  >,
+  biggestSentRow: unknown | undefined,
+  lastLoadRequestKey: string | undefined,
+  alias: string,
+  limit: number,
+):
+  | {
+      minValues: Array<unknown> | undefined
+      normalizedOrderBy: OrderBy
+      loadRequestKey: string
+    }
+  | undefined {
+  const { orderBy, valueExtractorForRawRow, offset } = orderByInfo
+
+  // Extract all orderBy column values from the biggest sent row
+  // For single-column: returns single value, for multi-column: returns array
+  const extractedValues = biggestSentRow
+    ? valueExtractorForRawRow(biggestSentRow as Record<string, unknown>)
+    : undefined
+
+  // Normalize to array format for minValues
+  let minValues: Array<unknown> | undefined
+  if (extractedValues !== undefined) {
+    minValues = Array.isArray(extractedValues)
+      ? extractedValues
+      : [extractedValues]
+  }
+
+  // Deduplicate: skip if we already issued an identical load request
+  const loadRequestKey = serializeValue({
+    minValues: minValues ?? null,
+    offset,
+    limit,
+  })
+  if (lastLoadRequestKey === loadRequestKey) {
+    return undefined
+  }
+
+  const normalizedOrderBy = normalizeOrderByPaths(orderBy, alias)
+
+  return { minValues, normalizedOrderBy, loadRequestKey }
+}