Add onLoadMore callback and call it from the requestSnapshot methods

kevin-dp · samwillis · commit 7bb75b5edfa6 · 2025-09-27T12:19:41.000+01:00
diff --git a/packages/db/src/collection/subscription.ts b/packages/db/src/collection/subscription.ts
@@ -4,7 +4,9 @@ import {
   createFilterFunctionFromExpression,
   createFilteredCallback,
 } from "./change-events.js"
-import type { BasicExpression } from "../query/ir.js"
+import { ensureIndexForExpression } from "./indexes/auto-index.js"
+import { and } from "./query/index.js"
+import type { BasicExpression, OrderBy } from "../query/ir.js"
 import type { BaseIndex } from "../indexes/base-index.js"
 import type { ChangeMessage } from "../types.js"
 import type { CollectionImpl } from "./index.js"
@@ -15,8 +17,9 @@ type RequestSnapshotOptions = {
 }
 
 type RequestLimitedSnapshotOptions = {
-  minValue?: any
+  orderBy: OrderBy
   limit: number
+  minValue?: any
 }
 
 type CollectionSubscriptionOptions = {
@@ -117,6 +120,13 @@ export class CollectionSubscription {
       this.loadedInitialState = true
     }
 
+    // Request the sync layer to load more data
+    // don't await it, we will load the data into the collection when it comes in
+    this.collection.syncMore({
+      where: stateOpts.where,
+    })
+
+    // Also load data immediately from the collection
     const snapshot = this.collection.currentStateAsChanges(stateOpts)
 
     if (snapshot === undefined) {
@@ -140,7 +150,11 @@ export class CollectionSubscription {
    * 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) {
+  requestLimitedSnapshot({
+    orderBy,
+    limit,
+    minValue,
+  }: RequestLimitedSnapshotOptions) {
     if (!limit) throw new Error(`limit is required`)
 
     if (!this.orderByIndex) {
@@ -190,6 +204,14 @@ export class CollectionSubscription {
     }
 
     this.callback(changes)
+
+    // Request the sync layer to load more data
+    // don't await it, we will load the data into the collection when it comes in
+    this.collection.syncMore({
+      where,
+      limit,
+      orderBy,
+    })
   }
 
   /**
diff --git a/packages/db/src/collection/sync.ts b/packages/db/src/collection/sync.ts
@@ -8,7 +8,13 @@ import {
   SyncTransactionAlreadyCommittedWriteError,
 } from "../errors"
 import type { StandardSchemaV1 } from "@standard-schema/spec"
-import type { ChangeMessage, CollectionConfig } from "../types"
+import type {
+  ChangeMessage,
+  CollectionConfig,
+  CleanupFn,
+  OnLoadMoreOptions,
+  SyncConfigRes,
+} from "../types"
 import type { CollectionImpl } from "./index.js"
 import type { CollectionStateManager } from "./state"
 import type { CollectionLifecycleManager } from "./lifecycle"
@@ -27,6 +33,9 @@ export class CollectionSyncManager<
 
   public preloadPromise: Promise<void> | null = null
   public syncCleanupFn: (() => void) | null = null
+  public syncOnLoadMoreFn:
+    | ((options: OnLoadMoreOptions) => void | Promise<void>)
+    | null = null
 
   /**
    * Creates a new CollectionSyncManager instance
@@ -62,106 +71,111 @@ export class CollectionSyncManager<
     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 syncRes = normalizeSyncFnResult(
+        this.config.sync.sync({
+          collection: this,
+          begin: () => {
+            this.pendingSyncedTransactions.push({
+              committed: false,
+              operations: [],
+              deletedKeys: new Set(),
+            })
+          },
+          write: (messageWithoutKey: Omit<ChangeMessage<TOutput>, `key`>) => {
+            const pendingTransaction =
+              this.pendingSyncedTransactions[
+                this.pendingSyncedTransactions.length - 1
+              ]
+            if (!pendingTransaction) {
+              throw new NoPendingSyncTransactionWriteError()
             }
-          }
-
-          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
-        },
-      })
+            if (pendingTransaction.committed) {
+              throw new SyncTransactionAlreadyCommittedWriteError()
+            }
+            const key = this.getKeyFromItem(messageWithoutKey.value)
+
+            // Check if an item with this key already exists when inserting
+            if (messageWithoutKey.type === `insert`) {
+              const insertingIntoExistingSynced = this.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 =
+              this.pendingSyncedTransactions[
+                this.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._status === `loading`) {
+              this.setStatus(`initialCommit`)
+            }
+
+            this.commitPendingTransactions()
+          },
+          markReady: () => {
+            this.markReady()
+          },
+          truncate: () => {
+            const pendingTransaction =
+              this.pendingSyncedTransactions[
+                this.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
+      this.syncCleanupFn = syncRes?.cleanup ?? null
+
+      // Store onLoadMore function if provided
+      this.syncOnLoadMoreFn = syncRes?.onLoadMore ?? null
     } catch (error) {
       this.lifecycle.setStatus(`error`)
       throw error
@@ -210,6 +224,18 @@ export class CollectionSyncManager<
     return this.preloadPromise
   }
 
+  /**
+   * Requests the sync layer to load more data.
+   * @param options Options to control what data is being loaded
+   * @returns If data loading is asynchronous, this method returns a promise that resolves when the data is loaded.
+   *          If data loading is synchronous, the data is loaded when the method returns.
+   */
+  public syncMore(options: OnLoadMoreOptions): void | Promise<void> {
+    if (this.syncOnLoadMoreFn) {
+      return this.syncOnLoadMoreFn(options)
+    }
+  }
+
   public cleanup(): void {
     try {
       if (this.syncCleanupFn) {
@@ -233,3 +259,15 @@ export class CollectionSyncManager<
     this.preloadPromise = null
   }
 }
+
+function normalizeSyncFnResult(result: void | CleanupFn | SyncConfigRes) {
+  if (typeof result === `function`) {
+    return { cleanup: result }
+  }
+
+  if (typeof result === `object`) {
+    return result
+  }
+
+  return undefined
+}
diff --git a/packages/db/src/query/compiler/order-by.ts b/packages/db/src/query/compiler/order-by.ts
@@ -7,13 +7,14 @@ import { compileExpression } from "./evaluators.js"
 import { replaceAggregatesByRefs } from "./group-by.js"
 import { followRef } from "./index.js"
 import type { CompiledSingleRowExpression } from "./evaluators.js"
-import type { OrderByClause, QueryIR, Select } from "../ir.js"
+import type { OrderBy, OrderByClause, QueryIR, Select } from "../ir.js"
 import type { NamespacedAndKeyedStream, NamespacedRow } from "../../types.js"
 import type { IStreamBuilder, KeyValue } from "@tanstack/db-ivm"
 import type { BaseIndex } from "../../indexes/base-index.js"
 import type { Collection } from "../../collection/index.js"
 
 export type OrderByOptimizationInfo = {
+  orderBy: OrderBy
   offset: number
   limit: number
   comparator: (
@@ -164,6 +165,7 @@ export function processOrderBy(
           comparator,
           valueExtractorForRawRow,
           index,
+          orderBy: orderByClause,
         }
 
         optimizableOrderByCollections[followRefCollection.id] =
diff --git a/packages/db/src/query/live/collection-subscriber.ts b/packages/db/src/query/live/collection-subscriber.ts
@@ -128,7 +128,7 @@ export class CollectionSubscriber<
   private subscribeToOrderedChanges(
     whereExpression: BasicExpression<boolean> | undefined
   ) {
-    const { offset, limit, comparator, dataNeeded, index } =
+    const { orderBy, offset, limit, comparator, dataNeeded, index } =
       this.collectionConfigBuilder.optimizableOrderByCollections[
         this.collectionId
       ]!
@@ -168,6 +168,7 @@ export class CollectionSubscriber<
     // i.e. the K items from the collection that fall into the requested range: [offset, offset + limit[
     subscription.requestLimitedSnapshot({
       limit: offset + limit,
+      orderBy,
     })
 
     return subscription
@@ -225,7 +226,7 @@ export class CollectionSubscriber<
   // Loads the next `n` items from the collection
   // starting from the biggest item it has sent
   private loadNextItems(n: number, subscription: CollectionSubscription) {
-    const { valueExtractorForRawRow } =
+    const { orderBy, valueExtractorForRawRow } =
       this.collectionConfigBuilder.optimizableOrderByCollections[
         this.collectionId
       ]!
@@ -235,6 +236,7 @@ export class CollectionSubscriber<
       : biggestSentRow
     // Take the `n` items after the biggest sent value
     subscription.requestLimitedSnapshot({
+      orderBy,
       limit: n,
       minValue: biggestSentValue,
     })
diff --git a/packages/db/src/types.ts b/packages/db/src/types.ts
