orez 0.0.18 → 0.0.19

package/src/integration/integration.test.ts

@@ -9,13 +9,18 @@
 
 import { describe, expect, test, beforeAll, afterAll, beforeEach } from 'vitest'
 import WebSocket from 'ws'
+
 import { startZeroLite } from '../index.js'
+
 import type { PGlite } from '@electric-sql/pglite'
 
 // simple async queue for collecting websocket messages
 class Queue<T> {
   private items: T[] = []
-  private waiters: Array<{ resolve: (v: T) => void; timer?: ReturnType<typeof setTimeout> }> = []
+  private waiters: Array<{
+    resolve: (v: T) => void
+    timer?: ReturnType<typeof setTimeout>
+  }> = []
 
   enqueue(item: T) {
     const waiter = this.waiters.shift()
@@ -113,7 +118,7 @@ describe('orez integration', { timeout: 120000 }, () => {
   test('zero-cache starts and accepts websocket connections', async () => {
     const ws = new WebSocket(
       `ws://localhost:${zeroPort}/sync/v4/connect` +
-        `?clientGroupID=test-cg&clientID=test-client&wsid=ws1&schemaVersion=1&baseCookie=&ts=${Date.now()}&lmid=0`,
+        `?clientGroupID=test-cg&clientID=test-client&wsid=ws1&schemaVersion=1&baseCookie=&ts=${Date.now()}&lmid=0`
     )
 
     const connected = new Promise<void>((resolve, reject) => {
@@ -161,7 +166,7 @@ describe('orez integration', { timeout: 120000 }, () => {
             value: 'hello',
           }),
         }),
-      ]),
+      ])
     )
 
     ws.close()
@@ -196,7 +201,7 @@ describe('orez integration', { timeout: 120000 }, () => {
             value: 'live-value',
           }),
         }),
-      ]),
+      ])
     )
 
     ws.close()
@@ -236,7 +241,7 @@ describe('orez integration', { timeout: 120000 }, () => {
             value: 'updated',
           }),
         }),
-      ]),
+      ])
     )
 
     ws.close()
@@ -267,7 +272,7 @@ describe('orez integration', { timeout: 120000 }, () => {
           op: 'del',
           tableName: 'foo',
         }),
-      ]),
+      ])
     )
 
     ws.close()
@@ -289,8 +294,8 @@ describe('orez integration', { timeout: 120000 }, () => {
           `concurrent-${i}`,
           `value-${i}`,
           i,
-        ]),
-      ),
+        ])
+      )
     )
 
     // collect all poke parts within a window
@@ -316,11 +321,11 @@ describe('orez integration', { timeout: 120000 }, () => {
   function connectAndSubscribe(
     port: number,
     downstream: Queue<unknown>,
-    query: Record<string, unknown>,
+    query: Record<string, unknown>
   ): WebSocket {
     const ws = new WebSocket(
       `ws://localhost:${port}/sync/v4/connect` +
-        `?clientGroupID=test-cg-${Date.now()}&clientID=test-client&wsid=ws1&schemaVersion=1&baseCookie=&ts=${Date.now()}&lmid=0`,
+        `?clientGroupID=test-cg-${Date.now()}&clientID=test-client&wsid=ws1&schemaVersion=1&baseCookie=&ts=${Date.now()}&lmid=0`
     )
 
     ws.on('message', (data) => {
@@ -334,7 +339,7 @@ describe('orez integration', { timeout: 120000 }, () => {
           {
             desiredQueriesPatch: [{ op: 'put', hash: 'q1', ast: query }],
           },
-        ]),
+        ])
       )
     })
 
@@ -363,7 +368,7 @@ describe('orez integration', { timeout: 120000 }, () => {
 
   async function waitForPokePart(
     downstream: Queue<unknown>,
-    timeoutMs = 10000,
+    timeoutMs = 10000
   ): Promise<Record<string, any>> {
     const deadline = Date.now() + timeoutMs
     while (Date.now() < deadline) {
@@ -379,7 +384,7 @@ describe('orez integration', { timeout: 120000 }, () => {
 
   async function collectPokeRows(
     downstream: Queue<unknown>,
-    windowMs = 5000,
+    windowMs = 5000
   ): Promise<any[]> {
     const rows: any[] = []
     const deadline = Date.now() + windowMs
@@ -392,7 +397,12 @@ describe('orez integration', { timeout: 120000 }, () => {
         rows.push(...msg[1].rowsPatch)
         // check if more poke parts come quickly
         const more = (await downstream.dequeue('timeout' as any, 2000)) as any
-        if (more !== 'timeout' && Array.isArray(more) && more[0] === 'pokePart' && more[1]?.rowsPatch) {
+        if (
+          more !== 'timeout' &&
+          Array.isArray(more) &&
+          more[0] === 'pokePart' &&
+          more[1]?.rowsPatch
+        ) {
           rows.push(...more[1].rowsPatch)
         }
         break

package/src/mutex.ts

@@ -0,0 +1,27 @@
+// simple mutex for serializing pglite access
+export class Mutex {
+  private locked = false
+  private queue: Array<() => void> = []
+
+  async acquire(): Promise<void> {
+    if (!this.locked) {
+      this.locked = true
+      return
+    }
+    return new Promise<void>((resolve) => {
+      this.queue.push(resolve)
+    })
+  }
+
+  release(): void {
+    const next = this.queue.shift()
+    if (next) {
+      next()
+    } else {
+      this.locked = false
+    }
+  }
+}
+
+// shared mutex instance for serializing pglite access across proxy and replication handler
+export const pgMutex = new Mutex()

package/src/pg-proxy.ts

@@ -32,26 +32,33 @@ const QUERY_REWRITES: Array<{ match: RegExp; replace: string }> = [
     match: /current_setting\s*\(\s*'wal_level'\s*\)/gi,
     replace: "'logical'::text",
   },
-  // strip READ ONLY from BEGIN
+  // strip READ ONLY from BEGIN (pglite is single-session, no read-only transactions)
   {
     match: /\bREAD\s+ONLY\b/gi,
     replace: '',
   },
+  // strip ISOLATION LEVEL from any query (pglite is single-session, isolation is meaningless)
+  // catches: SET TRANSACTION ISOLATION LEVEL SERIALIZABLE, BEGIN ISOLATION LEVEL SERIALIZABLE, etc.
+  {
+    match:
+      /\bISOLATION\s+LEVEL\s+(SERIALIZABLE|REPEATABLE\s+READ|READ\s+COMMITTED|READ\s+UNCOMMITTED)\b/gi,
+    replace: '',
+  },
+  // strip bare SET TRANSACTION (after ISOLATION LEVEL is removed, this becomes a no-op statement)
+  {
+    match: /\bSET\s+TRANSACTION\s*;/gi,
+    replace: ';',
+  },
   // redirect pg_replication_slots to our fake table
   {
     match: /\bpg_replication_slots\b/g,
     replace: 'public._zero_replication_slots',
   },
-  // SET TRANSACTION queries: rewrite to SELECT 1 (PGlite doesn't support them,
-  // and no-op interception doesn't work for extended protocol Bind/Execute)
-  {
-    match: /^\s*SET\s+TRANSACTION\b.*$/gis,
-    replace: 'SELECT 1',
-  },
 ]
 
-// queries to intercept and return no-op success
-const NOOP_QUERY_PATTERNS: RegExp[] = []
+// queries to intercept and return no-op success (synthetic SET response)
+// pglite rejects SET TRANSACTION if any query (e.g. SET search_path) ran first
+const NOOP_QUERY_PATTERNS: RegExp[] = [/^\s*SET\s+TRANSACTION\b/i, /^\s*SET\s+SESSION\b/i]
 
 /**
  * extract query text from a Parse message (0x50).
@@ -239,32 +246,9 @@ function stripReadyForQuery(data: Uint8Array): Uint8Array {
   return result
 }
 
-// simple mutex for serializing pglite access
-class Mutex {
-  private locked = false
-  private queue: Array<() => void> = []
-
-  async acquire(): Promise<void> {
-    if (!this.locked) {
-      this.locked = true
-      return
-    }
-    return new Promise<void>((resolve) => {
-      this.queue.push(resolve)
-    })
-  }
-
-  release(): void {
-    const next = this.queue.shift()
-    if (next) {
-      next()
-    } else {
-      this.locked = false
-    }
-  }
-}
+import { pgMutex } from './mutex.js'
 
-const mutex = new Mutex()
+const mutex = pgMutex
 
 // module-level search_path tracking
 let currentSearchPath = 'public'
@@ -434,6 +418,9 @@ async function handleReplicationMessage(
   const response = await handleReplicationQuery(query, db)
   if (response) return response
 
+  // apply query rewrites (e.g. SET TRANSACTION → SELECT 1) before forwarding
+  data = interceptQuery(data)
+
   // fall through to pglite for unrecognized queries
   await mutex.acquire()
   try {

package/src/replication/change-tracker.test.ts

@@ -33,7 +33,7 @@ describe('change-tracker', () => {
     const changes = await getChangesSince(db, 0)
     expect(changes).toHaveLength(1)
     expect(changes[0].op).toBe('INSERT')
-    expect(changes[0].table_name).toBe('items')
+    expect(changes[0].table_name).toBe('public.items')
     expect(changes[0].row_data).toMatchObject({ name: 'a', value: 1 })
     expect(changes[0].old_data).toBeNull()
   })
@@ -116,8 +116,8 @@ describe('change-tracker', () => {
 
     const changes = await getChangesSince(db, 0)
     const tables = new Set(changes.map((c) => c.table_name))
-    expect(tables).toContain('items')
-    expect(tables).toContain('other')
+    expect(tables).toContain('public.items')
+    expect(tables).toContain('public.other')
   })
 
   it('handles rapid inserts (50 rows)', async () => {
@@ -178,7 +178,7 @@ describe('change-tracker', () => {
     await db.exec(`INSERT INTO public."my""table" (val) VALUES ('works')`)
 
     const changes = await getChangesSince(db, 0)
-    const special = changes.filter((c) => c.table_name === 'my"table')
+    const special = changes.filter((c) => c.table_name === 'public.my"table')
     expect(special).toHaveLength(1)
     expect(special[0].op).toBe('INSERT')
     expect(special[0].row_data).toMatchObject({ val: 'works' })

package/src/replication/change-tracker.ts

@@ -45,18 +45,21 @@ export async function installChangeTracking(db: PGlite): Promise<void> {
   // create trigger function
   await db.exec(`
     CREATE OR REPLACE FUNCTION public._zero_track_change() RETURNS TRIGGER AS $$
+    DECLARE
+      qualified_name TEXT;
     BEGIN
+      qualified_name := TG_TABLE_SCHEMA || '.' || TG_TABLE_NAME;
       IF TG_OP = 'DELETE' THEN
         INSERT INTO public._zero_changes (table_name, op, old_data)
-        VALUES (TG_TABLE_NAME, 'DELETE', row_to_json(OLD)::jsonb);
+        VALUES (qualified_name, 'DELETE', row_to_json(OLD)::jsonb);
         RETURN OLD;
       ELSIF TG_OP = 'UPDATE' THEN
         INSERT INTO public._zero_changes (table_name, op, row_data, old_data)
-        VALUES (TG_TABLE_NAME, 'UPDATE', row_to_json(NEW)::jsonb, row_to_json(OLD)::jsonb);
+        VALUES (qualified_name, 'UPDATE', row_to_json(NEW)::jsonb, row_to_json(OLD)::jsonb);
         RETURN NEW;
       ELSIF TG_OP = 'INSERT' THEN
         INSERT INTO public._zero_changes (table_name, op, row_data)
-        VALUES (TG_TABLE_NAME, 'INSERT', row_to_json(NEW)::jsonb);
+        VALUES (qualified_name, 'INSERT', row_to_json(NEW)::jsonb);
         RETURN NEW;
       END IF;
       RETURN NULL;
@@ -137,6 +140,49 @@ async function installTriggersOnAllTables(db: PGlite): Promise<void> {
   log.debug.pglite(`installed change tracking triggers on ${count} tables`)
 }
 
+/**
+ * install change tracking triggers on tables in shard schemas.
+ * zero-cache creates shard schemas (e.g. chat_0) with clients/mutations
+ * tables that track mutation confirmations. these must be replicated
+ * for .server promises to resolve.
+ */
+export async function installTriggersOnShardTables(db: PGlite): Promise<void> {
+  const result = await db.query<{ nspname: string }>(
+    `SELECT nspname FROM pg_namespace
+     WHERE nspname NOT IN ('pg_catalog', 'information_schema', 'pg_toast', 'public')
+       AND nspname NOT LIKE 'pg_%'
+       AND nspname NOT LIKE 'zero_%'
+       AND nspname NOT LIKE '_zero_%'
+       AND nspname NOT LIKE '%/%'`
+  )
+
+  if (result.rows.length === 0) return
+
+  let count = 0
+  for (const { nspname } of result.rows) {
+    const tables = await db.query<{ tablename: string }>(
+      `SELECT tablename FROM pg_tables WHERE schemaname = $1`,
+      [nspname]
+    )
+
+    for (const { tablename } of tables.rows) {
+      const quotedSchema = quoteIdent(nspname)
+      const quotedTable = quoteIdent(tablename)
+      await db.exec(`
+        DROP TRIGGER IF EXISTS _zero_change_trigger ON ${quotedSchema}.${quotedTable};
+        CREATE TRIGGER _zero_change_trigger
+          AFTER INSERT OR UPDATE OR DELETE ON ${quotedSchema}.${quotedTable}
+          FOR EACH ROW EXECUTE FUNCTION public._zero_track_change();
+      `)
+      count++
+    }
+  }
+
+  if (count > 0) {
+    log.debug.pglite(`installed change tracking on ${count} shard tables`)
+  }
+}
+
 export async function getChangesSince(
   db: PGlite,
   watermark: number,