electric-db-collection: support for snapshots in awaitTxid (#648)

* first pass at adding support for snapshots to awaitTxid

* add tests

* fix cleanup

* changeset

* address review

Author: samwillis

.changeset/eager-wings-jog.md

@@ -0,0 +1,5 @@
+---
+"@tanstack/electric-db-collection": patch
+---
+
+The awaitTxId utility now resolves transaction IDs based on snapshot-end message metadata (xmin, xmax, xip_list) in addition to explicit txid arrays, enabling matching on the initial snapshot at the start of a new shape.

packages/electric-db-collection/package.json

@@ -3,8 +3,8 @@
   "description": "ElectricSQL collection for TanStack DB",
   "version": "0.1.28",
   "dependencies": {
-    "@standard-schema/spec": "^1.0.0",
     "@electric-sql/client": "^1.0.14",
+    "@standard-schema/spec": "^1.0.0",
     "@tanstack/db": "workspace:*",
     "@tanstack/store": "^0.7.7",
     "debug": "^4.4.3"

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

@@ -2,6 +2,7 @@ import {
   ShapeStream,
   isChangeMessage,
   isControlMessage,
+  isVisibleInSnapshot,
 } from "@electric-sql/client"
 import { Store } from "@tanstack/store"
 import DebugModule from "debug"
@@ -27,6 +28,7 @@ import type {
   ControlMessage,
   GetExtensions,
   Message,
+  PostgresSnapshot,
   Row,
   ShapeStreamOptions,
 } from "@electric-sql/client"
@@ -38,6 +40,23 @@ const debug = DebugModule.debug(`ts/db:electric`)
  */
 export type Txid = number
 
+/**
+ * Type representing the result of an insert, update, or delete handler
+ */
+type MaybeTxId =
+  | {
+      txid?: Txid | Array<Txid>
+    }
+  | undefined
+  | null
+
+/**
+ * Type representing a snapshot end message
+ */
+type SnapshotEndMessage = ControlMessage & {
+  headers: { control: `snapshot-end` }
+}
+
 // The `InferSchemaOutput` and `ResolveType` are copied from the `@tanstack/db` package
 // but we modified `InferSchemaOutput` slightly to restrict the schema output to `Row<unknown>`
 // This is needed in order for `GetExtensions` to be able to infer the parser extensions type from the schema
@@ -80,6 +99,20 @@ function isMustRefetchMessage<T extends Row<unknown>>(
   return isControlMessage(message) && message.headers.control === `must-refetch`
 }
 
+function isSnapshotEndMessage<T extends Row<unknown>>(
+  message: Message<T>
+): message is SnapshotEndMessage {
+  return isControlMessage(message) && message.headers.control === `snapshot-end`
+}
+
+function parseSnapshotMessage(message: SnapshotEndMessage): PostgresSnapshot {
+  return {
+    xmin: message.headers.xmin,
+    xmax: message.headers.xmax,
+    xip_list: message.headers.xip_list,
+  }
+}
+
 // Check if a message contains txids in its headers
 function hasTxids<T extends Row<unknown>>(
   message: Message<T>
@@ -139,8 +172,10 @@ export function electricCollectionOptions(
   schema?: any
 } {
   const seenTxids = new Store<Set<Txid>>(new Set([]))
+  const seenSnapshots = new Store<Array<PostgresSnapshot>>([])
   const sync = createElectricSync<any>(config.shapeOptions, {
     seenTxids,
+    seenSnapshots,
   })
 
   /**
@@ -158,20 +193,46 @@ export function electricCollectionOptions(
       throw new ExpectedNumberInAwaitTxIdError(typeof txId)
     }
 
+    // First check if the txid is in the seenTxids store
     const hasTxid = seenTxids.state.has(txId)
     if (hasTxid) return true
 
+    // Then check if the txid is in any of the seen snapshots
+    const hasSnapshot = seenSnapshots.state.some((snapshot) =>
+      isVisibleInSnapshot(txId, snapshot)
+    )
+    if (hasSnapshot) return true
+
     return new Promise((resolve, reject) => {
       const timeoutId = setTimeout(() => {
-        unsubscribe()
+        unsubscribeSeenTxids()
+        unsubscribeSeenSnapshots()
         reject(new TimeoutWaitingForTxIdError(txId))
       }, timeout)
 
-      const unsubscribe = seenTxids.subscribe(() => {
+      const unsubscribeSeenTxids = seenTxids.subscribe(() => {
         if (seenTxids.state.has(txId)) {
           debug(`awaitTxId found match for txid %o`, txId)
           clearTimeout(timeoutId)
-          unsubscribe()
+          unsubscribeSeenTxids()
+          unsubscribeSeenSnapshots()
+          resolve(true)
+        }
+      })
+
+      const unsubscribeSeenSnapshots = seenSnapshots.subscribe(() => {
+        const visibleSnapshot = seenSnapshots.state.find((snapshot) =>
+          isVisibleInSnapshot(txId, snapshot)
+        )
+        if (visibleSnapshot) {
+          debug(
+            `awaitTxId found match for txid %o in snapshot %o`,
+            txId,
+            visibleSnapshot
+          )
+          clearTimeout(timeoutId)
+          unsubscribeSeenSnapshots()
+          unsubscribeSeenTxids()
           resolve(true)
         }
       })
@@ -183,8 +244,9 @@ export function electricCollectionOptions(
     ? async (params: InsertMutationFnParams<any>) => {
         // Runtime check (that doesn't follow type)
 
-        const handlerResult = (await config.onInsert!(params)) ?? {}
-        const txid = (handlerResult as { txid?: Txid | Array<Txid> }).txid
+        const handlerResult =
+          ((await config.onInsert!(params)) as MaybeTxId) ?? {}
+        const txid = handlerResult.txid
 
         if (!txid) {
           throw new ElectricInsertHandlerMustReturnTxIdError()
@@ -205,8 +267,9 @@ export function electricCollectionOptions(
     ? async (params: UpdateMutationFnParams<any>) => {
         // Runtime check (that doesn't follow type)
 
-        const handlerResult = (await config.onUpdate!(params)) ?? {}
-        const txid = (handlerResult as { txid?: Txid | Array<Txid> }).txid
+        const handlerResult =
+          ((await config.onUpdate!(params)) as MaybeTxId) ?? {}
+        const txid = handlerResult.txid
 
         if (!txid) {
           throw new ElectricUpdateHandlerMustReturnTxIdError()
@@ -269,9 +332,11 @@ function createElectricSync<T extends Row<unknown>>(
   shapeOptions: ShapeStreamOptions<GetExtensions<T>>,
   options: {
     seenTxids: Store<Set<Txid>>
+    seenSnapshots: Store<Array<PostgresSnapshot>>
   }
 ): SyncConfig<T> {
   const { seenTxids } = options
+  const { seenSnapshots } = options
 
   // Store for the relation schema information
   const relationSchema = new Store<string | undefined>(undefined)
@@ -342,6 +407,7 @@ function createElectricSync<T extends Row<unknown>>(
       })
       let transactionStarted = false
       const newTxids = new Set<Txid>()
+      const newSnapshots: Array<PostgresSnapshot> = []
 
       unsubscribeStream = stream.subscribe((messages: Array<Message<T>>) => {
         let hasUpToDate = false
@@ -373,6 +439,8 @@ function createElectricSync<T extends Row<unknown>>(
                 ...message.headers,
               },
             })
+          } else if (isSnapshotEndMessage(message)) {
+            newSnapshots.push(parseSnapshotMessage(message))
           } else if (isUpToDateMessage(message)) {
             hasUpToDate = true
           } else if (isMustRefetchMessage(message)) {
@@ -413,6 +481,16 @@ function createElectricSync<T extends Row<unknown>>(
             newTxids.clear()
             return clonedSeen
           })
+
+          // Always commit snapshots when we receive up-to-date, regardless of transaction state
+          seenSnapshots.setState((currentSnapshots) => {
+            const seen = [...currentSnapshots, ...newSnapshots]
+            newSnapshots.forEach((snapshot) =>
+              debug(`new snapshot synced from pg %o`, snapshot)
+            )
+            newSnapshots.length = 0
+            return seen
+          })
         }
       })