Add onLoadMore callback and call it from the requestSnapshot methods

Author: kevin-dp

packages/db/src/collection-subscription.ts

@@ -4,7 +4,7 @@ import {
 } from "./change-events.js"
 import { ensureIndexForExpression } from "./indexes/auto-index.js"
 import { and } from "./query/index.js"
-import type { BasicExpression } from "./query/ir.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 { Collection } from "./collection.js"
@@ -15,8 +15,9 @@ type RequestSnapshotOptions = {
 }
 
 type RequestLimitedSnapshotOptions = {
-  minValue?: any
+  orderBy: OrderBy
   limit: number
+  minValue?: any
 }
 
 type CollectionSubscriptionOptions = {
@@ -117,6 +118,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 +148,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 +202,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,
+    })
   }
 
   /**

packages/db/src/collection.ts

@@ -37,29 +37,32 @@ import {
   UpdateKeyNotFoundError,
 } from "./errors"
 import { CollectionEvents } from "./collection-events.js"
+import { currentStateAsChanges } from "./change-events"
+import { CollectionSubscription } from "./collection-subscription.js"
 import type {
   AllCollectionEvents,
   CollectionEventHandler,
 } from "./collection-events.js"
-import { currentStateAsChanges } from "./change-events"
-import { CollectionSubscription } from "./collection-subscription.js"
 import type { Transaction } from "./transactions"
 import type { StandardSchemaV1 } from "@standard-schema/spec"
 import type { SingleRowRefProxy } from "./query/builder/ref-proxy"
 import type {
   ChangeMessage,
+  CleanupFn,
   CollectionConfig,
   CollectionStatus,
   CurrentStateAsChangesOptions,
   Fn,
   InferSchemaInput,
   InferSchemaOutput,
   InsertConfig,
+  OnLoadMoreOptions,
   OperationConfig,
   OptimisticChangeMessage,
   PendingMutation,
   StandardSchema,
   SubscribeChangesOptions,
+  SyncConfigRes,
   Transaction as TransactionType,
   TransactionWithMutations,
   UtilsRecord,
@@ -266,6 +269,9 @@ export class CollectionImpl<
   private gcTimeoutId: ReturnType<typeof setTimeout> | null = null
   private preloadPromise: Promise<void> | null = null
   private syncCleanupFn: (() => void) | null = null
+  private syncOnLoadMoreFn:
+    | ((options: OnLoadMoreOptions) => void | Promise<void>)
+    | null = null
 
   // Event system
   private events: CollectionEvents
@@ -488,106 +494,111 @@ export class CollectionImpl<
     this.setStatus(`loading`)
 
     try {
-      const cleanupFn = 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()
-          }
-          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 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()
+            }
+            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)
+            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()
-          }
+            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
+            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()
-          }
+            // 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`)
+            }
 
-          // Clear all operations from the current transaction
-          pendingTransaction.operations = []
-          pendingTransaction.deletedKeys.clear()
+            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()
+            }
 
-          // 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
-        },
-      })
+            // 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.setStatus(`error`)
       throw error
@@ -633,6 +644,18 @@ export class CollectionImpl<
     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)
+    }
+  }
+
   /**
    * Clean up the collection by stopping sync and clearing data
    * This can be called manually or automatically by garbage collection
@@ -2478,3 +2501,15 @@ export class CollectionImpl<
     return this.events.waitFor(event, timeout)
   }
 }
+
+function normalizeSyncFnResult(result: void | CleanupFn | SyncConfigRes) {
+  if (typeof result === `function`) {
+    return { cleanup: result }
+  }
+
+  if (typeof result === `object`) {
+    return result
+  }
+
+  return undefined
+}

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.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] =