Handle pushed down predicates in Electric collection

Co-authored-by: Kevin De Porre <[email protected]>
Co-authored-by: Sam Willis <[email protected]>

Author: samwillis

.changeset/tender-carpets-cheat.md

@@ -0,0 +1,5 @@
+---
+"@tanstack/electric-db-collection": patch
+---
+
+Handle predicates that are pushed down.

packages/electric-db-collection/src/electric.ts

@@ -12,12 +12,15 @@ import {
   TimeoutWaitingForMatchError,
   TimeoutWaitingForTxIdError,
 } from "./errors"
+import { compileSQL } from "./sql-compiler"
 import type {
   BaseCollectionConfig,
   CollectionConfig,
   DeleteMutationFnParams,
   InsertMutationFnParams,
+  LoadSubsetOptions,
   SyncConfig,
+  SyncMode,
   UpdateMutationFnParams,
   UtilsRecord,
 } from "@tanstack/db"
@@ -72,6 +75,24 @@ type InferSchemaOutput<T> = T extends StandardSchemaV1
     : Record<string, unknown>
   : Record<string, unknown>
 
+/**
+ * The mode of sync to use for the collection.
+ * @default `eager`
+ * @description
+ * - `eager`:
+ *   - syncs all data immediately on preload
+ *   - collection will be marked as ready once the sync is complete
+ *   - there is no incremental sync
+ * - `on-demand`:
+ *   - syncs data in incremental snapshots when the collection is queried
+ *   - collection will be marked as ready immediately after the first snapshot is synced
+ * - `progressive`:
+ *   - syncs all data for the collection in the background
+ *   - uses incremental snapshots during the initial sync to provide a fast path to the data required for queries
+ *   - collection will be marked as ready once the full sync is complete
+ */
+export type ElectricSyncMode = SyncMode | `progressive`
+
 /**
  * Configuration interface for Electric collection options
  * @template T - The type of items in the collection
@@ -82,12 +103,13 @@ export interface ElectricCollectionConfig<
   TSchema extends StandardSchemaV1 = never,
 > extends Omit<
     BaseCollectionConfig<T, string | number, TSchema, UtilsRecord, any>,
-    `onInsert` | `onUpdate` | `onDelete`
+    `onInsert` | `onUpdate` | `onDelete` | `syncMode`
   > {
   /**
    * Configuration options for the ElectricSQL ShapeStream
    */
   shapeOptions: ShapeStreamOptions<GetExtensions<T>>
+  syncMode?: ElectricSyncMode
 
   /**
    * Optional asynchronous handler function called before an insert operation
@@ -281,6 +303,8 @@ export function electricCollectionOptions(
 } {
   const seenTxids = new Store<Set<Txid>>(new Set([]))
   const seenSnapshots = new Store<Array<PostgresSnapshot>>([])
+  const syncMode = config.syncMode ?? `eager`
+  const finalSyncMode = syncMode === `progressive` ? `on-demand` : syncMode
   const pendingMatches = new Store<
     Map<
       string,
@@ -331,6 +355,7 @@ export function electricCollectionOptions(
   const sync = createElectricSync<any>(config.shapeOptions, {
     seenTxids,
     seenSnapshots,
+    syncMode,
     pendingMatches,
     currentBatchMessages,
     removePendingMatches,
@@ -550,6 +575,7 @@ export function electricCollectionOptions(
 
   return {
     ...restConfig,
+    syncMode: finalSyncMode,
     sync,
     onInsert: wrappedOnInsert,
     onUpdate: wrappedOnUpdate,
@@ -567,6 +593,7 @@ export function electricCollectionOptions(
 function createElectricSync<T extends Row<unknown>>(
   shapeOptions: ShapeStreamOptions<GetExtensions<T>>,
   options: {
+    syncMode: ElectricSyncMode
     seenTxids: Store<Set<Txid>>
     seenSnapshots: Store<Array<PostgresSnapshot>>
     pendingMatches: Store<
@@ -590,6 +617,7 @@ function createElectricSync<T extends Row<unknown>>(
   const {
     seenTxids,
     seenSnapshots,
+    syncMode,
     pendingMatches,
     currentBatchMessages,
     removePendingMatches,
@@ -653,6 +681,15 @@ function createElectricSync<T extends Row<unknown>>(
 
       const stream = new ShapeStream({
         ...shapeOptions,
+        // In on-demand mode, we only want to sync changes, so we set the log to `changes_only`
+        log: syncMode === `on-demand` ? `changes_only` : undefined,
+        // In on-demand mode, we only need the changes from the point of time the collection was created
+        // so we default to `now` when there is no saved offset.
+        offset: shapeOptions.offset
+          ? shapeOptions.offset
+          : syncMode === `on-demand`
+            ? `now`
+            : undefined,
         signal: abortController.signal,
         onError: (errorParams) => {
           // Just immediately mark ready if there's an error to avoid blocking
@@ -679,9 +716,11 @@ function createElectricSync<T extends Row<unknown>>(
       let transactionStarted = false
       const newTxids = new Set<Txid>()
       const newSnapshots: Array<PostgresSnapshot> = []
+      let hasReceivedUpToDate = false // Track if we've completed initial sync in progressive mode
 
       unsubscribeStream = stream.subscribe((messages: Array<Message<T>>) => {
         let hasUpToDate = false
+        let hasSnapshotEnd = false
 
         for (const message of messages) {
           // Add message to current batch buffer (for race condition handling)
@@ -746,6 +785,7 @@ function createElectricSync<T extends Row<unknown>>(
             })
           } else if (isSnapshotEndMessage(message)) {
             newSnapshots.push(parseSnapshotMessage(message))
+            hasSnapshotEnd = true
           } else if (isUpToDateMessage(message)) {
             hasUpToDate = true
           } else if (isMustRefetchMessage(message)) {
@@ -761,12 +801,14 @@ function createElectricSync<T extends Row<unknown>>(
 
             truncate()
 
-            // Reset hasUpToDate so we continue accumulating changes until next up-to-date
+            // Reset flags so we continue accumulating changes until next up-to-date
             hasUpToDate = false
+            hasSnapshotEnd = false
+            hasReceivedUpToDate = false // Reset for progressive mode - we're starting a new sync
           }
         }
 
-        if (hasUpToDate) {
+        if (hasUpToDate || hasSnapshotEnd) {
           // Clear the current batch buffer since we're now up-to-date
           currentBatchMessages.setState(() => [])
 
@@ -776,8 +818,15 @@ function createElectricSync<T extends Row<unknown>>(
             transactionStarted = false
           }
 
-          // Mark the collection as ready now that sync is up to date
-          markReady()
+          if (hasUpToDate || (hasSnapshotEnd && syncMode === `on-demand`)) {
+            // Mark the collection as ready now that sync is up to date
+            markReady()
+          }
+
+          // Track that we've received the first up-to-date for progressive mode
+          if (hasUpToDate) {
+            hasReceivedUpToDate = true
+          }
 
           // Always commit txids when we receive up-to-date, regardless of transaction state
           seenTxids.setState((currentTxids) => {
@@ -811,12 +860,29 @@ function createElectricSync<T extends Row<unknown>>(
         }
       })
 
-      // Return the unsubscribe function
-      return () => {
-        // Unsubscribe from the stream
-        unsubscribeStream()
-        // Abort the abort controller to stop the stream
-        abortController.abort()
+      // Only set onLoadSubset if the sync mode is not eager, this indicates to the sync
+      // layer that it can load more data on demand via the requestSnapshot method when,
+      // the syncMode = `on-demand` or `progressive`
+      const loadSubset =
+        syncMode === `eager`
+          ? undefined
+          : async (opts: LoadSubsetOptions) => {
+              // In progressive mode, stop requesting snapshots once full sync is complete
+              if (syncMode === `progressive` && hasReceivedUpToDate) {
+                return
+              }
+              const snapshotParams = compileSQL<T>(opts)
+              await stream.requestSnapshot(snapshotParams)
+            }
+
+      return {
+        loadSubset,
+        cleanup: () => {
+          // Unsubscribe from the stream
+          unsubscribeStream()
+          // Abort the abort controller to stop the stream
+          abortController.abort()
+        },
       }
     },
     // Expose the getSyncMetadata function

packages/electric-db-collection/src/pg-serializer.ts

@@ -0,0 +1,27 @@
+export function serialize(value: unknown): string {
+  if (typeof value === `string`) {
+    return `'${value}'`
+  }
+
+  if (typeof value === `number`) {
+    return value.toString()
+  }
+
+  if (value === null || value === undefined) {
+    return `NULL`
+  }
+
+  if (typeof value === `boolean`) {
+    return value ? `true` : `false`
+  }
+
+  if (value instanceof Date) {
+    return `'${value.toISOString()}'`
+  }
+
+  if (Array.isArray(value)) {
+    return `ARRAY[${value.map(serialize).join(`,`)}]`
+  }
+
+  throw new Error(`Cannot serialize value: ${JSON.stringify(value)}`)
+}