Handle pushed down predicates in Electric collection

Co-authored-by: Kevin De Porre <kevin@electric-sql.com>
Co-authored-by: Sam Willis <sam.willis@gmail.com>

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,9 @@ export function electricCollectionOptions(
 } {
   const seenTxids = new Store<Set<Txid>>(new Set([]))
   const seenSnapshots = new Store<Array<PostgresSnapshot>>([])
+  const internalSyncMode = config.syncMode ?? `eager`
+  const finalSyncMode =
+    internalSyncMode === `progressive` ? `on-demand` : internalSyncMode
   const pendingMatches = new Store<
     Map<
       string,
@@ -331,6 +356,7 @@ export function electricCollectionOptions(
   const sync = createElectricSync<any>(config.shapeOptions, {
     seenTxids,
     seenSnapshots,
+    syncMode: internalSyncMode,
     pendingMatches,
     currentBatchMessages,
     removePendingMatches,
@@ -550,6 +576,7 @@ export function electricCollectionOptions(
 
   return {
     ...restConfig,
+    syncMode: finalSyncMode,
     sync,
     onInsert: wrappedOnInsert,
     onUpdate: wrappedOnUpdate,
@@ -567,6 +594,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 +618,7 @@ function createElectricSync<T extends Row<unknown>>(
   const {
     seenTxids,
     seenSnapshots,
+    syncMode,
     pendingMatches,
     currentBatchMessages,
     removePendingMatches,
@@ -653,6 +682,12 @@ 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 ?? (syncMode === `on-demand` ? `now` : undefined),
         signal: abortController.signal,
         onError: (errorParams) => {
           // Just immediately mark ready if there's an error to avoid blocking
@@ -679,9 +714,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 +783,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 +799,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 +816,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 +858,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)}`)
+}

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

@@ -0,0 +1,163 @@
+import { serialize } from "./pg-serializer"
+import type { SubsetParams } from "@electric-sql/client"
+import type { IR, LoadSubsetOptions } from "@tanstack/db"
+
+export type CompiledSqlRecord = Omit<SubsetParams, `params`> & {
+  params?: Array<unknown>
+}
+
+export function compileSQL<T>(options: LoadSubsetOptions): SubsetParams {
+  const { where, orderBy, limit } = options
+
+  const params: Array<T> = []
+  const compiledSQL: CompiledSqlRecord = { params }
+
+  if (where) {
+    // TODO: this only works when the where expression's PropRefs directly reference a column of the collection
+    //       doesn't work if it goes through aliases because then we need to know the entire query to be able to follow the reference until the base collection (cf. followRef function)
+    compiledSQL.where = compileBasicExpression(where, params)
+  }
+
+  if (orderBy) {
+    compiledSQL.orderBy = compileOrderBy(orderBy, params)
+  }
+
+  if (limit) {
+    compiledSQL.limit = limit
+  }
+
+  // Serialize the values in the params array into PG formatted strings
+  // and transform the array into a Record<string, string>
+  const paramsRecord = params.reduce(
+    (acc, param, index) => {
+      acc[`${index + 1}`] = serialize(param)
+      return acc
+    },
+    {} as Record<string, string>
+  )
+
+  return {
+    ...compiledSQL,
+    params: paramsRecord,
+  }
+}
+
+/**
+ * Compiles the expression to a SQL string and mutates the params array with the values.
+ * @param exp - The expression to compile
+ * @param params - The params array
+ * @returns The compiled SQL string
+ */
+function compileBasicExpression(
+  exp: IR.BasicExpression<unknown>,
+  params: Array<unknown>
+): string {
+  switch (exp.type) {
+    case `val`:
+      params.push(exp.value)
+      return `$${params.length}`
+    case `ref`:
+      // TODO: doesn't yet support JSON(B) values which could be accessed with nested props
+      if (exp.path.length !== 1) {
+        throw new Error(
+          `Compiler can't handle nested properties: ${exp.path.join(`.`)}`
+        )
+      }
+      return exp.path[0]!
+    case `func`:
+      return compileFunction(exp, params)
+    default:
+      throw new Error(`Unknown expression type`)
+  }
+}
+
+function compileOrderBy(orderBy: IR.OrderBy, params: Array<unknown>): string {
+  const compiledOrderByClauses = orderBy.map((clause: IR.OrderByClause) =>
+    compileOrderByClause(clause, params)
+  )
+  return compiledOrderByClauses.join(`,`)
+}
+
+function compileOrderByClause(
+  clause: IR.OrderByClause,
+  params: Array<unknown>
+): string {
+  // TODO: what to do with stringSort and locale?
+  //       Correctly supporting them is tricky as it depends on Postgres' collation
+  const { expression, compareOptions } = clause
+  let sql = compileBasicExpression(expression, params)
+
+  if (compareOptions.direction === `desc`) {
+    sql = `${sql} DESC`
+  }
+
+  if (compareOptions.nulls === `first`) {
+    sql = `${sql} NULLS FIRST`
+  }
+
+  if (compareOptions.nulls === `last`) {
+    sql = `${sql} NULLS LAST`
+  }
+
+  return sql
+}
+
+function compileFunction(
+  exp: IR.Func<unknown>,
+  params: Array<unknown> = []
+): string {
+  const { name, args } = exp
+
+  const opName = getOpName(name)
+
+  const compiledArgs = args.map((arg: IR.BasicExpression) =>
+    compileBasicExpression(arg, params)
+  )
+
+  if (isBinaryOp(name)) {
+    if (compiledArgs.length !== 2) {
+      throw new Error(`Binary operator ${name} expects 2 arguments`)
+    }
+    const [lhs, rhs] = compiledArgs
+    return `${lhs} ${opName} ${rhs}`
+  }
+
+  return `${opName}(${compiledArgs.join(`,`)})`
+}
+
+function isBinaryOp(name: string): boolean {
+  const binaryOps = [`eq`, `gt`, `gte`, `lt`, `lte`, `and`, `or`]
+  return binaryOps.includes(name)
+}
+
+function getOpName(name: string): string {
+  const opNames = {
+    eq: `=`,
+    gt: `>`,
+    gte: `>=`,
+    lt: `<`,
+    lte: `<=`,
+    add: `+`,
+    and: `AND`,
+    or: `OR`,
+    not: `NOT`,
+    isUndefined: `IS NULL`,
+    isNull: `IS NULL`,
+    in: `IN`,
+    like: `LIKE`,
+    ilike: `ILIKE`,
+    upper: `UPPER`,
+    lower: `LOWER`,
+    length: `LENGTH`,
+    concat: `CONCAT`,
+    coalesce: `COALESCE`,
+  }
+
+  const opName = opNames[name as keyof typeof opNames]
+
+  if (!opName) {
+    throw new Error(`Unknown operator/function: ${name}`)
+  }
+
+  return opName
+}