@tanstack/db 0.3.2 → 0.4.1

package/src/collection/subscription.ts

@@ -0,0 +1,239 @@
+import { ensureIndexForExpression } from "../indexes/auto-index.js"
+import { and } from "../query/index.js"
+import {
+  createFilterFunctionFromExpression,
+  createFilteredCallback,
+} from "./change-events.js"
+import type { BasicExpression } from "../query/ir.js"
+import type { BaseIndex } from "../indexes/base-index.js"
+import type { ChangeMessage } from "../types.js"
+import type { CollectionImpl } from "./index.js"
+
+type RequestSnapshotOptions = {
+  where?: BasicExpression<boolean>
+  optimizedOnly?: boolean
+}
+
+type RequestLimitedSnapshotOptions = {
+  minValue?: any
+  limit: number
+}
+
+type CollectionSubscriptionOptions = {
+  /** Pre-compiled expression for filtering changes */
+  whereExpression?: BasicExpression<boolean>
+  /** Callback to call when the subscription is unsubscribed */
+  onUnsubscribe?: () => void
+}
+
+export class CollectionSubscription {
+  private loadedInitialState = false
+
+  // Flag to indicate that we have sent at least 1 snapshot.
+  // While `snapshotSent` is false we filter out all changes from subscription to the collection.
+  private snapshotSent = false
+
+  // Keep track of the keys we've sent (needed for join and orderBy optimizations)
+  private sentKeys = new Set<string | number>()
+
+  private filteredCallback: (changes: Array<ChangeMessage<any, any>>) => void
+
+  private orderByIndex: BaseIndex<string | number> | undefined
+
+  constructor(
+    private collection: CollectionImpl<any, any, any, any, any>,
+    private callback: (changes: Array<ChangeMessage<any, any>>) => void,
+    private options: CollectionSubscriptionOptions
+  ) {
+    // Auto-index for where expressions if enabled
+    if (options.whereExpression) {
+      ensureIndexForExpression(options.whereExpression, this.collection)
+    }
+
+    const callbackWithSentKeysTracking = (
+      changes: Array<ChangeMessage<any, any>>
+    ) => {
+      callback(changes)
+      this.trackSentKeys(changes)
+    }
+
+    this.callback = callbackWithSentKeysTracking
+
+    // Create a filtered callback if where clause is provided
+    this.filteredCallback = options.whereExpression
+      ? createFilteredCallback(this.callback, options)
+      : this.callback
+  }
+
+  setOrderByIndex(index: BaseIndex<any>) {
+    this.orderByIndex = index
+  }
+
+  hasLoadedInitialState() {
+    return this.loadedInitialState
+  }
+
+  hasSentAtLeastOneSnapshot() {
+    return this.snapshotSent
+  }
+
+  emitEvents(changes: Array<ChangeMessage<any, any>>) {
+    const newChanges = this.filterAndFlipChanges(changes)
+    this.filteredCallback(newChanges)
+  }
+
+  /**
+   * Sends the snapshot to the callback.
+   * Returns a boolean indicating if it succeeded.
+   * It can only fail if there is no index to fulfill the request
+   * and the optimizedOnly option is set to true,
+   * or, the entire state was already loaded.
+   */
+  requestSnapshot(opts?: RequestSnapshotOptions): boolean {
+    if (this.loadedInitialState) {
+      // Subscription was deoptimized so we already sent the entire initial state
+      return false
+    }
+
+    const stateOpts: RequestSnapshotOptions = {
+      where: this.options.whereExpression,
+      optimizedOnly: opts?.optimizedOnly ?? false,
+    }
+
+    if (opts) {
+      if (`where` in opts) {
+        const snapshotWhereExp = opts.where
+        if (stateOpts.where) {
+          // Combine the two where expressions
+          const subWhereExp = stateOpts.where
+          const combinedWhereExp = and(subWhereExp, snapshotWhereExp)
+          stateOpts.where = combinedWhereExp
+        } else {
+          stateOpts.where = snapshotWhereExp
+        }
+      }
+    } else {
+      // No options provided so it's loading the entire initial state
+      this.loadedInitialState = true
+    }
+
+    const snapshot = this.collection.currentStateAsChanges(stateOpts)
+
+    if (snapshot === undefined) {
+      // Couldn't load from indexes
+      return false
+    }
+
+    // Only send changes that have not been sent yet
+    const filteredSnapshot = snapshot.filter(
+      (change) => !this.sentKeys.has(change.key)
+    )
+
+    this.snapshotSent = true
+    this.callback(filteredSnapshot)
+    return true
+  }
+
+  /**
+   * Sends a snapshot that is limited to the first `limit` rows that fulfill the `where` clause and are bigger than `minValue`.
+   * Requires a range index to be set with `setOrderByIndex` prior to calling this method.
+   * It uses that range index to load the items in the order of the index.
+   * Note: it does not send keys that have already been sent before.
+   */
+  requestLimitedSnapshot({ limit, minValue }: RequestLimitedSnapshotOptions) {
+    if (!limit) throw new Error(`limit is required`)
+
+    if (!this.orderByIndex) {
+      throw new Error(
+        `Ordered snapshot was requested but no index was found. You have to call setOrderByIndex before requesting an ordered snapshot.`
+      )
+    }
+
+    const index = this.orderByIndex
+    const where = this.options.whereExpression
+    const whereFilterFn = where
+      ? createFilterFunctionFromExpression(where)
+      : undefined
+
+    const filterFn = (key: string | number): boolean => {
+      if (this.sentKeys.has(key)) {
+        return false
+      }
+
+      const value = this.collection.get(key)
+      if (value === undefined) {
+        return false
+      }
+
+      return whereFilterFn?.(value) ?? true
+    }
+
+    let biggestObservedValue = minValue
+    const changes: Array<ChangeMessage<any, string | number>> = []
+    let keys: Array<string | number> = index.take(limit, minValue, filterFn)
+
+    const valuesNeeded = () => Math.max(limit - changes.length, 0)
+    const collectionExhausted = () => keys.length === 0
+
+    while (valuesNeeded() > 0 && !collectionExhausted()) {
+      for (const key of keys) {
+        const value = this.collection.get(key)!
+        changes.push({
+          type: `insert`,
+          key,
+          value,
+        })
+        biggestObservedValue = value
+      }
+
+      keys = index.take(valuesNeeded(), biggestObservedValue, filterFn)
+    }
+
+    this.callback(changes)
+  }
+
+  /**
+   * Filters and flips changes for keys that have not been sent yet.
+   * Deletes are filtered out for keys that have not been sent yet.
+   * Updates are flipped into inserts for keys that have not been sent yet.
+   */
+  private filterAndFlipChanges(changes: Array<ChangeMessage<any, any>>) {
+    if (this.loadedInitialState) {
+      // We loaded the entire initial state
+      // so no need to filter or flip changes
+      return changes
+    }
+
+    const newChanges = []
+    for (const change of changes) {
+      let newChange = change
+      if (!this.sentKeys.has(change.key)) {
+        if (change.type === `update`) {
+          newChange = { ...change, type: `insert`, previousValue: undefined }
+        } else if (change.type === `delete`) {
+          // filter out deletes for keys that have not been sent
+          continue
+        }
+        this.sentKeys.add(change.key)
+      }
+      newChanges.push(newChange)
+    }
+    return newChanges
+  }
+
+  private trackSentKeys(changes: Array<ChangeMessage<any, string | number>>) {
+    if (this.loadedInitialState) {
+      // No need to track sent keys if we loaded the entire state.
+      // Since we sent everything, all keys must have been observed.
+      return
+    }
+
+    for (const change of changes) {
+      this.sentKeys.add(change.key)
+    }
+  }
+
+  unsubscribe() {
+    this.options.onUnsubscribe?.()
+  }
+}

package/src/collection/sync.ts

@@ -0,0 +1,235 @@
+import {
+  CollectionIsInErrorStateError,
+  DuplicateKeySyncError,
+  NoPendingSyncTransactionCommitError,
+  NoPendingSyncTransactionWriteError,
+  SyncCleanupError,
+  SyncTransactionAlreadyCommittedError,
+  SyncTransactionAlreadyCommittedWriteError,
+} from "../errors"
+import type { StandardSchemaV1 } from "@standard-schema/spec"
+import type { ChangeMessage, CollectionConfig } from "../types"
+import type { CollectionImpl } from "./index.js"
+import type { CollectionStateManager } from "./state"
+import type { CollectionLifecycleManager } from "./lifecycle"
+
+export class CollectionSyncManager<
+  TOutput extends object = Record<string, unknown>,
+  TKey extends string | number = string | number,
+  TSchema extends StandardSchemaV1 = StandardSchemaV1,
+  TInput extends object = TOutput,
+> {
+  private collection!: CollectionImpl<TOutput, TKey, any, TSchema, TInput>
+  private state!: CollectionStateManager<TOutput, TKey, TSchema, TInput>
+  private lifecycle!: CollectionLifecycleManager<TOutput, TKey, TSchema, TInput>
+  private config!: CollectionConfig<TOutput, TKey, TSchema>
+  private id: string
+
+  public preloadPromise: Promise<void> | null = null
+  public syncCleanupFn: (() => void) | null = null
+
+  /**
+   * Creates a new CollectionSyncManager instance
+   */
+  constructor(config: CollectionConfig<TOutput, TKey, TSchema>, id: string) {
+    this.config = config
+    this.id = id
+  }
+
+  setDeps(deps: {
+    collection: CollectionImpl<TOutput, TKey, any, TSchema, TInput>
+    state: CollectionStateManager<TOutput, TKey, TSchema, TInput>
+    lifecycle: CollectionLifecycleManager<TOutput, TKey, TSchema, TInput>
+  }) {
+    this.collection = deps.collection
+    this.state = deps.state
+    this.lifecycle = deps.lifecycle
+  }
+
+  /**
+   * Start the sync process for this collection
+   * This is called when the collection is first accessed or preloaded
+   */
+  public startSync(): void {
+    const state = this.state
+    if (
+      this.lifecycle.status !== `idle` &&
+      this.lifecycle.status !== `cleaned-up`
+    ) {
+      return // Already started or in progress
+    }
+
+    this.lifecycle.setStatus(`loading`)
+
+    try {
+      const cleanupFn = this.config.sync.sync({
+        collection: this.collection,
+        begin: () => {
+          state.pendingSyncedTransactions.push({
+            committed: false,
+            operations: [],
+            deletedKeys: new Set(),
+          })
+        },
+        write: (messageWithoutKey: Omit<ChangeMessage<TOutput>, `key`>) => {
+          const pendingTransaction =
+            state.pendingSyncedTransactions[
+              state.pendingSyncedTransactions.length - 1
+            ]
+          if (!pendingTransaction) {
+            throw new NoPendingSyncTransactionWriteError()
+          }
+          if (pendingTransaction.committed) {
+            throw new SyncTransactionAlreadyCommittedWriteError()
+          }
+          const key = this.config.getKey(messageWithoutKey.value)
+
+          // Check if an item with this key already exists when inserting
+          if (messageWithoutKey.type === `insert`) {
+            const insertingIntoExistingSynced = state.syncedData.has(key)
+            const hasPendingDeleteForKey =
+              pendingTransaction.deletedKeys.has(key)
+            const isTruncateTransaction = pendingTransaction.truncate === true
+            // Allow insert after truncate in the same transaction even if it existed in syncedData
+            if (
+              insertingIntoExistingSynced &&
+              !hasPendingDeleteForKey &&
+              !isTruncateTransaction
+            ) {
+              throw new DuplicateKeySyncError(key, this.id)
+            }
+          }
+
+          const message: ChangeMessage<TOutput> = {
+            ...messageWithoutKey,
+            key,
+          }
+          pendingTransaction.operations.push(message)
+
+          if (messageWithoutKey.type === `delete`) {
+            pendingTransaction.deletedKeys.add(key)
+          }
+        },
+        commit: () => {
+          const pendingTransaction =
+            state.pendingSyncedTransactions[
+              state.pendingSyncedTransactions.length - 1
+            ]
+          if (!pendingTransaction) {
+            throw new NoPendingSyncTransactionCommitError()
+          }
+          if (pendingTransaction.committed) {
+            throw new SyncTransactionAlreadyCommittedError()
+          }
+
+          pendingTransaction.committed = true
+
+          // Update status to initialCommit when transitioning from loading
+          // This indicates we're in the process of committing the first transaction
+          if (this.lifecycle.status === `loading`) {
+            this.lifecycle.setStatus(`initialCommit`)
+          }
+
+          state.commitPendingTransactions()
+        },
+        markReady: () => {
+          this.lifecycle.markReady()
+        },
+        truncate: () => {
+          const pendingTransaction =
+            state.pendingSyncedTransactions[
+              state.pendingSyncedTransactions.length - 1
+            ]
+          if (!pendingTransaction) {
+            throw new NoPendingSyncTransactionWriteError()
+          }
+          if (pendingTransaction.committed) {
+            throw new SyncTransactionAlreadyCommittedWriteError()
+          }
+
+          // Clear all operations from the current transaction
+          pendingTransaction.operations = []
+          pendingTransaction.deletedKeys.clear()
+
+          // Mark the transaction as a truncate operation. During commit, this triggers:
+          // - Delete events for all previously synced keys (excluding optimistic-deleted keys)
+          // - Clearing of syncedData/syncedMetadata
+          // - Subsequent synced ops applied on the fresh base
+          // - Finally, optimistic mutations re-applied on top (single batch)
+          pendingTransaction.truncate = true
+        },
+      })
+
+      // Store cleanup function if provided
+      this.syncCleanupFn = typeof cleanupFn === `function` ? cleanupFn : null
+    } catch (error) {
+      this.lifecycle.setStatus(`error`)
+      throw error
+    }
+  }
+
+  /**
+   * Preload the collection data by starting sync if not already started
+   * Multiple concurrent calls will share the same promise
+   */
+  public preload(): Promise<void> {
+    if (this.preloadPromise) {
+      return this.preloadPromise
+    }
+
+    this.preloadPromise = new Promise<void>((resolve, reject) => {
+      if (this.lifecycle.status === `ready`) {
+        resolve()
+        return
+      }
+
+      if (this.lifecycle.status === `error`) {
+        reject(new CollectionIsInErrorStateError())
+        return
+      }
+
+      // Register callback BEFORE starting sync to avoid race condition
+      this.lifecycle.onFirstReady(() => {
+        resolve()
+      })
+
+      // Start sync if collection hasn't started yet or was cleaned up
+      if (
+        this.lifecycle.status === `idle` ||
+        this.lifecycle.status === `cleaned-up`
+      ) {
+        try {
+          this.startSync()
+        } catch (error) {
+          reject(error)
+          return
+        }
+      }
+    })
+
+    return this.preloadPromise
+  }
+
+  public cleanup(): void {
+    try {
+      if (this.syncCleanupFn) {
+        this.syncCleanupFn()
+        this.syncCleanupFn = null
+      }
+    } catch (error) {
+      // Re-throw in a microtask to surface the error after cleanup completes
+      queueMicrotask(() => {
+        if (error instanceof Error) {
+          // Preserve the original error and stack trace
+          const wrappedError = new SyncCleanupError(this.id, error)
+          wrappedError.cause = error
+          wrappedError.stack = error.stack
+          throw wrappedError
+        } else {
+          throw new SyncCleanupError(this.id, error as Error | string)
+        }
+      })
+    }
+    this.preloadPromise = null
+  }
+}

package/src/index.ts

@@ -1,5 +1,5 @@
 // Re-export all public APIs
-export * from "./collection"
+export * from "./collection/index.js"
 export * from "./SortedMap"
 export * from "./transactions"
 export * from "./types"
@@ -17,4 +17,4 @@ export * from "./indexes/lazy-index.js"
 export { type IndexOptions } from "./indexes/index-options.js"
 
 // Re-export some stuff explicitly to ensure the type & value is exported
-export type { Collection } from "./collection"
+export type { Collection } from "./collection/index.js"

package/src/indexes/auto-index.ts

@@ -1,6 +1,6 @@
 import { BTreeIndex } from "./btree-index"
 import type { BasicExpression } from "../query/ir"
-import type { CollectionImpl } from "../collection"
+import type { CollectionImpl } from "../collection/index.js"
 
 export interface AutoIndexConfig {
   autoIndex?: `off` | `eager`

package/src/indexes/base-index.ts

@@ -1,6 +1,6 @@
 import { compileSingleRowExpression } from "../query/compiler/evaluators.js"
 import { comparisonFunctions } from "../query/builder/functions.js"
-import type { BasicExpression, OrderByDirection } from "../query/ir.js"
+import type { BasicExpression } from "../query/ir.js"
 
 /**
  * Operations that indexes can support, imported from available comparison functions
@@ -58,8 +58,8 @@ export abstract class BaseIndex<
   abstract lookup(operation: IndexOperation, value: any): Set<TKey>
   abstract take(
     n: number,
-    direction?: OrderByDirection,
-    from?: TKey
+    from?: TKey,
+    filterFn?: (key: TKey) => boolean
   ): Array<TKey>
   abstract get keyCount(): number
 

package/src/indexes/btree-index.ts

@@ -236,7 +236,7 @@ export class BTreeIndex<
    * @param from - The item to start from (exclusive). Starts from the smallest item (inclusive) if not provided.
    * @returns The next n items after the provided key. Returns the first n items if no from item is provided.
    */
-  take(n: number, from?: any): Array<TKey> {
+  take(n: number, from?: any, filterFn?: (key: TKey) => boolean): Array<TKey> {
     const keysInResult: Set<TKey> = new Set()
     const result: Array<TKey> = []
     const nextKey = (k?: any) => this.orderedEntries.nextHigherKey(k)
@@ -248,7 +248,7 @@ export class BTreeIndex<
         const it = keys.values()
         let ks: TKey | undefined
         while (result.length < n && (ks = it.next().value)) {
-          if (!keysInResult.has(ks)) {
+          if (!keysInResult.has(ks) && (filterFn?.(ks) ?? true)) {
             result.push(ks)
             keysInResult.add(ks)
           }

package/src/query/builder/index.ts

@@ -1,4 +1,4 @@
-import { CollectionImpl } from "../../collection.js"
+import { CollectionImpl } from "../../collection/index.js"
 import {
   Aggregate as AggregateExpr,
   CollectionRef,

package/src/query/builder/types.ts

@@ -1,4 +1,4 @@
-import type { CollectionImpl } from "../../collection.js"
+import type { CollectionImpl } from "../../collection/index.js"
 import type {
   Aggregate,
   BasicExpression,

package/src/query/compiler/index.ts

@@ -13,6 +13,7 @@ import { processJoins } from "./joins.js"
 import { processGroupBy } from "./group-by.js"
 import { processOrderBy } from "./order-by.js"
 import { processSelect } from "./select.js"
+import type { CollectionSubscription } from "../../collection/subscription.js"
 import type { OrderByOptimizationInfo } from "./order-by.js"
 import type {
   BasicExpression,
@@ -21,7 +22,7 @@ import type {
   QueryRef,
 } from "../ir.js"
 import type { LazyCollectionCallbacks } from "./joins.js"
-import type { Collection } from "../../collection.js"
+import type { Collection } from "../../collection/index.js"
 import type {
   KeyedStream,
   NamespacedAndKeyedStream,
@@ -53,6 +54,7 @@ export function compileQuery(
   rawQuery: QueryIR,
   inputs: Record<string, KeyedStream>,
   collections: Record<string, Collection<any, any, any, any, any>>,
+  subscriptions: Record<string, CollectionSubscription>,
   callbacks: Record<string, LazyCollectionCallbacks>,
   lazyCollections: Set<string>,
   optimizableOrderByCollections: Record<string, OrderByOptimizationInfo>,
@@ -88,6 +90,7 @@ export function compileQuery(
     query.from,
     allInputs,
     collections,
+    subscriptions,
     callbacks,
     lazyCollections,
     optimizableOrderByCollections,
@@ -120,6 +123,7 @@ export function compileQuery(
       cache,
       queryMapping,
       collections,
+      subscriptions,
       callbacks,
       lazyCollections,
       optimizableOrderByCollections,
@@ -323,6 +327,7 @@ function processFrom(
   from: CollectionRef | QueryRef,
   allInputs: Record<string, KeyedStream>,
   collections: Record<string, Collection>,
+  subscriptions: Record<string, CollectionSubscription>,
   callbacks: Record<string, LazyCollectionCallbacks>,
   lazyCollections: Set<string>,
   optimizableOrderByCollections: Record<string, OrderByOptimizationInfo>,
@@ -346,6 +351,7 @@ function processFrom(
         originalQuery,
         allInputs,
         collections,
+        subscriptions,
         callbacks,
         lazyCollections,
         optimizableOrderByCollections,