Track state of rows in TX to be able to delete rows on a move out

Author: kevin-dp

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

@@ -883,7 +883,7 @@ function createElectricSync<T extends Row<unknown>>(
       }
 
       rowTagSet.add(tag)
-      addTagToIndex(parsedTag, rowId, tagIndex, tagLength)
+      addTagToIndex(parsedTag, rowId, tagIndex, tagLength)      
     }
   }
 
@@ -1015,6 +1015,7 @@ function createElectricSync<T extends Row<unknown>>(
     begin: () => void,
     write: (message: Omit<ChangeMessage<T>, `key`>) => void,
     transactionStarted: boolean,
+    transactionState: Map<RowId, T>
   ): boolean => {
     if (tagLength === undefined) {
       debug(
@@ -1038,12 +1039,15 @@ function createElectricSync<T extends Row<unknown>>(
             txStarted = true
           }
 
-          const rowValue = collection.get(rowId)
+          // Get row value from transaction state (uncommitted) or collection
+          const rowValue = transactionState.get(rowId) ?? collection.get(rowId)
           if (rowValue !== undefined) {
             write({
               type: `delete`,
               value: rowValue,
             })
+            // Remove from transaction state since we're deleting it
+            transactionState.delete(rowId)
           }
         }
       }
@@ -1164,6 +1168,10 @@ function createElectricSync<T extends Row<unknown>>(
         syncMode === `progressive` && !hasReceivedUpToDate
       const bufferedMessages: Array<Message<T>> = [] // Buffer change messages during initial sync
 
+      // Track row state during the current transaction to access uncommitted row values
+      // This allows us to handle partial updates correctly by merging with existing state
+      const transactionState = new Map<RowId, T>()
+
       /**
        * Process a change message: handle tags and write the mutation
        */
@@ -1178,20 +1186,62 @@ function createElectricSync<T extends Row<unknown>>(
         const hasTags = tags || removedTags
 
         const rowId = collection.getKeyFromItem(changeMessage.value)
+        const operation = changeMessage.headers.operation
+
+        if (operation === `insert`) {
+          // For insert, store the full row in transaction state
+          transactionState.set(rowId, changeMessage.value)
+
+          if (hasTags) {
+            processTagsForChangeMessage(tags, removedTags, rowId)
+          }
 
-        if (changeMessage.headers.operation === `delete`) {
+          write({
+            type: `insert`,
+            value: changeMessage.value,
+            metadata: {
+              ...changeMessage.headers,
+            },
+          })
+        } else if (operation === `update`) {
+          // For update, merge with existing state (from transaction state or collection)
+          const existingValue =
+            transactionState.get(rowId) ?? collection.get(rowId)
+
+          // Merge the update with existing value (handles partial updates)
+          const updatedValue =
+            existingValue !== undefined
+              ? Object.assign({}, existingValue, changeMessage.value)
+              : changeMessage.value
+
+          // Store the merged result in transaction state
+          transactionState.set(rowId, updatedValue)
+
+          if (hasTags) {
+            processTagsForChangeMessage(tags, removedTags, rowId)
+          }
+
+          write({
+            type: `update`,
+            value: updatedValue,
+            metadata: {
+              ...changeMessage.headers,
+            },
+          })
+        } else {
+          // Operation is delete
           clearTagsForRow(rowId)
-        } else if (hasTags) {
-          processTagsForChangeMessage(tags, removedTags, rowId)
+          // Remove from transaction state
+          transactionState.delete(rowId)
+
+          write({
+            type: `delete`,
+            value: changeMessage.value,
+            metadata: {
+              ...changeMessage.headers,
+            },
+          })
         }
-
-        write({
-          type: changeMessage.headers.operation,
-          value: changeMessage.value,
-          metadata: {
-            ...changeMessage.headers,
-          },
-        })
       }
 
       // Create deduplicated loadSubset wrapper for non-eager modes
@@ -1213,7 +1263,7 @@ function createElectricSync<T extends Row<unknown>>(
 
         for (const message of messages) {
           // Add message to current batch buffer (for race condition handling)
-          if (isChangeMessage(message)) {
+          if (isChangeMessage(message) || isMoveOutMessage(message)) {
             currentBatchMessages.setState((currentBuffer) => {
               const newBuffer = [...currentBuffer, message]
               // Limit buffer size for safety
@@ -1281,14 +1331,20 @@ function createElectricSync<T extends Row<unknown>>(
           } else if (isUpToDateMessage(message)) {
             hasUpToDate = true
           } else if (isMoveOutMessage(message)) {
-            // Handle move-out event: remove matching tags from rows
-            transactionStarted = processMoveOutEvent(
-              message.headers.patterns,
-              collection,
-              begin,
-              write,
-              transactionStarted,
-            )
+            // Handle move-out event: buffer if buffering, otherwise process immediately
+            if (isBufferingInitialSync()) {
+              bufferedMessages.push(message)
+            } else {
+              // Normal processing: process move-out immediately
+              transactionStarted = processMoveOutEvent(
+                message.headers.patterns,
+                collection,
+                begin,
+                write,
+                transactionStarted,
+                transactionState
+              )
+            }
           } else if (isMustRefetchMessage(message)) {
             debug(
               `${collectionId ? `[${collectionId}] ` : ``}Received must-refetch message, starting transaction with truncate`,
@@ -1347,11 +1403,22 @@ function createElectricSync<T extends Row<unknown>>(
               } else if (isSnapshotEndMessage(bufferedMsg)) {
                 // Extract snapshots from buffered messages (will be committed to store after transaction)
                 newSnapshots.push(parseSnapshotMessage(bufferedMsg))
+              } else if (isMoveOutMessage(bufferedMsg)) {
+                // Process buffered move-out messages during atomic swap
+                processMoveOutEvent(
+                  bufferedMsg.headers.patterns,
+                  collection,
+                  begin,
+                  write,
+                  transactionStarted,
+                  transactionState
+                )
               }
             }
 
             // Commit the atomic swap
             commit()
+            transactionState.clear()
 
             // Exit buffering phase by marking that we've received up-to-date
             // isBufferingInitialSync() will now return false
@@ -1371,6 +1438,7 @@ function createElectricSync<T extends Row<unknown>>(
             if (transactionStarted && shouldCommit) {
               commit()
               transactionStarted = false
+              transactionState.clear()
             }
           }