From 78924fe26341152e576bd146e6f575faa4cef334 Mon Sep 17 00:00:00 2001
From: Sam Willis <sam.willis@gmail.com>
Date: Sun, 8 Dec 2024 14:32:04 +0000
Subject: [PATCH 01/12] Add `commitGranularity` parameter in the
 `syncShapeToTable`

---
 .changeset/happy-jokes-yawn.md    |   5 +
 packages/pglite-sync/src/index.ts | 298 +++++++++++++++++++-----------
 2 files changed, 191 insertions(+), 112 deletions(-)
 create mode 100644 .changeset/happy-jokes-yawn.md

diff --git a/.changeset/happy-jokes-yawn.md b/.changeset/happy-jokes-yawn.md
new file mode 100644
index 000000000..6756f0f48
--- /dev/null
+++ b/.changeset/happy-jokes-yawn.md
@@ -0,0 +1,5 @@
+---
+'@electric-sql/pglite-sync': patch
+---
+
+Add options for the `commitGranularity` parameter in the `syncShapeToTable` function, enabling the user to choose how often the sync should commit.
diff --git a/packages/pglite-sync/src/index.ts b/packages/pglite-sync/src/index.ts
index 8d88fe7f5..5cd61be2c 100644
--- a/packages/pglite-sync/src/index.ts
+++ b/packages/pglite-sync/src/index.ts
@@ -21,6 +21,20 @@ type InsertChangeMessage = ChangeMessage<any> & {
   headers: { operation: 'insert' }
 }
 
+/**
+ * The granularity of the commit operation.
+ * - `up-to-date`: Commit all messages when the `up-to-date` message is received.
+ * - `transaction`: Commit all messages within transactions inferred from the LSN prefix of the offset.
+ * - `operation`: Commit each message in its own transaction.
+ * - `number`: Commit every N messages.
+ * Note a commit will always be performed on the `up-to-date` message.
+ */
+export type CommitGranularity =
+  | 'up-to-date'
+  | 'transaction'
+  | 'operation'
+  | number
+
 export interface SyncShapeToTableOptions {
   shape: ShapeStreamOptions
   table: string
@@ -29,6 +43,9 @@ export interface SyncShapeToTableOptions {
   primaryKey: string[]
   shapeKey?: ShapeKey
   useCopy?: boolean
+  commitGranularity?: CommitGranularity
+  commitThrottle?: number
+  onInitialSync?: () => void
 }
 
 export interface SyncShapeToTableResult {
@@ -64,6 +81,11 @@ async function createPlugin(
     syncShapeToTable: async (
       options: SyncShapeToTableOptions,
     ): Promise<SyncShapeToTableResult> => {
+      await firstRun()
+      options = {
+        commitGranularity: 'up-to-date',
+        ...options,
+      }
       if (shapePerTableLock.has(options.table)) {
         throw new Error('Already syncing shape for table ' + options.table)
       }
@@ -113,123 +135,165 @@ async function createPlugin(
       // or use a separate connection to hold a long transaction
       let messageAggregator: ChangeMessage<any>[] = []
       let truncateNeeded = false
+      let lastLSN: string | null = null
+      let lastCommitAt: number = 0
+
+      const commit = async () => {
+        if (messageAggregator.length === 0 && !truncateNeeded) return
+        await pg.transaction(async (tx) => {
+          if (debug) {
+            console.log('committing message batch', messageAggregator.length)
+            console.time('commit')
+          }
+
+          // Set the syncing flag to true during this transaction so that
+          // user defined triggers on the table are able to chose how to run
+          // during a sync
+          tx.exec(`SET LOCAL ${metadataSchema}.syncing = true;`)
+
+          if (truncateNeeded) {
+            truncateNeeded = false
+            // TODO: sync into shadow table and reference count
+            // for now just clear the whole table - will break
+            // cases with multiple shapes on the same table
+            await tx.exec(`DELETE FROM ${options.table};`)
+            if (options.shapeKey) {
+              await deleteShapeSubscriptionState({
+                pg: tx,
+                metadataSchema,
+                shapeKey: options.shapeKey,
+              })
+            }
+          }
+
+          if (doCopy) {
+            // We can do a `COPY FROM` to insert the initial data
+            // Split messageAggregator into initial inserts and remaining messages
+            const initialInserts: InsertChangeMessage[] = []
+            const remainingMessages: ChangeMessage<any>[] = []
+            let foundNonInsert = false
+            for (const message of messageAggregator) {
+              if (!foundNonInsert && message.headers.operation === 'insert') {
+                initialInserts.push(message as InsertChangeMessage)
+              } else {
+                foundNonInsert = true
+                remainingMessages.push(message)
+              }
+            }
+            if (initialInserts.length > 0) {
+              // As `COPY FROM` doesn't trigger a NOTIFY, we pop
+              // the last insert message and and add it to the be beginning
+              // of the remaining messages to be applied after the `COPY FROM`
+              remainingMessages.unshift(initialInserts.pop()!)
+            }
+            messageAggregator = remainingMessages
+
+            // Do the `COPY FROM` with initial inserts
+            if (initialInserts.length > 0) {
+              applyMessagesToTableWithCopy({
+                pg: tx,
+                table: options.table,
+                schema: options.schema,
+                messages: initialInserts as InsertChangeMessage[],
+                mapColumns: options.mapColumns,
+                primaryKey: options.primaryKey,
+                debug,
+              })
+              // We don't want to do a `COPY FROM` again after that
+              doCopy = false
+            }
+          }
+
+          for (const changeMessage of messageAggregator) {
+            await applyMessageToTable({
+              pg: tx,
+              table: options.table,
+              schema: options.schema,
+              message: changeMessage,
+              mapColumns: options.mapColumns,
+              primaryKey: options.primaryKey,
+              debug,
+            })
+          }
+
+          if (
+            options.shapeKey &&
+            messageAggregator.length > 0 &&
+            stream.shapeHandle !== undefined
+          ) {
+            await updateShapeSubscriptionState({
+              pg: tx,
+              metadataSchema,
+              shapeKey: options.shapeKey,
+              shapeId: stream.shapeHandle,
+              lastOffset:
+                messageAggregator[messageAggregator.length - 1].offset,
+            })
+          }
+        })
+        if (debug) console.timeEnd('commit')
+        messageAggregator = []
+        // Await a timeout to start a new task and  allow other connections to do work
+        await new Promise((resolve) => setTimeout(resolve, 0))
+      }
+
+      const throttledCommit = async () => {
+        if (
+          options.commitThrottle &&
+          Date.now() - lastCommitAt < options.commitThrottle
+        ) {
+          return
+        }
+        lastCommitAt = Date.now()
+        await commit()
+      }
 
       stream.subscribe(async (messages) => {
         if (debug) console.log('sync messages received', messages)
 
         for (const message of messages) {
-          // accumulate change messages for committing all at once
           if (isChangeMessage(message)) {
-            messageAggregator.push(message)
-            continue
-          }
-
-          // perform actual DB operations upon receiving control messages
-          if (!isControlMessage(message)) continue
-          switch (message.headers.control) {
-            // mark table as needing truncation before next batch commit
-            case 'must-refetch':
-              if (debug) console.log('refetching shape')
-              truncateNeeded = true
-              messageAggregator = []
-
-              break
-
-            // perform all accumulated changes and store stream state
-            case 'up-to-date':
-              await pg.transaction(async (tx) => {
-                if (debug) console.log('up-to-date, committing all messages')
-
-                // Set the syncing flag to true during this transaction so that
-                // user defined triggers on the table are able to chose how to run
-                // during a sync
-                tx.exec(`SET LOCAL ${metadataSchema}.syncing = true;`)
-
-                if (truncateNeeded) {
-                  truncateNeeded = false
-                  // TODO: sync into shadow table and reference count
-                  // for now just clear the whole table - will break
-                  // cases with multiple shapes on the same table
-                  await tx.exec(`DELETE FROM ${options.table};`)
-                  if (options.shapeKey) {
-                    await deleteShapeSubscriptionState({
-                      pg: tx,
-                      metadataSchema,
-                      shapeKey: options.shapeKey,
-                    })
-                  }
-                }
-
-                if (doCopy) {
-                  // We can do a `COPY FROM` to insert the initial data
-                  // Split messageAggregator into initial inserts and remaining messages
-                  const initialInserts: InsertChangeMessage[] = []
-                  const remainingMessages: ChangeMessage<any>[] = []
-                  let foundNonInsert = false
-                  for (const message of messageAggregator) {
-                    if (
-                      !foundNonInsert &&
-                      message.headers.operation === 'insert'
-                    ) {
-                      initialInserts.push(message as InsertChangeMessage)
-                    } else {
-                      foundNonInsert = true
-                      remainingMessages.push(message)
-                    }
-                  }
-                  if (initialInserts.length > 0) {
-                    // As `COPY FROM` doesn't trigger a NOTIFY, we pop
-                    // the last insert message and and add it to the be beginning
-                    // of the remaining messages to be applied after the `COPY FROM`
-                    remainingMessages.unshift(initialInserts.pop()!)
-                  }
-                  messageAggregator = remainingMessages
-
-                  // Do the `COPY FROM` with initial inserts
-                  if (initialInserts.length > 0) {
-                    applyMessagesToTableWithCopy({
-                      pg: tx,
-                      table: options.table,
-                      schema: options.schema,
-                      messages: initialInserts as InsertChangeMessage[],
-                      mapColumns: options.mapColumns,
-                      primaryKey: options.primaryKey,
-                      debug,
-                    })
-                    // We don't want to do a `COPY FROM` again after that
-                    doCopy = false
-                  }
-                }
+            const newLSN = message.offset.split('_')[0]
+            if (newLSN !== lastLSN) {
+              // If the LSN has changed and granularity is set to transaction
+              // we need to commit the current batch.
+              // This is done before we accumulate any more messages as they are
+              // part of the next transaction batch.
+              if (options.commitGranularity === 'transaction') {
+                await throttledCommit()
+              }
+              lastLSN = newLSN
+            }
 
-                for (const changeMessage of messageAggregator) {
-                  await applyMessageToTable({
-                    pg: tx,
-                    table: options.table,
-                    schema: options.schema,
-                    message: changeMessage,
-                    mapColumns: options.mapColumns,
-                    primaryKey: options.primaryKey,
-                    debug,
-                  })
-                }
+            // accumulate change messages for committing all at once or in batches
+            messageAggregator.push(message)
 
-                if (
-                  options.shapeKey &&
-                  messageAggregator.length > 0 &&
-                  stream.shapeHandle !== undefined
-                ) {
-                  await updateShapeSubscriptionState({
-                    pg: tx,
-                    metadataSchema,
-                    shapeKey: options.shapeKey,
-                    shapeId: stream.shapeHandle,
-                    lastOffset:
-                      messageAggregator[messageAggregator.length - 1].offset,
-                  })
+            if (options.commitGranularity === 'operation') {
+              // commit after each operation if granularity is set to operation
+              await throttledCommit()
+            } else if (typeof options.commitGranularity === 'number') {
+              // commit after every N messages if granularity is set to a number
+              if (messageAggregator.length >= options.commitGranularity) {
+                await throttledCommit()
+              }
+            }
+          } else if (isControlMessage(message)) {
+            switch (message.headers.control) {
+              case 'must-refetch':
+                // mark table as needing truncation before next batch commit
+                if (debug) console.log('refetching shape')
+                truncateNeeded = true
+                messageAggregator = []
+                break
+
+              case 'up-to-date':
+                // perform all accumulated changes and store stream state
+                await commit() // not throttled, we want this to happen ASAP
+                if (isNewSubscription && options.onInitialSync) {
+                  options.onInitialSync()
                 }
-              })
-              messageAggregator = []
-              break
+                break
+            }
           }
         }
       })
@@ -270,7 +334,11 @@ async function createPlugin(
     }
   }
 
-  const init = async () => {
+  let firstRunDone = false
+
+  const firstRun = async () => {
+    if (firstRunDone) return
+    firstRunDone = true
     await migrateShapeMetadataTables({
       pg,
       metadataSchema,
@@ -280,19 +348,25 @@ async function createPlugin(
   return {
     namespaceObj,
     close,
-    init,
   }
 }
 
+export type SyncNamespaceObj = Awaited<
+  ReturnType<typeof createPlugin>
+>['namespaceObj']
+
+export type PGliteWithSync = PGliteInterface & {
+  sync: SyncNamespaceObj
+}
+
 export function electricSync(options?: ElectricSyncOptions) {
   return {
     name: 'ElectricSQL Sync',
     setup: async (pg: PGliteInterface) => {
-      const { namespaceObj, close, init } = await createPlugin(pg, options)
+      const { namespaceObj, close } = await createPlugin(pg, options)
       return {
         namespaceObj,
         close,
-        init,
       }
     },
   } satisfies Extension

From c29cd8b61c06999d9eb35738b43ed12a8b3b19a0 Mon Sep 17 00:00:00 2001
From: Sam Willis <sam.willis@gmail.com>
Date: Sun, 8 Dec 2024 14:53:15 +0000
Subject: [PATCH 02/12] Fix tests

---
 packages/pglite-sync/src/index.ts      | 24 ++++++++++++------------
 packages/pglite-sync/test/sync.test.ts |  3 ++-
 2 files changed, 14 insertions(+), 13 deletions(-)

diff --git a/packages/pglite-sync/src/index.ts b/packages/pglite-sync/src/index.ts
index 5cd61be2c..8cec03f44 100644
--- a/packages/pglite-sync/src/index.ts
+++ b/packages/pglite-sync/src/index.ts
@@ -77,11 +77,22 @@ async function createPlugin(
   // resolved by using reference counting in shadow tables
   const shapePerTableLock = new Map<string, void>()
 
+  let initMetadataTablesDone = false
+  const initMetadataTables = async () => {
+    if (initMetadataTablesDone) return
+    initMetadataTablesDone = true
+    await migrateShapeMetadataTables({
+      pg,
+      metadataSchema,
+    })
+  }
+
   const namespaceObj = {
+    initMetadataTables,
     syncShapeToTable: async (
       options: SyncShapeToTableOptions,
     ): Promise<SyncShapeToTableResult> => {
-      await firstRun()
+      await initMetadataTables()
       options = {
         commitGranularity: 'up-to-date',
         ...options,
@@ -334,17 +345,6 @@ async function createPlugin(
     }
   }
 
-  let firstRunDone = false
-
-  const firstRun = async () => {
-    if (firstRunDone) return
-    firstRunDone = true
-    await migrateShapeMetadataTables({
-      pg,
-      metadataSchema,
-    })
-  }
-
   return {
     namespaceObj,
     close,
diff --git a/packages/pglite-sync/test/sync.test.ts b/packages/pglite-sync/test/sync.test.ts
index 2f03baa51..82de3e810 100644
--- a/packages/pglite-sync/test/sync.test.ts
+++ b/packages/pglite-sync/test/sync.test.ts
@@ -371,6 +371,7 @@ describe('pglite-sync', () => {
         }),
       },
     })
+    await db.electric.initMetadataTables()
 
     const result = await db.query(
       `SELECT schema_name FROM information_schema.schemata WHERE schema_name = $1`,
@@ -533,7 +534,7 @@ describe('pglite-sync', () => {
     // Check the flag is not set outside of a sync
     const result0 =
       await pg.sql`SELECT current_setting('electric.syncing', true)`
-    expect(result0.rows[0]).toEqual({ current_setting: 'false' })
+    expect(result0.rows[0]).toEqual({ current_setting: null }) // not set yet as syncShapeToTable hasn't been called
 
     const shape = await pg.electric.syncShapeToTable({
       shape: {

From 5a6bf0c8cab75f2f4ed80bf8a5176558ff8b57e0 Mon Sep 17 00:00:00 2001
From: Sam Willis <sam.willis@gmail.com>
Date: Sun, 8 Dec 2024 15:17:41 +0000
Subject: [PATCH 03/12] Docs

---
 docs/docs/sync.md                 | 14 ++++++++++++++
 packages/pglite-sync/src/index.ts |  2 +-
 2 files changed, 15 insertions(+), 1 deletion(-)

diff --git a/docs/docs/sync.md b/docs/docs/sync.md
index 3462fe621..70d0cd74c 100644
--- a/docs/docs/sync.md
+++ b/docs/docs/sync.md
@@ -94,6 +94,20 @@ It takes the following options as an object:
 - `useCopy: boolean`<br>
   Whether to use the `COPY FROM` command to insert the initial data, defaults to `false`. This process may be faster than inserting row by row as it combines the inserts into a CSV to be passed to Postgres.
 
+- `commitGranularity: CommitGranularity`<br>
+  The granularity of the commit operation, defaults to `"up-to-date"`. Note that a commit will always be performed immediately on the `up-to-date` message.
+  Options:
+  - `"up-to-date"`: Commit all messages when the `up-to-date` message is received.
+  - `"transaction"`: Commit all messages within transactions as they were applied to the source Postgres.
+  - `"operation"`: Commit each message in its own transaction.
+  - `number`: Commit every N messages.  
+
+- `commitThrottle: number`<br>
+  The number of milliseconds to wait between commits, defaults to `0`.
+
+- `onInitialSync: () => void`<br>
+  A callback that is called when the initial sync is complete.
+
 The returned `shape` object from the `syncShapeToTable` call has the following methods:
 
 - `isUpToDate: boolean`<br>
diff --git a/packages/pglite-sync/src/index.ts b/packages/pglite-sync/src/index.ts
index 8cec03f44..70069d8a7 100644
--- a/packages/pglite-sync/src/index.ts
+++ b/packages/pglite-sync/src/index.ts
@@ -24,7 +24,7 @@ type InsertChangeMessage = ChangeMessage<any> & {
 /**
  * The granularity of the commit operation.
  * - `up-to-date`: Commit all messages when the `up-to-date` message is received.
- * - `transaction`: Commit all messages within transactions inferred from the LSN prefix of the offset.
+ * - `transaction`: Commit all messages within transactions as they were applied to the source Postgres.
  * - `operation`: Commit each message in its own transaction.
  * - `number`: Commit every N messages.
  * Note a commit will always be performed on the `up-to-date` message.

From 1990212e2a242fec2adc8469daf720e74a7d7894 Mon Sep 17 00:00:00 2001
From: Sam Willis <sam.willis@gmail.com>
Date: Sun, 8 Dec 2024 15:48:46 +0000
Subject: [PATCH 04/12] add tests for commitGranularity

---
 packages/pglite-sync/test/sync.test.ts | 336 +++++++++++++++++++++++++
 1 file changed, 336 insertions(+)

diff --git a/packages/pglite-sync/test/sync.test.ts b/packages/pglite-sync/test/sync.test.ts
index 82de3e810..bc0e82202 100644
--- a/packages/pglite-sync/test/sync.test.ts
+++ b/packages/pglite-sync/test/sync.test.ts
@@ -725,4 +725,340 @@ describe('pglite-sync', () => {
 
     shape.unsubscribe()
   })
+
+  it('respects numeric batch commit granularity settings', async () => {
+    let feedMessages: (messages: Message[]) => Promise<void> = async (_) => {}
+    MockShapeStream.mockImplementation(() => ({
+      subscribe: vi.fn((cb: (messages: Message[]) => Promise<void>) => {
+        feedMessages = (messages) => cb([...messages, upToDateMsg])
+      }),
+      unsubscribeAll: vi.fn(),
+    }))
+
+    // Create a trigger to notify on transaction commit
+    await pg.exec(`
+      CREATE OR REPLACE FUNCTION notify_transaction()
+      RETURNS TRIGGER AS $$
+      BEGIN
+        PERFORM pg_notify('transaction_commit', TG_TABLE_NAME);
+        RETURN NEW;
+      END;
+      $$ LANGUAGE plpgsql;
+
+      CREATE TRIGGER todo_transaction_trigger
+      AFTER INSERT ON todo
+      FOR EACH STATEMENT
+      EXECUTE FUNCTION notify_transaction();
+    `)
+
+    const commits: string[] = []
+    const unsubscribe = await pg.listen('transaction_commit', (payload) => {
+      commits.push(payload)
+    })
+
+    const batchSize = 5
+    const shape = await pg.electric.syncShapeToTable({
+      shape: { url: 'http://localhost:3000/v1/shape', table: 'todo' },
+      table: 'todo',
+      primaryKey: ['id'],
+      commitGranularity: batchSize,
+    })
+
+    // Create test messages - 7 total (should see batch of 5, then 2)
+    const messages = Array.from(
+      { length: 7 },
+      (_, idx) =>
+        ({
+          headers: { operation: 'insert' },
+          offset: `1_${idx}`,
+          key: `id${idx}`,
+          value: {
+            id: idx,
+            task: `task${idx}`,
+            done: false,
+          },
+        }) satisfies Message,
+    )
+
+    await feedMessages(messages)
+
+    // Wait for all inserts to complete
+    await vi.waitUntil(async () => {
+      const result = await pg.sql<{ count: number }>`
+        SELECT COUNT(*) as count FROM todo;
+      `
+      return result.rows[0].count === 7
+    })
+
+    // Verify all rows were inserted
+    const result = await pg.sql`
+      SELECT * FROM todo ORDER BY id;
+    `
+    expect(result.rows).toEqual(
+      messages.map((m) => ({
+        id: m.value.id,
+        task: m.value.task,
+        done: m.value.done,
+      })),
+    )
+
+    // Should have received 2 commit notifications:
+    // - One for the first batch of 5
+    // - One for the remaining 2 (triggered by up-to-date message)
+    expect(commits).toHaveLength(2)
+    expect(commits).toEqual(['todo', 'todo'])
+
+    await unsubscribe()
+    shape.unsubscribe()
+  })
+
+  it('respects transaction commit granularity', async () => {
+    let feedMessages: (messages: Message[]) => Promise<void> = async (_) => {}
+    MockShapeStream.mockImplementation(() => ({
+      subscribe: vi.fn((cb: (messages: Message[]) => Promise<void>) => {
+        feedMessages = (messages) => cb([...messages, upToDateMsg])
+      }),
+      unsubscribeAll: vi.fn(),
+    }))
+
+    // Create a trigger to notify on transaction commit
+    await pg.exec(`
+      CREATE OR REPLACE FUNCTION notify_transaction()
+      RETURNS TRIGGER AS $$
+      BEGIN
+        PERFORM pg_notify('transaction_commit', TG_TABLE_NAME);
+        RETURN NEW;
+      END;
+      $$ LANGUAGE plpgsql;
+
+      CREATE TRIGGER todo_transaction_trigger
+      AFTER INSERT ON todo
+      FOR EACH STATEMENT
+      EXECUTE FUNCTION notify_transaction();
+    `)
+
+    // Track transaction commits
+    const transactionCommits: string[] = []
+    const unsubscribe = await pg.listen('transaction_commit', (payload) => {
+      transactionCommits.push(payload)
+    })
+
+    const shape = await pg.electric.syncShapeToTable({
+      shape: { url: 'http://localhost:3000/v1/shape', table: 'todo' },
+      table: 'todo',
+      primaryKey: ['id'],
+      commitGranularity: 'transaction',
+    })
+
+    // Send messages with different LSNs (first part of offset before _)
+    await feedMessages([
+      {
+        headers: { operation: 'insert' },
+        offset: '1_1', // Transaction 1
+        key: 'id1',
+        value: {
+          id: 1,
+          task: 'task1',
+          done: false,
+        },
+      },
+      {
+        headers: { operation: 'insert' },
+        offset: '1_2', // Same transaction
+        key: 'id2',
+        value: {
+          id: 2,
+          task: 'task2',
+          done: false,
+        },
+      },
+      {
+        headers: { operation: 'insert' },
+        offset: '2_1', // New transaction
+        key: 'id3',
+        value: {
+          id: 3,
+          task: 'task3',
+          done: false,
+        },
+      },
+    ])
+
+    // Wait for all inserts to complete
+    await vi.waitUntil(async () => {
+      const result = await pg.sql<{ count: number }>`
+        SELECT COUNT(*) as count FROM todo;
+      `
+      return result.rows[0].count === 3
+    })
+
+    // Verify all rows were inserted
+    const result = await pg.sql`
+      SELECT * FROM todo ORDER BY id;
+    `
+    expect(result.rows).toEqual([
+      { id: 1, task: 'task1', done: false },
+      { id: 2, task: 'task2', done: false },
+      { id: 3, task: 'task3', done: false },
+    ])
+
+    // Should have received 2 transaction notifications
+    // One for LSN 1 (containing 2 inserts) and one for LSN 2 (containing 1 insert)
+    expect(transactionCommits).toHaveLength(2)
+    expect(transactionCommits).toEqual(['todo', 'todo'])
+
+    await unsubscribe()
+    shape.unsubscribe()
+  })
+
+  it('respects up-to-date commit granularity settings', async () => {
+    let feedMessages: (messages: Message[]) => Promise<void> = async (_) => {}
+    MockShapeStream.mockImplementation(() => ({
+      subscribe: vi.fn((cb: (messages: Message[]) => Promise<void>) => {
+        feedMessages = (messages) => cb([...messages, upToDateMsg])
+      }),
+      unsubscribeAll: vi.fn(),
+    }))
+
+    // Create a trigger to notify on transaction commit
+    await pg.exec(`
+      CREATE OR REPLACE FUNCTION notify_transaction()
+      RETURNS TRIGGER AS $$
+      BEGIN
+        PERFORM pg_notify('transaction_commit', TG_TABLE_NAME);
+        RETURN NEW;
+      END;
+      $$ LANGUAGE plpgsql;
+
+      CREATE TRIGGER todo_transaction_trigger
+      AFTER INSERT ON todo
+      FOR EACH STATEMENT
+      EXECUTE FUNCTION notify_transaction();
+    `)
+
+    const commits: string[] = []
+    const unsubscribe = await pg.listen('transaction_commit', (payload) => {
+      commits.push(payload)
+    })
+
+    const shape = await pg.electric.syncShapeToTable({
+      shape: { url: 'http://localhost:3000/v1/shape', table: 'todo' },
+      table: 'todo',
+      primaryKey: ['id'],
+      commitGranularity: 'up-to-date',
+    })
+
+    // Send multiple messages
+    await feedMessages([
+      {
+        headers: { operation: 'insert' },
+        offset: '1_1',
+        key: 'id1',
+        value: { id: 1, task: 'task1', done: false },
+      },
+      {
+        headers: { operation: 'insert' },
+        offset: '2_1',
+        key: 'id2',
+        value: { id: 2, task: 'task2', done: false },
+      },
+      {
+        headers: { operation: 'insert' },
+        offset: '3_1',
+        key: 'id3',
+        value: { id: 3, task: 'task3', done: false },
+      },
+    ])
+
+    // Wait for all inserts to complete
+    await vi.waitUntil(async () => {
+      const result = await pg.sql<{ count: number }>`
+        SELECT COUNT(*) as count FROM todo;
+      `
+      return result.rows[0].count === 3
+    })
+
+    // Should have received only one commit notification since all operations
+    // are committed together when up-to-date message is received
+    expect(commits).toHaveLength(1)
+    expect(commits).toEqual(['todo'])
+
+    await unsubscribe()
+    shape.unsubscribe()
+  })
+
+  it('respects operation commit granularity settings', async () => {
+    let feedMessages: (messages: Message[]) => Promise<void> = async (_) => {}
+    MockShapeStream.mockImplementation(() => ({
+      subscribe: vi.fn((cb: (messages: Message[]) => Promise<void>) => {
+        feedMessages = (messages) => cb([...messages, upToDateMsg])
+      }),
+      unsubscribeAll: vi.fn(),
+    }))
+
+    // Create a trigger to notify on transaction commit
+    await pg.exec(`
+      CREATE OR REPLACE FUNCTION notify_transaction()
+      RETURNS TRIGGER AS $$
+      BEGIN
+        PERFORM pg_notify('transaction_commit', TG_TABLE_NAME);
+        RETURN NEW;
+      END;
+      $$ LANGUAGE plpgsql;
+
+      CREATE TRIGGER todo_transaction_trigger
+      AFTER INSERT ON todo
+      FOR EACH STATEMENT
+      EXECUTE FUNCTION notify_transaction();
+    `)
+
+    const commits: string[] = []
+    const unsubscribe = await pg.listen('transaction_commit', (payload) => {
+      commits.push(payload)
+    })
+
+    const shape = await pg.electric.syncShapeToTable({
+      shape: { url: 'http://localhost:3000/v1/shape', table: 'todo' },
+      table: 'todo',
+      primaryKey: ['id'],
+      commitGranularity: 'operation',
+    })
+
+    // Send multiple messages
+    await feedMessages([
+      {
+        headers: { operation: 'insert' },
+        offset: '1_1',
+        key: 'id1',
+        value: { id: 1, task: 'task1', done: false },
+      },
+      {
+        headers: { operation: 'insert' },
+        offset: '1_2',
+        key: 'id2',
+        value: { id: 2, task: 'task2', done: false },
+      },
+      {
+        headers: { operation: 'insert' },
+        offset: '1_3',
+        key: 'id3',
+        value: { id: 3, task: 'task3', done: false },
+      },
+    ])
+
+    // Wait for all inserts to complete
+    await vi.waitUntil(async () => {
+      const result = await pg.sql<{ count: number }>`
+        SELECT COUNT(*) as count FROM todo;
+      `
+      return result.rows[0].count === 3
+    })
+
+    // Should have received a notification for each operation
+    expect(commits).toHaveLength(3)
+    expect(commits).toEqual(['todo', 'todo', 'todo'])
+
+    await unsubscribe()
+    shape.unsubscribe()
+  })
 })

From e786972757246c3fbb833616d56ff5b2148ace6f Mon Sep 17 00:00:00 2001
From: Sam Willis <sam.willis@gmail.com>
Date: Sun, 8 Dec 2024 16:20:50 +0000
Subject: [PATCH 05/12] Add test of the commitThrottle option

---
 packages/pglite-sync/src/index.ts      |  16 +++-
 packages/pglite-sync/test/sync.test.ts | 110 +++++++++++++++++++++++++
 2 files changed, 124 insertions(+), 2 deletions(-)

diff --git a/packages/pglite-sync/src/index.ts b/packages/pglite-sync/src/index.ts
index 70069d8a7..6a684b2e1 100644
--- a/packages/pglite-sync/src/index.ts
+++ b/packages/pglite-sync/src/index.ts
@@ -249,13 +249,25 @@ async function createPlugin(
       }
 
       const throttledCommit = async () => {
+        const now = Date.now()
+        if (options.commitThrottle && debug)
+          console.log(
+            'throttled commit: now:',
+            now,
+            'lastCommitAt:',
+            lastCommitAt,
+            'diff:',
+            now - lastCommitAt,
+          )
         if (
           options.commitThrottle &&
-          Date.now() - lastCommitAt < options.commitThrottle
+          now - lastCommitAt < options.commitThrottle
         ) {
+          // Skip this commit - messages will be caught by next commit or up-to-date
+          if (debug) console.log('skipping commit due to throttle')
           return
         }
-        lastCommitAt = Date.now()
+        lastCommitAt = now
         await commit()
       }
 
diff --git a/packages/pglite-sync/test/sync.test.ts b/packages/pglite-sync/test/sync.test.ts
index bc0e82202..6a77e926e 100644
--- a/packages/pglite-sync/test/sync.test.ts
+++ b/packages/pglite-sync/test/sync.test.ts
@@ -1061,4 +1061,114 @@ describe('pglite-sync', () => {
     await unsubscribe()
     shape.unsubscribe()
   })
+
+  // Skip this test as it's flaky in CI, timing is sensitive
+  it.skip('respects commitThrottle with operation commit granularity', async () => {
+    let feedMessages: (messages: Message[]) => Promise<void> = async (_) => {}
+    MockShapeStream.mockImplementation(() => ({
+      subscribe: vi.fn((cb: (messages: Message[]) => Promise<void>) => {
+        feedMessages = (messages) => cb([...messages])
+      }),
+      unsubscribeAll: vi.fn(),
+    }))
+
+    // Create a trigger to notify on transaction commit
+    await pg.exec(`
+      CREATE OR REPLACE FUNCTION notify_transaction()
+      RETURNS TRIGGER AS $$
+      BEGIN
+        PERFORM pg_notify('transaction_commit', 
+          TG_TABLE_NAME || '_' || 
+          (SELECT COUNT(*) FROM todo)::text || '_' ||
+          EXTRACT(MILLISECONDS FROM NOW())::text
+        );
+        RETURN NEW;
+      END;
+      $$ LANGUAGE plpgsql;
+
+      CREATE TRIGGER todo_transaction_trigger
+      AFTER INSERT ON todo
+      FOR EACH STATEMENT
+      EXECUTE FUNCTION notify_transaction();
+    `)
+
+    const commits: string[] = []
+    const unsubscribe = await pg.listen('transaction_commit', (payload) => {
+      commits.push(payload)
+    })
+
+    const throttleMs = 15 // Short throttle for testing
+    const shape = await pg.electric.syncShapeToTable({
+      shape: { url: 'http://localhost:3000/v1/shape', table: 'todo' },
+      table: 'todo',
+      primaryKey: ['id'],
+      commitGranularity: 'operation',
+      commitThrottle: throttleMs,
+    })
+
+    // Send messages with 10ms delays between them
+    for (const message of [
+      {
+        headers: { operation: 'insert' as const },
+        offset: '1_1' as const,
+        key: 'id1',
+        value: { id: 1, task: 'task1', done: false },
+      },
+      {
+        headers: { operation: 'insert' as const },
+        offset: '1_2' as const, 
+        key: 'id2',
+        value: { id: 2, task: 'task2', done: false },
+      },
+      {
+        headers: { operation: 'insert' as const },
+        offset: '1_3' as const,
+        key: 'id3',
+        value: { id: 3, task: 'task3', done: false },
+      },
+      {
+        headers: { operation: 'insert' as const },
+        offset: '1_4' as const,
+        key: 'id4',
+        value: { id: 4, task: 'task4', done: false },
+      },
+      upToDateMsg,
+    ]) {
+      await feedMessages([message])
+      await new Promise(resolve => setTimeout(resolve, 10))
+    }
+
+    // Wait for all inserts to complete
+    await vi.waitUntil(async () => {
+      const result = await pg.sql<{ count: number }>`
+        SELECT COUNT(*) as count FROM todo;
+      `
+      return result.rows[0].count === 4
+    })
+
+    console.log(commits)
+
+    // Extract row counts and timestamps from commit notifications
+    const commitInfo = commits.map(commit => {
+      const [_, rowCount, timestamp] = commit.split('_')
+      return {
+        rowCount: parseInt(rowCount),
+        timestamp: parseFloat(timestamp)
+      }
+    })
+
+    // Verify we got 4 operation messages
+    expect(commitInfo.length).toBe(4)
+
+    // Check timestamps are at least 15ms apart for first 3
+    expect(commitInfo[1].timestamp - commitInfo[0].timestamp).toBeGreaterThanOrEqual(15)
+    expect(commitInfo[2].timestamp - commitInfo[1].timestamp).toBeGreaterThanOrEqual(15) 
+
+    // Last 2 operation messages should have same timestamp since they're batched
+    expect(commitInfo[3].timestamp).toBe(commitInfo[2].timestamp)
+
+    await unsubscribe()
+    shape.unsubscribe()
+  })
 })
+

From f69b7b24d5df416d93bd6084540bcd11597afe83 Mon Sep 17 00:00:00 2001
From: Sam Willis <sam.willis@gmail.com>
Date: Sun, 8 Dec 2024 16:26:08 +0000
Subject: [PATCH 06/12] Fix onInitialSync and add test

---
 packages/pglite-sync/src/index.ts      | 10 ++-
 packages/pglite-sync/test/sync.test.ts | 87 +++++++++++++++++++++++---
 2 files changed, 89 insertions(+), 8 deletions(-)

diff --git a/packages/pglite-sync/src/index.ts b/packages/pglite-sync/src/index.ts
index 6a684b2e1..29fceb4de 100644
--- a/packages/pglite-sync/src/index.ts
+++ b/packages/pglite-sync/src/index.ts
@@ -125,6 +125,9 @@ async function createPlugin(
       // may overlap and so the insert logic will be wrong.
       let doCopy = isNewSubscription && options.useCopy
 
+      // Track if onInitialSync has been called
+      let onInitialSyncCalled = false
+
       const aborter = new AbortController()
       if (options.shape.signal) {
         // we new to have our own aborter to be able to abort the stream
@@ -312,8 +315,13 @@ async function createPlugin(
               case 'up-to-date':
                 // perform all accumulated changes and store stream state
                 await commit() // not throttled, we want this to happen ASAP
-                if (isNewSubscription && options.onInitialSync) {
+                if (
+                  isNewSubscription &&
+                  !onInitialSyncCalled &&
+                  options.onInitialSync
+                ) {
                   options.onInitialSync()
+                  onInitialSyncCalled = true
                 }
                 break
             }
diff --git a/packages/pglite-sync/test/sync.test.ts b/packages/pglite-sync/test/sync.test.ts
index 6a77e926e..72c9744c2 100644
--- a/packages/pglite-sync/test/sync.test.ts
+++ b/packages/pglite-sync/test/sync.test.ts
@@ -1116,7 +1116,7 @@ describe('pglite-sync', () => {
       },
       {
         headers: { operation: 'insert' as const },
-        offset: '1_2' as const, 
+        offset: '1_2' as const,
         key: 'id2',
         value: { id: 2, task: 'task2', done: false },
       },
@@ -1135,7 +1135,7 @@ describe('pglite-sync', () => {
       upToDateMsg,
     ]) {
       await feedMessages([message])
-      await new Promise(resolve => setTimeout(resolve, 10))
+      await new Promise((resolve) => setTimeout(resolve, 10))
     }
 
     // Wait for all inserts to complete
@@ -1149,11 +1149,11 @@ describe('pglite-sync', () => {
     console.log(commits)
 
     // Extract row counts and timestamps from commit notifications
-    const commitInfo = commits.map(commit => {
+    const commitInfo = commits.map((commit) => {
       const [_, rowCount, timestamp] = commit.split('_')
       return {
         rowCount: parseInt(rowCount),
-        timestamp: parseFloat(timestamp)
+        timestamp: parseFloat(timestamp),
       }
     })
 
@@ -1161,8 +1161,12 @@ describe('pglite-sync', () => {
     expect(commitInfo.length).toBe(4)
 
     // Check timestamps are at least 15ms apart for first 3
-    expect(commitInfo[1].timestamp - commitInfo[0].timestamp).toBeGreaterThanOrEqual(15)
-    expect(commitInfo[2].timestamp - commitInfo[1].timestamp).toBeGreaterThanOrEqual(15) 
+    expect(
+      commitInfo[1].timestamp - commitInfo[0].timestamp,
+    ).toBeGreaterThanOrEqual(15)
+    expect(
+      commitInfo[2].timestamp - commitInfo[1].timestamp,
+    ).toBeGreaterThanOrEqual(15)
 
     // Last 2 operation messages should have same timestamp since they're batched
     expect(commitInfo[3].timestamp).toBe(commitInfo[2].timestamp)
@@ -1170,5 +1174,74 @@ describe('pglite-sync', () => {
     await unsubscribe()
     shape.unsubscribe()
   })
-})
 
+  it('calls onInitialSync callback after initial sync', async () => {
+    let feedMessages: (messages: Message[]) => Promise<void> = async (_) => {}
+    MockShapeStream.mockImplementation(() => ({
+      subscribe: vi.fn((cb: (messages: Message[]) => Promise<void>) => {
+        feedMessages = (messages) => cb([...messages, upToDateMsg])
+      }),
+      unsubscribeAll: vi.fn(),
+    }))
+
+    const onInitialSync = vi.fn()
+    const shape = await pg.electric.syncShapeToTable({
+      shape: { url: 'http://localhost:3000/v1/shape', table: 'todo' },
+      table: 'todo',
+      primaryKey: ['id'],
+      onInitialSync,
+    })
+
+    // Send some initial data
+    await feedMessages([
+      {
+        headers: { operation: 'insert' },
+        offset: '1_1',
+        key: 'id1',
+        value: {
+          id: 1,
+          task: 'task1',
+          done: false,
+        },
+      },
+      {
+        headers: { operation: 'insert' },
+        offset: '1_2',
+        key: 'id2',
+        value: {
+          id: 2,
+          task: 'task2',
+          done: true,
+        },
+      },
+    ])
+
+    // Verify callback was called once
+    expect(onInitialSync).toHaveBeenCalledTimes(1)
+
+    // Send more data - callback should not be called again
+    await feedMessages([
+      {
+        headers: { operation: 'insert' },
+        offset: '1_3',
+        key: 'id3',
+        value: {
+          id: 3,
+          task: 'task3',
+          done: false,
+        },
+      },
+    ])
+
+    // Verify callback was still only called once
+    expect(onInitialSync).toHaveBeenCalledTimes(1)
+
+    // Verify all data was inserted
+    expect(
+      (await pg.sql<{ count: number }>`SELECT COUNT(*) as count FROM todo;`)
+        .rows[0].count,
+    ).toBe(3)
+
+    shape.unsubscribe()
+  })
+})

From 1ffac5763a5d9afa7f4c4e238b2d9f4da0e60e18 Mon Sep 17 00:00:00 2001
From: Yacine <ybsoft2011@live.com>
Date: Sun, 8 Dec 2024 15:53:50 +0100
Subject: [PATCH 07/12] doc: export return type for syncShapeToTable (#452)

* doc: export return type for syncShapeToTable

* feat: expose stream object as of for useShape hook

* Changeset

---------

Co-authored-by: Sam Willis <sam.willis@gmail.com>
---
 packages/pglite-sync/src/index.ts | 8 ++++++++
 1 file changed, 8 insertions(+)

diff --git a/packages/pglite-sync/src/index.ts b/packages/pglite-sync/src/index.ts
index 29fceb4de..82f4d3718 100644
--- a/packages/pglite-sync/src/index.ts
+++ b/packages/pglite-sync/src/index.ts
@@ -56,6 +56,14 @@ export interface SyncShapeToTableResult {
   stream: ShapeStreamInterface
 }
 
+export interface SyncShapeToTableResult {
+  unsubscribe: () => void
+  readonly isUpToDate: boolean
+  readonly shapeId: string
+  subscribe: (cb: () => void, error: (err: Error) => void) => () => void
+  stream: ShapeStreamInterface
+}
+
 export interface ElectricSyncOptions {
   debug?: boolean
   metadataSchema?: string

From bf150359d37f2262a4b9310c1d2e56a46fbbadba Mon Sep 17 00:00:00 2001
From: Sam Willis <sam.willis@gmail.com>
Date: Sun, 8 Dec 2024 18:42:43 +0000
Subject: [PATCH 08/12] Remove part of test thats flaky on CI

---
 packages/pglite-sync/test/sync.test.ts | 12 ++++++------
 1 file changed, 6 insertions(+), 6 deletions(-)

diff --git a/packages/pglite-sync/test/sync.test.ts b/packages/pglite-sync/test/sync.test.ts
index 72c9744c2..e945f88dd 100644
--- a/packages/pglite-sync/test/sync.test.ts
+++ b/packages/pglite-sync/test/sync.test.ts
@@ -153,11 +153,11 @@ describe('pglite-sync', () => {
       )
     }
 
-    let timeToProcessMicrotask = Infinity
-    const startTime = performance.now()
-    Promise.resolve().then(() => {
-      timeToProcessMicrotask = performance.now() - startTime
-    })
+    // let timeToProcessMicrotask = Infinity
+    // const startTime = performance.now()
+    // Promise.resolve().then(() => {
+    //   timeToProcessMicrotask = performance.now() - startTime
+    // })
 
     let numItemsInserted = 0
     await vi.waitUntil(async () => {
@@ -175,7 +175,7 @@ describe('pglite-sync', () => {
     expect(numItemsInserted).toBe(numInserts)
 
     // should have processed microtask within few ms, not blocking main loop
-    expect(timeToProcessMicrotask).toBeLessThan(15)
+    // expect(timeToProcessMicrotask).toBeLessThan(15) // TODO: flaky on CI
 
     await shape.unsubscribe()
   })

From 8b2568fc04eee2a172ba28b8ef52940600fa9231 Mon Sep 17 00:00:00 2001
From: Sam Willis <sam.willis@gmail.com>
Date: Sun, 8 Dec 2024 19:22:27 +0000
Subject: [PATCH 09/12] formatting

---
 docs/docs/sync.md | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/docs/docs/sync.md b/docs/docs/sync.md
index 70d0cd74c..67a99029d 100644
--- a/docs/docs/sync.md
+++ b/docs/docs/sync.md
@@ -97,10 +97,11 @@ It takes the following options as an object:
 - `commitGranularity: CommitGranularity`<br>
   The granularity of the commit operation, defaults to `"up-to-date"`. Note that a commit will always be performed immediately on the `up-to-date` message.
   Options:
+
   - `"up-to-date"`: Commit all messages when the `up-to-date` message is received.
   - `"transaction"`: Commit all messages within transactions as they were applied to the source Postgres.
   - `"operation"`: Commit each message in its own transaction.
-  - `number`: Commit every N messages.  
+  - `number`: Commit every N messages.
 
 - `commitThrottle: number`<br>
   The number of milliseconds to wait between commits, defaults to `0`.

From 4278d46b6fcd96f77f51df44e5b354f8b1fe9caa Mon Sep 17 00:00:00 2001
From: Sam Willis <sam.willis@gmail.com>
Date: Sun, 8 Dec 2024 19:35:35 +0000
Subject: [PATCH 10/12] Fix types

---
 packages/pglite-sync/test/sync.test.ts | 30 ++++++++++++++++++++------
 1 file changed, 24 insertions(+), 6 deletions(-)

diff --git a/packages/pglite-sync/test/sync.test.ts b/packages/pglite-sync/test/sync.test.ts
index e945f88dd..9f8560914 100644
--- a/packages/pglite-sync/test/sync.test.ts
+++ b/packages/pglite-sync/test/sync.test.ts
@@ -758,7 +758,10 @@ describe('pglite-sync', () => {
 
     const batchSize = 5
     const shape = await pg.electric.syncShapeToTable({
-      shape: { url: 'http://localhost:3000/v1/shape', table: 'todo' },
+      shape: {
+        url: 'http://localhost:3000/v1/shape',
+        params: { table: 'todo' },
+      },
       table: 'todo',
       primaryKey: ['id'],
       commitGranularity: batchSize,
@@ -844,7 +847,10 @@ describe('pglite-sync', () => {
     })
 
     const shape = await pg.electric.syncShapeToTable({
-      shape: { url: 'http://localhost:3000/v1/shape', table: 'todo' },
+      shape: {
+        url: 'http://localhost:3000/v1/shape',
+        params: { table: 'todo' },
+      },
       table: 'todo',
       primaryKey: ['id'],
       commitGranularity: 'transaction',
@@ -942,7 +948,10 @@ describe('pglite-sync', () => {
     })
 
     const shape = await pg.electric.syncShapeToTable({
-      shape: { url: 'http://localhost:3000/v1/shape', table: 'todo' },
+      shape: {
+        url: 'http://localhost:3000/v1/shape',
+        params: { table: 'todo' },
+      },
       table: 'todo',
       primaryKey: ['id'],
       commitGranularity: 'up-to-date',
@@ -1018,7 +1027,10 @@ describe('pglite-sync', () => {
     })
 
     const shape = await pg.electric.syncShapeToTable({
-      shape: { url: 'http://localhost:3000/v1/shape', table: 'todo' },
+      shape: {
+        url: 'http://localhost:3000/v1/shape',
+        params: { table: 'todo' },
+      },
       table: 'todo',
       primaryKey: ['id'],
       commitGranularity: 'operation',
@@ -1099,7 +1111,10 @@ describe('pglite-sync', () => {
 
     const throttleMs = 15 // Short throttle for testing
     const shape = await pg.electric.syncShapeToTable({
-      shape: { url: 'http://localhost:3000/v1/shape', table: 'todo' },
+      shape: {
+        url: 'http://localhost:3000/v1/shape',
+        params: { table: 'todo' },
+      },
       table: 'todo',
       primaryKey: ['id'],
       commitGranularity: 'operation',
@@ -1186,7 +1201,10 @@ describe('pglite-sync', () => {
 
     const onInitialSync = vi.fn()
     const shape = await pg.electric.syncShapeToTable({
-      shape: { url: 'http://localhost:3000/v1/shape', table: 'todo' },
+      shape: {
+        url: 'http://localhost:3000/v1/shape',
+        params: { table: 'todo' },
+      },
       table: 'todo',
       primaryKey: ['id'],
       onInitialSync,

From b37985e138443a2074a9e10c2330ecaa717290f1 Mon Sep 17 00:00:00 2001
From: Sam Willis <sam.willis@gmail.com>
Date: Sun, 8 Dec 2024 22:24:54 +0000
Subject: [PATCH 11/12] Remove transaction commit granularity until Electric
 has stabilised on LSN metadata

---
 docs/docs/sync.md                      |   2 +-
 packages/pglite-sync/src/index.ts      |  28 ++--
 packages/pglite-sync/test/sync.test.ts | 203 +++++++++++++------------
 3 files changed, 117 insertions(+), 116 deletions(-)

diff --git a/docs/docs/sync.md b/docs/docs/sync.md
index 67a99029d..d52a47bfc 100644
--- a/docs/docs/sync.md
+++ b/docs/docs/sync.md
@@ -99,7 +99,7 @@ It takes the following options as an object:
   Options:
 
   - `"up-to-date"`: Commit all messages when the `up-to-date` message is received.
-  - `"transaction"`: Commit all messages within transactions as they were applied to the source Postgres.
+  <!-- - `"transaction"`: Commit all messages within transactions as they were applied to the source Postgres. -->
   - `"operation"`: Commit each message in its own transaction.
   - `number`: Commit every N messages.
 
diff --git a/packages/pglite-sync/src/index.ts b/packages/pglite-sync/src/index.ts
index 82f4d3718..b2a8607b2 100644
--- a/packages/pglite-sync/src/index.ts
+++ b/packages/pglite-sync/src/index.ts
@@ -24,14 +24,13 @@ type InsertChangeMessage = ChangeMessage<any> & {
 /**
  * The granularity of the commit operation.
  * - `up-to-date`: Commit all messages when the `up-to-date` message is received.
- * - `transaction`: Commit all messages within transactions as they were applied to the source Postgres.
  * - `operation`: Commit each message in its own transaction.
  * - `number`: Commit every N messages.
  * Note a commit will always be performed on the `up-to-date` message.
  */
 export type CommitGranularity =
   | 'up-to-date'
-  | 'transaction'
+  // | 'transaction'  // Removed until Electric has stabilised on LSN metadata
   | 'operation'
   | number
 
@@ -157,7 +156,7 @@ async function createPlugin(
       // or use a separate connection to hold a long transaction
       let messageAggregator: ChangeMessage<any>[] = []
       let truncateNeeded = false
-      let lastLSN: string | null = null
+      // let lastLSN: string | null = null  // Removed until Electric has stabilised on LSN metadata
       let lastCommitAt: number = 0
 
       const commit = async () => {
@@ -287,17 +286,18 @@ async function createPlugin(
 
         for (const message of messages) {
           if (isChangeMessage(message)) {
-            const newLSN = message.offset.split('_')[0]
-            if (newLSN !== lastLSN) {
-              // If the LSN has changed and granularity is set to transaction
-              // we need to commit the current batch.
-              // This is done before we accumulate any more messages as they are
-              // part of the next transaction batch.
-              if (options.commitGranularity === 'transaction') {
-                await throttledCommit()
-              }
-              lastLSN = newLSN
-            }
+            // Removed until Electric has stabilised on LSN metadata
+            // const newLSN = message.offset.split('_')[0]
+            // if (newLSN !== lastLSN) {
+            //   // If the LSN has changed and granularity is set to transaction
+            //   // we need to commit the current batch.
+            //   // This is done before we accumulate any more messages as they are
+            //   // part of the next transaction batch.
+            //   if (options.commitGranularity === 'transaction') {
+            //     await throttledCommit()
+            //   }
+            //   lastLSN = newLSN
+            // }
 
             // accumulate change messages for committing all at once or in batches
             messageAggregator.push(message)
diff --git a/packages/pglite-sync/test/sync.test.ts b/packages/pglite-sync/test/sync.test.ts
index 9f8560914..60e209edc 100644
--- a/packages/pglite-sync/test/sync.test.ts
+++ b/packages/pglite-sync/test/sync.test.ts
@@ -815,107 +815,108 @@ describe('pglite-sync', () => {
     shape.unsubscribe()
   })
 
-  it('respects transaction commit granularity', async () => {
-    let feedMessages: (messages: Message[]) => Promise<void> = async (_) => {}
-    MockShapeStream.mockImplementation(() => ({
-      subscribe: vi.fn((cb: (messages: Message[]) => Promise<void>) => {
-        feedMessages = (messages) => cb([...messages, upToDateMsg])
-      }),
-      unsubscribeAll: vi.fn(),
-    }))
-
-    // Create a trigger to notify on transaction commit
-    await pg.exec(`
-      CREATE OR REPLACE FUNCTION notify_transaction()
-      RETURNS TRIGGER AS $$
-      BEGIN
-        PERFORM pg_notify('transaction_commit', TG_TABLE_NAME);
-        RETURN NEW;
-      END;
-      $$ LANGUAGE plpgsql;
-
-      CREATE TRIGGER todo_transaction_trigger
-      AFTER INSERT ON todo
-      FOR EACH STATEMENT
-      EXECUTE FUNCTION notify_transaction();
-    `)
-
-    // Track transaction commits
-    const transactionCommits: string[] = []
-    const unsubscribe = await pg.listen('transaction_commit', (payload) => {
-      transactionCommits.push(payload)
-    })
-
-    const shape = await pg.electric.syncShapeToTable({
-      shape: {
-        url: 'http://localhost:3000/v1/shape',
-        params: { table: 'todo' },
-      },
-      table: 'todo',
-      primaryKey: ['id'],
-      commitGranularity: 'transaction',
-    })
-
-    // Send messages with different LSNs (first part of offset before _)
-    await feedMessages([
-      {
-        headers: { operation: 'insert' },
-        offset: '1_1', // Transaction 1
-        key: 'id1',
-        value: {
-          id: 1,
-          task: 'task1',
-          done: false,
-        },
-      },
-      {
-        headers: { operation: 'insert' },
-        offset: '1_2', // Same transaction
-        key: 'id2',
-        value: {
-          id: 2,
-          task: 'task2',
-          done: false,
-        },
-      },
-      {
-        headers: { operation: 'insert' },
-        offset: '2_1', // New transaction
-        key: 'id3',
-        value: {
-          id: 3,
-          task: 'task3',
-          done: false,
-        },
-      },
-    ])
-
-    // Wait for all inserts to complete
-    await vi.waitUntil(async () => {
-      const result = await pg.sql<{ count: number }>`
-        SELECT COUNT(*) as count FROM todo;
-      `
-      return result.rows[0].count === 3
-    })
-
-    // Verify all rows were inserted
-    const result = await pg.sql`
-      SELECT * FROM todo ORDER BY id;
-    `
-    expect(result.rows).toEqual([
-      { id: 1, task: 'task1', done: false },
-      { id: 2, task: 'task2', done: false },
-      { id: 3, task: 'task3', done: false },
-    ])
-
-    // Should have received 2 transaction notifications
-    // One for LSN 1 (containing 2 inserts) and one for LSN 2 (containing 1 insert)
-    expect(transactionCommits).toHaveLength(2)
-    expect(transactionCommits).toEqual(['todo', 'todo'])
-
-    await unsubscribe()
-    shape.unsubscribe()
-  })
+  // Removed until Electric has stabilised on LSN metadata
+  // it('respects transaction commit granularity', async () => {
+  //   let feedMessages: (messages: Message[]) => Promise<void> = async (_) => {}
+  //   MockShapeStream.mockImplementation(() => ({
+  //     subscribe: vi.fn((cb: (messages: Message[]) => Promise<void>) => {
+  //       feedMessages = (messages) => cb([...messages, upToDateMsg])
+  //     }),
+  //     unsubscribeAll: vi.fn(),
+  //   }))
+
+  //   // Create a trigger to notify on transaction commit
+  //   await pg.exec(`
+  //     CREATE OR REPLACE FUNCTION notify_transaction()
+  //     RETURNS TRIGGER AS $$
+  //     BEGIN
+  //       PERFORM pg_notify('transaction_commit', TG_TABLE_NAME);
+  //       RETURN NEW;
+  //     END;
+  //     $$ LANGUAGE plpgsql;
+
+  //     CREATE TRIGGER todo_transaction_trigger
+  //     AFTER INSERT ON todo
+  //     FOR EACH STATEMENT
+  //     EXECUTE FUNCTION notify_transaction();
+  //   `)
+
+  //   // Track transaction commits
+  //   const transactionCommits: string[] = []
+  //   const unsubscribe = await pg.listen('transaction_commit', (payload) => {
+  //     transactionCommits.push(payload)
+  //   })
+
+  //   const shape = await pg.electric.syncShapeToTable({
+  //     shape: {
+  //       url: 'http://localhost:3000/v1/shape',
+  //       params: { table: 'todo' },
+  //     },
+  //     table: 'todo',
+  //     primaryKey: ['id'],
+  //     commitGranularity: 'transaction',
+  //   })
+
+  //   // Send messages with different LSNs (first part of offset before _)
+  //   await feedMessages([
+  //     {
+  //       headers: { operation: 'insert' },
+  //       offset: '1_1', // Transaction 1
+  //       key: 'id1',
+  //       value: {
+  //         id: 1,
+  //         task: 'task1',
+  //         done: false,
+  //       },
+  //     },
+  //     {
+  //       headers: { operation: 'insert' },
+  //       offset: '1_2', // Same transaction
+  //       key: 'id2',
+  //       value: {
+  //         id: 2,
+  //         task: 'task2',
+  //         done: false,
+  //       },
+  //     },
+  //     {
+  //       headers: { operation: 'insert' },
+  //       offset: '2_1', // New transaction
+  //       key: 'id3',
+  //       value: {
+  //         id: 3,
+  //         task: 'task3',
+  //         done: false,
+  //       },
+  //     },
+  //   ])
+
+  //   // Wait for all inserts to complete
+  //   await vi.waitUntil(async () => {
+  //     const result = await pg.sql<{ count: number }>`
+  //       SELECT COUNT(*) as count FROM todo;
+  //     `
+  //     return result.rows[0].count === 3
+  //   })
+
+  //   // Verify all rows were inserted
+  //   const result = await pg.sql`
+  //     SELECT * FROM todo ORDER BY id;
+  //   `
+  //   expect(result.rows).toEqual([
+  //     { id: 1, task: 'task1', done: false },
+  //     { id: 2, task: 'task2', done: false },
+  //     { id: 3, task: 'task3', done: false },
+  //   ])
+
+  //   // Should have received 2 transaction notifications
+  //   // One for LSN 1 (containing 2 inserts) and one for LSN 2 (containing 1 insert)
+  //   expect(transactionCommits).toHaveLength(2)
+  //   expect(transactionCommits).toEqual(['todo', 'todo'])
+
+  //   await unsubscribe()
+  //   shape.unsubscribe()
+  // })
 
   it('respects up-to-date commit granularity settings', async () => {
     let feedMessages: (messages: Message[]) => Promise<void> = async (_) => {}

From 7417c2196c1fe4fd06945fa666fcd41df3d1791f Mon Sep 17 00:00:00 2001
From: Sam Willis <sam.willis@gmail.com>
Date: Mon, 9 Dec 2024 09:23:16 +0000
Subject: [PATCH 12/12] Post review changes

---
 packages/pglite-sync/src/index.ts | 15 +++++++++++----
 1 file changed, 11 insertions(+), 4 deletions(-)

diff --git a/packages/pglite-sync/src/index.ts b/packages/pglite-sync/src/index.ts
index b2a8607b2..f3f8efe85 100644
--- a/packages/pglite-sync/src/index.ts
+++ b/packages/pglite-sync/src/index.ts
@@ -161,6 +161,7 @@ async function createPlugin(
 
       const commit = async () => {
         if (messageAggregator.length === 0 && !truncateNeeded) return
+        const shapeHandle = stream.shapeHandle // The shape handle could change while we are committing
         await pg.transaction(async (tx) => {
           if (debug) {
             console.log('committing message batch', messageAggregator.length)
@@ -240,13 +241,13 @@ async function createPlugin(
           if (
             options.shapeKey &&
             messageAggregator.length > 0 &&
-            stream.shapeHandle !== undefined
+            shapeHandle !== undefined
           ) {
             await updateShapeSubscriptionState({
               pg: tx,
               metadataSchema,
               shapeKey: options.shapeKey,
-              shapeId: stream.shapeHandle,
+              shapeId: shapeHandle,
               lastOffset:
                 messageAggregator[messageAggregator.length - 1].offset,
             })
@@ -258,8 +259,14 @@ async function createPlugin(
         await new Promise((resolve) => setTimeout(resolve, 0))
       }
 
-      const throttledCommit = async () => {
+      const throttledCommit = async ({
+        reset = false,
+      }: { reset?: boolean } = {}) => {
         const now = Date.now()
+        if (reset) {
+          // Reset the last commit time to 0, forcing the next commit to happen immediately
+          lastCommitAt = 0
+        }
         if (options.commitThrottle && debug)
           console.log(
             'throttled commit: now:',
@@ -322,7 +329,7 @@ async function createPlugin(
 
               case 'up-to-date':
                 // perform all accumulated changes and store stream state
-                await commit() // not throttled, we want this to happen ASAP
+                await throttledCommit({ reset: true }) // not throttled, we want this to happen ASAP
                 if (
                   isNewSubscription &&
                   !onInitialSyncCalled &&