orez 0.2.17 → 0.2.19

package/src/replication/handler.ts

@@ -134,6 +134,7 @@ export function resetReplicationState(): void {
   cachedTableKeyColumns = null
   cachedExcludedColumns = null
   cachedColumnTypeOids = null
+  cachedColumns.clear()
 }
 function nextLsn(): bigint {
   const floor = lsnFloorFromTime()
@@ -539,8 +540,8 @@ export async function handleStartReplication(
       if (shardClientSchemas.length > 0) {
         const shardTables = await db.query<{ schemaname: string; tablename: string }>(
           `SELECT schemaname, tablename FROM pg_tables
-           WHERE schemaname = ANY($1) AND tablename = 'clients'`,
-          [shardClientSchemas]
+           WHERE schemaname = ANY($1) AND tablename = ANY($2)`,
+          [shardClientSchemas, ['clients', 'mutations']]
         )
         for (const { schemaname, tablename } of shardTables.rows) {
           const qs = '"' + schemaname.replace(/"/g, '""') + '"'
@@ -578,8 +579,9 @@ export async function handleStartReplication(
         table_name: string
         column_name: string
         data_type: string | null
+        ordinal_position: number
       }>(
-        `SELECT 'pk' AS kind, tc.table_schema, tc.table_name, kcu.column_name, NULL AS data_type
+        `SELECT 'pk' AS kind, tc.table_schema, tc.table_name, kcu.column_name, NULL AS data_type, kcu.ordinal_position
          FROM information_schema.table_constraints tc
          JOIN information_schema.key_column_usage kcu
            ON tc.constraint_name = kcu.constraint_name
@@ -587,9 +589,10 @@ export async function handleStartReplication(
          WHERE tc.constraint_type = 'PRIMARY KEY'
            AND tc.table_schema = ANY($1)
          UNION ALL
-         SELECT 'col' AS kind, table_schema, table_name, column_name, data_type
+         SELECT 'col' AS kind, table_schema, table_name, column_name, data_type, ordinal_position
          FROM information_schema.columns
-         WHERE table_schema = ANY($1)`,
+         WHERE table_schema = ANY($1)
+         ORDER BY table_schema, table_name, kind, ordinal_position`,
         [relevantSchemas]
       )
 
@@ -834,9 +837,9 @@ export async function handleStartReplication(
             `found ${changes.length} changes [${tableSummary}] (wm ${lastWatermark}→${changes[changes.length - 1].watermark}) query=${queryMs.toFixed(1)}ms signal→query=${signalToQueryMs}ms`
           )
           // filter out shard tables that zero-cache doesn't expect.
-          // only `clients` is needed (for .server promise resolution).
-          // other shard tables (replicas, mutations) crash zero-cache
-          // with "Unknown table" in change-processor.
+          // `clients` advances lmid; `mutations` carries mutation results.
+          // other shard tables (e.g. replicas) crash zero-cache with
+          // "Unknown table" in the change processor.
           const batchEnd = changes[changes.length - 1].watermark
           const preFilterCount = changes.length
           changes = changes.filter((c) => {
@@ -845,7 +848,7 @@ export async function handleStartReplication(
             const schema = c.table_name.substring(0, dot)
             if (schema === 'public') return true
             const table = c.table_name.substring(dot + 1)
-            return table === 'clients'
+            return table === 'clients' || table === 'mutations'
           })
           log.debug.repl(`filter: ${preFilterCount} → ${changes.length} changes`)
 
@@ -859,6 +862,15 @@ export async function handleStartReplication(
             continue
           }
 
+          await ensureMetadataForChangedTables(
+            db,
+            mutex,
+            changes,
+            tableKeyColumns,
+            excludedColumns,
+            columnTypeOids
+          )
+
           log.debug.repl(`streaming ${changes.length} changes to writer`)
           await streamChanges(
             changes,
@@ -930,6 +942,106 @@ export async function handleStartReplication(
   log.repl('poll loop exited')
 }
 
+async function ensureMetadataForChangedTables(
+  db: PGlite,
+  mutex: Mutex,
+  changes: ChangeRecord[],
+  tableKeyColumns: Map<string, Set<string>>,
+  excludedColumns: Map<string, Set<string>>,
+  columnTypeOids: Map<string, Map<string, number>>
+): Promise<void> {
+  const missing = new Map<string, { schema: string; table: string }>()
+
+  for (const change of changes) {
+    if (
+      tableKeyColumns.has(change.table_name) ||
+      columnTypeOids.has(change.table_name) ||
+      excludedColumns.has(change.table_name)
+    ) {
+      continue
+    }
+
+    const dot = change.table_name.indexOf('.')
+    const schema = dot === -1 ? 'public' : change.table_name.substring(0, dot)
+    const table = dot === -1 ? change.table_name : change.table_name.substring(dot + 1)
+    missing.set(`${schema}.${table}`, { schema, table })
+  }
+
+  if (missing.size === 0) return
+
+  const schemas = [...new Set([...missing.values()].map((entry) => entry.schema))]
+  const tables = [...new Set([...missing.values()].map((entry) => entry.table))]
+
+  await mutex.acquire()
+  try {
+    const schemaResult = await db.query<{
+      kind: string
+      table_schema: string
+      table_name: string
+      column_name: string
+      data_type: string | null
+      ordinal_position: number
+    }>(
+      `SELECT 'pk' AS kind, tc.table_schema, tc.table_name, kcu.column_name, NULL AS data_type, kcu.ordinal_position
+       FROM information_schema.table_constraints tc
+       JOIN information_schema.key_column_usage kcu
+         ON tc.constraint_name = kcu.constraint_name
+         AND tc.table_schema = kcu.table_schema
+       WHERE tc.constraint_type = 'PRIMARY KEY'
+         AND tc.table_schema = ANY($1)
+         AND tc.table_name = ANY($2)
+       UNION ALL
+       SELECT 'col' AS kind, table_schema, table_name, column_name, data_type, ordinal_position
+       FROM information_schema.columns
+       WHERE table_schema = ANY($1)
+         AND table_name = ANY($2)
+       ORDER BY table_schema, table_name, kind, ordinal_position`,
+      [schemas, tables]
+    )
+
+    for (const row of schemaResult.rows) {
+      const key = `${row.table_schema}.${row.table_name}`
+      if (!missing.has(key)) continue
+
+      if (row.kind === 'pk') {
+        let keys = tableKeyColumns.get(key)
+        if (!keys) {
+          keys = new Set()
+          tableKeyColumns.set(key, keys)
+        }
+        keys.add(row.column_name)
+      } else {
+        if (row.data_type && UNSUPPORTED_TYPES.has(row.data_type)) {
+          let cols = excludedColumns.get(key)
+          if (!cols) {
+            cols = new Set()
+            excludedColumns.set(key, cols)
+          }
+          cols.add(row.column_name)
+        }
+        if (row.data_type) {
+          const oid = PG_DATA_TYPE_OIDS[row.data_type]
+          if (oid !== undefined) {
+            let cols = columnTypeOids.get(key)
+            if (!cols) {
+              cols = new Map()
+              columnTypeOids.set(key, cols)
+            }
+            cols.set(row.column_name, oid)
+          }
+        }
+      }
+    }
+
+    log.debug.repl(
+      `refreshed metadata for ${missing.size} late table(s): ${[...missing.keys()].join(',')}`
+    )
+    for (const key of missing.keys()) cachedColumns.delete(key)
+  } finally {
+    mutex.release()
+  }
+}
+
 // cache column info per table to avoid per-change allocation
 const cachedColumns = new Map<string, ReturnType<typeof inferColumns>>()
 

package/src/replication/zero-compat.test.ts

@@ -17,7 +17,11 @@ import { describe, it, expect, beforeEach, afterEach } from 'vitest'
 
 import { getConfig } from '../config'
 import { startPgProxy } from '../pg-proxy'
-import { installChangeTracking, resetShardSchemaCache } from './change-tracker'
+import {
+  installChangeTracking,
+  installTriggersOnShardTables,
+  resetShardSchemaCache,
+} from './change-tracker'
 import { signalReplicationChange, resetReplicationState } from './handler'
 
 import type { Server, AddressInfo } from 'node:net'
@@ -959,10 +963,10 @@ describe('zero-cache pgoutput compatibility', { timeout: 30000 }, () => {
     s.close()
   })
 
-  it('shard replicas/mutations changes NOT streamed (only clients)', async () => {
+  it('streams shard mutation-confirmation tables but filters replicas', async () => {
     // zero-cache creates shard schemas (chat_0) with clients, replicas, mutations.
-    // if we stream replicas/mutations changes, zero-cache crashes with
-    // "Unknown table chat_0.replicas". only clients changes should be streamed.
+    // clients advance lmid and mutations carry server results; replicas is
+    // internal shard state that zero-cache does not expect in the change stream.
     await db.exec(`
       CREATE SCHEMA chat_0;
       CREATE TABLE chat_0.clients (
@@ -976,9 +980,11 @@ describe('zero-cache pgoutput compatibility', { timeout: 30000 }, () => {
         version TEXT
       );
       CREATE TABLE chat_0.mutations (
-        id TEXT PRIMARY KEY,
-        "clientID" TEXT,
-        name TEXT
+        "clientGroupID" TEXT NOT NULL,
+        "clientID" TEXT NOT NULL,
+        "mutationID" BIGINT NOT NULL,
+        result JSON,
+        PRIMARY KEY ("clientGroupID", "clientID", "mutationID")
       );
     `)
     await installChangeTracking(db)
@@ -995,7 +1001,7 @@ describe('zero-cache pgoutput compatibility', { timeout: 30000 }, () => {
     )
     await db.exec(`INSERT INTO chat_0.replicas (id, version) VALUES ('r1', 'v1')`)
     await db.exec(
-      `INSERT INTO chat_0.mutations (id, "clientID", name) VALUES ('m1', 'c1', 'send')`
+      `INSERT INTO chat_0.mutations ("clientGroupID", "clientID", "mutationID", result) VALUES ('cg1', 'c1', 1, '{}')`
     )
     await db.exec(`INSERT INTO public.foo (id) VALUES ('normal')`)
 
@@ -1008,12 +1014,88 @@ describe('zero-cache pgoutput compatibility', { timeout: 30000 }, () => {
       if (m.tag === 'insert') inserts.push(m as ZcInsert)
     }
 
-    // should see clients + foo inserts, but NOT replicas or mutations
+    // should see clients + mutations + foo inserts, but NOT replicas.
     const streamedTables = inserts.map((i) => `${i.relation.schema}.${i.relation.name}`)
     expect(streamedTables).toContain('public.foo')
     expect(streamedTables).toContain('chat_0.clients')
+    expect(streamedTables).toContain('chat_0.mutations')
     expect(streamedTables).not.toContain('chat_0.replicas')
-    expect(streamedTables).not.toContain('chat_0.mutations')
+
+    s.close()
+  })
+
+  it('refreshes metadata for shard tables missing from cached setup', async () => {
+    // zero-cache can create shard schemas after the first replication stream
+    // cached metadata. the next stream must not encode bigint/json shard
+    // columns as text, or zero-cache cannot parse mutation confirmations.
+    const warmup = await stream()
+    warmup.close()
+    await new Promise((r) => setTimeout(r, 50))
+
+    await db.exec(`
+      CREATE SCHEMA chat_0;
+      CREATE TABLE chat_0.clients (
+        "clientGroupID" TEXT NOT NULL,
+        "clientID" TEXT NOT NULL,
+        "lastMutationID" BIGINT,
+        PRIMARY KEY ("clientGroupID", "clientID")
+      );
+      CREATE TABLE chat_0.mutations (
+        "clientGroupID" TEXT NOT NULL,
+        "clientID" TEXT NOT NULL,
+        "mutationID" BIGINT NOT NULL,
+        result JSON,
+        PRIMARY KEY ("clientGroupID", "clientID", "mutationID")
+      );
+    `)
+    await installTriggersOnShardTables(db)
+
+    const s = await stream()
+    const q = s.messages
+
+    await db.exec(
+      `INSERT INTO chat_0.clients ("clientGroupID", "clientID", "lastMutationID") VALUES ('cg1', 'c1', 42)`
+    )
+    await db.exec(
+      `INSERT INTO chat_0.mutations ("clientGroupID", "clientID", "mutationID", result) VALUES ('cg1', 'c1', 42, '{"data":"ok"}')`
+    )
+
+    let clientsInsert: ZcInsert | null = null
+    let mutationsInsert: ZcInsert | null = null
+    const deadline = Date.now() + 5000
+    while (Date.now() < deadline) {
+      const m = await q.dequeue(2000).catch(() => null)
+      if (!m) break
+      if (m.tag === 'insert') {
+        const ins = m as ZcInsert
+        if (ins.relation.schema === 'chat_0' && ins.relation.name === 'clients') {
+          clientsInsert = ins
+        }
+        if (ins.relation.schema === 'chat_0' && ins.relation.name === 'mutations') {
+          mutationsInsert = ins
+        }
+      }
+      if (clientsInsert && mutationsInsert) break
+    }
+
+    expect(clientsInsert).not.toBeNull()
+    expect(mutationsInsert).not.toBeNull()
+    expect([...clientsInsert!.relation.keyColumns].sort()).toEqual(
+      ['clientGroupID', 'clientID'].sort()
+    )
+    expect(
+      clientsInsert!.relation.columns.find((col) => col.name === 'lastMutationID')
+        ?.typeOid
+    ).toBe(20)
+    expect([...mutationsInsert!.relation.keyColumns].sort()).toEqual(
+      ['clientGroupID', 'clientID', 'mutationID'].sort()
+    )
+    expect(
+      mutationsInsert!.relation.columns.find((col) => col.name === 'mutationID')?.typeOid
+    ).toBe(20)
+    expect(
+      mutationsInsert!.relation.columns.find((col) => col.name === 'result')?.typeOid
+    ).toBe(114)
 
     s.close()
   })

package/src/worker/shims/ws-browser.test.ts

@@ -66,6 +66,27 @@ describe('messagePortToWs', () => {
     expect(spy).toHaveBeenCalledWith('hello')
   })
 
+  it('send() invokes the ws callback after posting', () => {
+    const [port1] = createMockPorts()
+    const ws = messagePortToWs(port1 as any)
+    const callback = vi.fn()
+
+    ws.send('hello', callback)
+
+    expect(callback).toHaveBeenCalledWith()
+  })
+
+  it('send() reports an error to the callback after close', () => {
+    const [port1] = createMockPorts()
+    const ws = messagePortToWs(port1 as any)
+    const callback = vi.fn()
+
+    ws.close()
+    ws.send('ignored', callback)
+
+    expect(callback).toHaveBeenCalledWith(expect.any(Error))
+  })
+
   it('forwards port messages as ws message events', () => {
     const [port1, port2] = createMockPorts()
     const ws1 = messagePortToWs(port1 as any)

package/src/worker/shims/ws-browser.ts

@@ -26,7 +26,10 @@ interface WsCompatible {
   CLOSING: number
   CLOSED: number
   readyState: number
-  send(data: string | ArrayBuffer | ArrayBufferView): void
+  send(
+    data: string | ArrayBuffer | ArrayBufferView,
+    callback?: (err?: Error) => void
+  ): void
   close(code?: number, reason?: string): void
   addEventListener(type: string, handler: (event: any) => void): void
   removeEventListener(type: string, handler: (event: any) => void): void
@@ -143,9 +146,24 @@ export function messagePortToWs(port: MessagePort): WsCompatible {
       return closed ? 3 : 1 // CLOSED or OPEN
     },
 
-    send(data: string | ArrayBuffer | ArrayBufferView) {
-      if (closed) return
-      port.postMessage(data)
+    send(data: string | ArrayBuffer | ArrayBufferView, callback?: (err?: Error) => void) {
+      if (closed) {
+        callback?.(new Error('WebSocket is closed'))
+        return
+      }
+
+      try {
+        port.postMessage(data)
+      } catch (error) {
+        const err = error instanceof Error ? error : new Error(String(error))
+        if (callback) {
+          callback(err)
+          return
+        }
+        throw err
+      }
+
+      callback?.()
     },
 
     close(code?: number, _reason?: string) {
@@ -179,8 +197,19 @@ export function browserWsToWs(ws: WebSocket): WsCompatible {
       return ws.readyState
     },
 
-    send(data: string | ArrayBuffer | ArrayBufferView) {
-      ws.send(data)
+    send(data: string | ArrayBuffer | ArrayBufferView, callback?: (err?: Error) => void) {
+      try {
+        ws.send(data)
+      } catch (error) {
+        const err = error instanceof Error ? error : new Error(String(error))
+        if (callback) {
+          callback(err)
+          return
+        }
+        throw err
+      }
+
+      callback?.()
     },
 
     close(code?: number, reason?: string) {