@tanstack/electric-db-collection 0.1.29 → 0.1.31

package/src/electric.ts

@@ -7,17 +7,15 @@ import {
 import { Store } from "@tanstack/store"
 import DebugModule from "debug"
 import {
-  ElectricDeleteHandlerMustReturnTxIdError,
-  ElectricInsertHandlerMustReturnTxIdError,
-  ElectricUpdateHandlerMustReturnTxIdError,
   ExpectedNumberInAwaitTxIdError,
+  StreamAbortedError,
+  TimeoutWaitingForMatchError,
   TimeoutWaitingForTxIdError,
 } from "./errors"
 import type {
   BaseCollectionConfig,
   CollectionConfig,
   DeleteMutationFnParams,
-  Fn,
   InsertMutationFnParams,
   SyncConfig,
   UpdateMutationFnParams,
@@ -33,6 +31,9 @@ import type {
   ShapeStreamOptions,
 } from "@electric-sql/client"
 
+// Re-export for user convenience in custom match functions
+export { isChangeMessage, isControlMessage } from "@electric-sql/client"
+
 const debug = DebugModule.debug(`ts/db:electric`)
 
 /**
@@ -41,14 +42,20 @@ const debug = DebugModule.debug(`ts/db:electric`)
 export type Txid = number
 
 /**
- * Type representing the result of an insert, update, or delete handler
+ * Custom match function type - receives stream messages and returns boolean
+ * indicating if the mutation has been synchronized
  */
-type MaybeTxId =
-  | {
-      txid?: Txid | Array<Txid>
-    }
-  | undefined
-  | null
+export type MatchFunction<T extends Row<unknown>> = (
+  message: Message<T>
+) => boolean
+
+/**
+ * Matching strategies for Electric synchronization
+ * Handlers can return:
+ * - Txid strategy: { txid: number | number[] } (recommended)
+ * - Void (no return value) - mutation completes without waiting
+ */
+export type MatchingStrategy = { txid: Txid | Array<Txid> } | void
 
 /**
  * Type representing a snapshot end message
@@ -56,7 +63,6 @@ type MaybeTxId =
 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
@@ -74,17 +80,109 @@ type InferSchemaOutput<T> = T extends StandardSchemaV1
 export interface ElectricCollectionConfig<
   T extends Row<unknown> = Row<unknown>,
   TSchema extends StandardSchemaV1 = never,
-> extends BaseCollectionConfig<
-    T,
-    string | number,
-    TSchema,
-    Record<string, Fn>,
-    { txid: Txid | Array<Txid> }
+> extends Omit<
+    BaseCollectionConfig<T, string | number, TSchema, UtilsRecord, any>,
+    `onInsert` | `onUpdate` | `onDelete`
   > {
   /**
    * Configuration options for the ElectricSQL ShapeStream
    */
   shapeOptions: ShapeStreamOptions<GetExtensions<T>>
+
+  /**
+   * Optional asynchronous handler function called before an insert operation
+   * @param params Object containing transaction and collection information
+   * @returns Promise resolving to { txid } or void
+   * @example
+   * // Basic Electric insert handler with txid (recommended)
+   * onInsert: async ({ transaction }) => {
+   *   const newItem = transaction.mutations[0].modified
+   *   const result = await api.todos.create({
+   *     data: newItem
+   *   })
+   *   return { txid: result.txid }
+   * }
+   *
+   * @example
+   * // Insert handler with multiple items - return array of txids
+   * onInsert: async ({ transaction }) => {
+   *   const items = transaction.mutations.map(m => m.modified)
+   *   const results = await Promise.all(
+   *     items.map(item => api.todos.create({ data: item }))
+   *   )
+   *   return { txid: results.map(r => r.txid) }
+   * }
+   *
+   * @example
+   * // Use awaitMatch utility for custom matching
+   * onInsert: async ({ transaction, collection }) => {
+   *   const newItem = transaction.mutations[0].modified
+   *   await api.todos.create({ data: newItem })
+   *   await collection.utils.awaitMatch(
+   *     (message) => isChangeMessage(message) &&
+   *                  message.headers.operation === 'insert' &&
+   *                  message.value.name === newItem.name
+   *   )
+   * }
+   */
+  onInsert?: (params: InsertMutationFnParams<T>) => Promise<MatchingStrategy>
+
+  /**
+   * Optional asynchronous handler function called before an update operation
+   * @param params Object containing transaction and collection information
+   * @returns Promise resolving to { txid } or void
+   * @example
+   * // Basic Electric update handler with txid (recommended)
+   * onUpdate: async ({ transaction }) => {
+   *   const { original, changes } = transaction.mutations[0]
+   *   const result = await api.todos.update({
+   *     where: { id: original.id },
+   *     data: changes
+   *   })
+   *   return { txid: result.txid }
+   * }
+   *
+   * @example
+   * // Use awaitMatch utility for custom matching
+   * onUpdate: async ({ transaction, collection }) => {
+   *   const { original, changes } = transaction.mutations[0]
+   *   await api.todos.update({ where: { id: original.id }, data: changes })
+   *   await collection.utils.awaitMatch(
+   *     (message) => isChangeMessage(message) &&
+   *                  message.headers.operation === 'update' &&
+   *                  message.value.id === original.id
+   *   )
+   * }
+   */
+  onUpdate?: (params: UpdateMutationFnParams<T>) => Promise<MatchingStrategy>
+
+  /**
+   * Optional asynchronous handler function called before a delete operation
+   * @param params Object containing transaction and collection information
+   * @returns Promise resolving to { txid } or void
+   * @example
+   * // Basic Electric delete handler with txid (recommended)
+   * onDelete: async ({ transaction }) => {
+   *   const mutation = transaction.mutations[0]
+   *   const result = await api.todos.delete({
+   *     id: mutation.original.id
+   *   })
+   *   return { txid: result.txid }
+   * }
+   *
+   * @example
+   * // Use awaitMatch utility for custom matching
+   * onDelete: async ({ transaction, collection }) => {
+   *   const mutation = transaction.mutations[0]
+   *   await api.todos.delete({ id: mutation.original.id })
+   *   await collection.utils.awaitMatch(
+   *     (message) => isChangeMessage(message) &&
+   *                  message.headers.operation === 'delete' &&
+   *                  message.value.id === mutation.original.id
+   *   )
+   * }
+   */
+  onDelete?: (params: DeleteMutationFnParams<T>) => Promise<MatchingStrategy>
 }
 
 function isUpToDateMessage<T extends Row<unknown>>(
@@ -125,11 +223,21 @@ function hasTxids<T extends Row<unknown>>(
  */
 export type AwaitTxIdFn = (txId: Txid, timeout?: number) => Promise<boolean>
 
+/**
+ * Type for the awaitMatch utility function
+ */
+export type AwaitMatchFn<T extends Row<unknown>> = (
+  matchFn: MatchFunction<T>,
+  timeout?: number
+) => Promise<boolean>
+
 /**
  * Electric collection utilities type
  */
-export interface ElectricCollectionUtils extends UtilsRecord {
+export interface ElectricCollectionUtils<T extends Row<unknown> = Row<unknown>>
+  extends UtilsRecord {
   awaitTxId: AwaitTxIdFn
+  awaitMatch: AwaitMatchFn<T>
 }
 
 /**
@@ -173,24 +281,79 @@ export function electricCollectionOptions(
 } {
   const seenTxids = new Store<Set<Txid>>(new Set([]))
   const seenSnapshots = new Store<Array<PostgresSnapshot>>([])
+  const pendingMatches = new Store<
+    Map<
+      string,
+      {
+        matchFn: (message: Message<any>) => boolean
+        resolve: (value: boolean) => void
+        reject: (error: Error) => void
+        timeoutId: ReturnType<typeof setTimeout>
+        matched: boolean
+      }
+    >
+  >(new Map())
+
+  // Buffer messages since last up-to-date to handle race conditions
+  const currentBatchMessages = new Store<Array<Message<any>>>([])
+
+  /**
+   * Helper function to remove multiple matches from the pendingMatches store
+   */
+  const removePendingMatches = (matchIds: Array<string>) => {
+    if (matchIds.length > 0) {
+      pendingMatches.setState((current) => {
+        const newMatches = new Map(current)
+        matchIds.forEach((id) => newMatches.delete(id))
+        return newMatches
+      })
+    }
+  }
+
+  /**
+   * Helper function to resolve and cleanup matched pending matches
+   */
+  const resolveMatchedPendingMatches = () => {
+    const matchesToResolve: Array<string> = []
+    pendingMatches.state.forEach((match, matchId) => {
+      if (match.matched) {
+        clearTimeout(match.timeoutId)
+        match.resolve(true)
+        matchesToResolve.push(matchId)
+        debug(
+          `${config.id ? `[${config.id}] ` : ``}awaitMatch resolved on up-to-date for match %s`,
+          matchId
+        )
+      }
+    })
+    removePendingMatches(matchesToResolve)
+  }
   const sync = createElectricSync<any>(config.shapeOptions, {
     seenTxids,
     seenSnapshots,
+    pendingMatches,
+    currentBatchMessages,
+    removePendingMatches,
+    resolveMatchedPendingMatches,
+    collectionId: config.id,
   })
 
   /**
    * Wait for a specific transaction ID to be synced
    * @param txId The transaction ID to wait for as a number
-   * @param timeout Optional timeout in milliseconds (defaults to 30000ms)
+   * @param timeout Optional timeout in milliseconds (defaults to 5000ms)
    * @returns Promise that resolves when the txId is synced
    */
   const awaitTxId: AwaitTxIdFn = async (
     txId: Txid,
-    timeout: number = 30000
+    timeout: number = 5000
   ): Promise<boolean> => {
-    debug(`awaitTxId called with txid %d`, txId)
+    debug(
+      `${config.id ? `[${config.id}] ` : ``}awaitTxId called with txid %d`,
+      txId
+    )
     if (typeof txId !== `number`) {
-      throw new ExpectedNumberInAwaitTxIdError(typeof txId)
+      throw new ExpectedNumberInAwaitTxIdError(typeof txId, config.id)
     }
 
     // First check if the txid is in the seenTxids store
@@ -207,12 +370,15 @@ export function electricCollectionOptions(
       const timeoutId = setTimeout(() => {
         unsubscribeSeenTxids()
         unsubscribeSeenSnapshots()
-        reject(new TimeoutWaitingForTxIdError(txId))
+        reject(new TimeoutWaitingForTxIdError(txId, config.id))
       }, timeout)
 
       const unsubscribeSeenTxids = seenTxids.subscribe(() => {
         if (seenTxids.state.has(txId)) {
-          debug(`awaitTxId found match for txid %o`, txId)
+          debug(
+            `${config.id ? `[${config.id}] ` : ``}awaitTxId found match for txid %o`,
+            txId
+          )
           clearTimeout(timeoutId)
           unsubscribeSeenTxids()
           unsubscribeSeenSnapshots()
@@ -226,7 +392,7 @@ export function electricCollectionOptions(
         )
         if (visibleSnapshot) {
           debug(
-            `awaitTxId found match for txid %o in snapshot %o`,
+            `${config.id ? `[${config.id}] ` : ``}awaitTxId found match for txid %o in snapshot %o`,
             txId,
             visibleSnapshot
           )
@@ -239,49 +405,128 @@ export function electricCollectionOptions(
     })
   }
 
-  // Create wrapper handlers for direct persistence operations that handle txid awaiting
-  const wrappedOnInsert = config.onInsert
-    ? async (params: InsertMutationFnParams<any>) => {
-        // Runtime check (that doesn't follow type)
+  /**
+   * Wait for a custom match function to find a matching message
+   * @param matchFn Function that returns true when a message matches
+   * @param timeout Optional timeout in milliseconds (defaults to 5000ms)
+   * @returns Promise that resolves when a matching message is found
+   */
+  const awaitMatch: AwaitMatchFn<any> = async (
+    matchFn: MatchFunction<any>,
+    timeout: number = 3000
+  ): Promise<boolean> => {
+    debug(
+      `${config.id ? `[${config.id}] ` : ``}awaitMatch called with custom function`
+    )
+
+    return new Promise((resolve, reject) => {
+      const matchId = Math.random().toString(36)
+
+      const cleanupMatch = () => {
+        pendingMatches.setState((current) => {
+          const newMatches = new Map(current)
+          newMatches.delete(matchId)
+          return newMatches
+        })
+      }
+
+      const onTimeout = () => {
+        cleanupMatch()
+        reject(new TimeoutWaitingForMatchError(config.id))
+      }
 
-        const handlerResult =
-          ((await config.onInsert!(params)) as MaybeTxId) ?? {}
-        const txid = handlerResult.txid
+      const timeoutId = setTimeout(onTimeout, timeout)
 
-        if (!txid) {
-          throw new ElectricInsertHandlerMustReturnTxIdError()
+      // We need access to the stream messages to check against the match function
+      // This will be handled by the sync configuration
+      const checkMatch = (message: Message<any>) => {
+        if (matchFn(message)) {
+          debug(
+            `${config.id ? `[${config.id}] ` : ``}awaitMatch found matching message, waiting for up-to-date`
+          )
+          // Mark as matched but don't resolve yet - wait for up-to-date
+          pendingMatches.setState((current) => {
+            const newMatches = new Map(current)
+            const existing = newMatches.get(matchId)
+            if (existing) {
+              newMatches.set(matchId, { ...existing, matched: true })
+            }
+            return newMatches
+          })
+          return true
         }
+        return false
+      }
 
-        // Handle both single txid and array of txids
-        if (Array.isArray(txid)) {
-          await Promise.all(txid.map((id) => awaitTxId(id)))
-        } else {
-          await awaitTxId(txid)
+      // Check against current batch messages first to handle race conditions
+      for (const message of currentBatchMessages.state) {
+        if (matchFn(message)) {
+          debug(
+            `${config.id ? `[${config.id}] ` : ``}awaitMatch found immediate match in current batch, waiting for up-to-date`
+          )
+          // Register match as already matched
+          pendingMatches.setState((current) => {
+            const newMatches = new Map(current)
+            newMatches.set(matchId, {
+              matchFn: checkMatch,
+              resolve,
+              reject,
+              timeoutId,
+              matched: true, // Already matched
+            })
+            return newMatches
+          })
+          return
         }
+      }
+
+      // Store the match function for the sync process to use
+      // We'll add this to a pending matches store
+      pendingMatches.setState((current) => {
+        const newMatches = new Map(current)
+        newMatches.set(matchId, {
+          matchFn: checkMatch,
+          resolve,
+          reject,
+          timeoutId,
+          matched: false,
+        })
+        return newMatches
+      })
+    })
+  }
 
+  /**
+   * Process matching strategy and wait for synchronization
+   */
+  const processMatchingStrategy = async (
+    result: MatchingStrategy
+  ): Promise<void> => {
+    // Only wait if result contains txid
+    if (result && `txid` in result) {
+      // Handle both single txid and array of txids
+      if (Array.isArray(result.txid)) {
+        await Promise.all(result.txid.map(awaitTxId))
+      } else {
+        await awaitTxId(result.txid)
+      }
+    }
+    // If result is void/undefined, don't wait - mutation completes immediately
+  }
+
+  // Create wrapper handlers for direct persistence operations that handle different matching strategies
+  const wrappedOnInsert = config.onInsert
+    ? async (params: InsertMutationFnParams<any>) => {
+        const handlerResult = await config.onInsert!(params)
+        await processMatchingStrategy(handlerResult)
         return handlerResult
       }
     : undefined
 
   const wrappedOnUpdate = config.onUpdate
     ? async (params: UpdateMutationFnParams<any>) => {
-        // Runtime check (that doesn't follow type)
-
-        const handlerResult =
-          ((await config.onUpdate!(params)) as MaybeTxId) ?? {}
-        const txid = handlerResult.txid
-
-        if (!txid) {
-          throw new ElectricUpdateHandlerMustReturnTxIdError()
-        }
-
-        // Handle both single txid and array of txids
-        if (Array.isArray(txid)) {
-          await Promise.all(txid.map((id) => awaitTxId(id)))
-        } else {
-          await awaitTxId(txid)
-        }
-
+        const handlerResult = await config.onUpdate!(params)
+        await processMatchingStrategy(handlerResult)
         return handlerResult
       }
     : undefined
@@ -289,17 +534,7 @@ export function electricCollectionOptions(
   const wrappedOnDelete = config.onDelete
     ? async (params: DeleteMutationFnParams<any>) => {
         const handlerResult = await config.onDelete!(params)
-        if (!handlerResult.txid) {
-          throw new ElectricDeleteHandlerMustReturnTxIdError()
-        }
-
-        // Handle both single txid and array of txids
-        if (Array.isArray(handlerResult.txid)) {
-          await Promise.all(handlerResult.txid.map((id) => awaitTxId(id)))
-        } else {
-          await awaitTxId(handlerResult.txid)
-        }
-
+        await processMatchingStrategy(handlerResult)
         return handlerResult
       }
     : undefined
@@ -321,7 +556,8 @@ export function electricCollectionOptions(
     onDelete: wrappedOnDelete,
     utils: {
       awaitTxId,
-    },
+      awaitMatch,
+    } as ElectricCollectionUtils<any>,
   }
 }
 
@@ -333,10 +569,34 @@ function createElectricSync<T extends Row<unknown>>(
   options: {
     seenTxids: Store<Set<Txid>>
     seenSnapshots: Store<Array<PostgresSnapshot>>
+    pendingMatches: Store<
+      Map<
+        string,
+        {
+          matchFn: (message: Message<T>) => boolean
+          resolve: (value: boolean) => void
+          reject: (error: Error) => void
+          timeoutId: ReturnType<typeof setTimeout>
+          matched: boolean
+        }
+      >
+    >
+    currentBatchMessages: Store<Array<Message<T>>>
+    removePendingMatches: (matchIds: Array<string>) => void
+    resolveMatchedPendingMatches: () => void
+    collectionId?: string
   }
 ): SyncConfig<T> {
-  const { seenTxids } = options
-  const { seenSnapshots } = options
+  const {
+    seenTxids,
+    seenSnapshots,
+    pendingMatches,
+    currentBatchMessages,
+    removePendingMatches,
+    resolveMatchedPendingMatches,
+    collectionId,
+  } = options
+  const MAX_BATCH_MESSAGES = 1000 // Safety limit for message buffer
 
   // Store for the relation schema information
   const relationSchema = new Store<string | undefined>(undefined)
@@ -380,6 +640,17 @@ function createElectricSync<T extends Row<unknown>>(
         }
       }
 
+      // Cleanup pending matches on abort
+      abortController.signal.addEventListener(`abort`, () => {
+        pendingMatches.setState((current) => {
+          current.forEach((match) => {
+            clearTimeout(match.timeoutId)
+            match.reject(new StreamAbortedError())
+          })
+          return new Map() // Clear all pending matches
+        })
+      })
+
       const stream = new ShapeStream({
         ...shapeOptions,
         signal: abortController.signal,
@@ -413,11 +684,45 @@ function createElectricSync<T extends Row<unknown>>(
         let hasUpToDate = false
 
         for (const message of messages) {
+          // Add message to current batch buffer (for race condition handling)
+          if (isChangeMessage(message)) {
+            currentBatchMessages.setState((currentBuffer) => {
+              const newBuffer = [...currentBuffer, message]
+              // Limit buffer size for safety
+              if (newBuffer.length > MAX_BATCH_MESSAGES) {
+                newBuffer.splice(0, newBuffer.length - MAX_BATCH_MESSAGES)
+              }
+              return newBuffer
+            })
+          }
+
           // Check for txids in the message and add them to our store
           if (hasTxids(message)) {
             message.headers.txids?.forEach((txid) => newTxids.add(txid))
           }
 
+          // Check pending matches against this message
+          // Note: matchFn will mark matches internally, we don't resolve here
+          const matchesToRemove: Array<string> = []
+          pendingMatches.state.forEach((match, matchId) => {
+            if (!match.matched) {
+              try {
+                match.matchFn(message)
+              } catch (err) {
+                // If matchFn throws, clean up and reject the promise
+                clearTimeout(match.timeoutId)
+                match.reject(
+                  err instanceof Error ? err : new Error(String(err))
+                )
+                matchesToRemove.push(matchId)
+                debug(`matchFn error: %o`, err)
+              }
+            }
+          })
+
+          // Remove matches that errored
+          removePendingMatches(matchesToRemove)
+
           if (isChangeMessage(message)) {
             // Check if the message contains schema information
             const schema = message.headers.schema
@@ -445,7 +750,7 @@ function createElectricSync<T extends Row<unknown>>(
             hasUpToDate = true
           } else if (isMustRefetchMessage(message)) {
             debug(
-              `Received must-refetch message, starting transaction with truncate`
+              `${collectionId ? `[${collectionId}] ` : ``}Received must-refetch message, starting transaction with truncate`
             )
 
             // Start a transaction and truncate the collection
@@ -462,6 +767,9 @@ function createElectricSync<T extends Row<unknown>>(
         }
 
         if (hasUpToDate) {
+          // Clear the current batch buffer since we're now up-to-date
+          currentBatchMessages.setState(() => [])
+
           // Commit transaction if one was started
           if (transactionStarted) {
             commit()
@@ -475,7 +783,10 @@ function createElectricSync<T extends Row<unknown>>(
           seenTxids.setState((currentTxids) => {
             const clonedSeen = new Set<Txid>(currentTxids)
             if (newTxids.size > 0) {
-              debug(`new txids synced from pg %O`, Array.from(newTxids))
+              debug(
+                `${collectionId ? `[${collectionId}] ` : ``}new txids synced from pg %O`,
+                Array.from(newTxids)
+              )
             }
             newTxids.forEach((txid) => clonedSeen.add(txid))
             newTxids.clear()
@@ -486,11 +797,17 @@ function createElectricSync<T extends Row<unknown>>(
           seenSnapshots.setState((currentSnapshots) => {
             const seen = [...currentSnapshots, ...newSnapshots]
             newSnapshots.forEach((snapshot) =>
-              debug(`new snapshot synced from pg %o`, snapshot)
+              debug(
+                `${collectionId ? `[${collectionId}] ` : ``}new snapshot synced from pg %o`,
+                snapshot
+              )
             )
             newSnapshots.length = 0
             return seen
           })
+
+          // Resolve all matched pending matches on up-to-date
+          resolveMatchedPendingMatches()
         }
       })