@dotdo/postgres 0.1.0

package/src/worker/wal.ts

@@ -0,0 +1,697 @@
+/**
+ * WAL (Write-Ahead Log) Capture for PostgresDO
+ * Task: postgres-ctf - WAL capture mechanism in PostgresDO
+ *
+ * Implements Write-Ahead Log capture for streaming to Iceberg for time travel.
+ * Each PostgresDO instance captures its mutations as WAL entries that can be
+ * streamed to R2/Iceberg for point-in-time recovery and data lineage.
+ */
+import type { CDCEventType } from './types'
+
+/**
+ * WAL operation types (same as CDC event types)
+ */
+export type WALOperation = CDCEventType
+
+/**
+ * WAL capture mode
+ * - 'full': Capture full row data for each operation
+ * - 'delta': Capture only changed columns for UPDATE operations
+ */
+export type WALCaptureMode = 'full' | 'delta'
+
+/**
+ * WAL entry representing a single database mutation
+ */
+export interface WALEntry {
+  /** Log Sequence Number - unique, monotonically increasing identifier */
+  lsn: bigint
+  /** Type of operation */
+  operation: WALOperation
+  /** Schema name (default: 'public') */
+  schema: string
+  /** Table name */
+  table: string
+  /** Primary key values (if available) */
+  primaryKey?: Record<string, unknown> | undefined
+  /** New row data (for INSERT and UPDATE) */
+  newRow?: Record<string, unknown> | undefined
+  /** Old row data (for UPDATE and DELETE) */
+  oldRow?: Record<string, unknown> | undefined
+  /** Changed columns (for delta mode UPDATE) */
+  changedColumns?: string[] | undefined
+  /** Transaction ID (if within a transaction) */
+  transactionId?: string | undefined
+  /** Timestamp when the change occurred */
+  timestamp: number
+  /** SQL that caused this change (for debugging, optional) */
+  sql?: string | undefined
+}
+
+/**
+ * Serialized WAL entry for storage/transmission
+ * Uses string for LSN since bigint doesn't serialize to JSON
+ */
+export interface SerializedWALEntry {
+  lsn: string
+  operation: WALOperation
+  schema: string
+  table: string
+  primaryKey?: Record<string, unknown> | undefined
+  newRow?: Record<string, unknown> | undefined
+  oldRow?: Record<string, unknown> | undefined
+  changedColumns?: string[] | undefined
+  transactionId?: string | undefined
+  timestamp: number
+  sql?: string | undefined
+}
+
+/**
+ * WAL buffer statistics
+ */
+export interface WALBufferStats {
+  /** Number of entries in the buffer */
+  entryCount: number
+  /** Estimated memory usage in bytes */
+  memoryBytes: number
+  /** Oldest LSN in the buffer */
+  oldestLsn: bigint | null
+  /** Newest LSN in the buffer */
+  newestLsn: bigint | null
+  /** Total entries captured since initialization */
+  totalEntriesCaptured: bigint
+  /** Entries flushed to storage */
+  entriesFlushed: bigint
+  /** Entries dropped due to buffer overflow */
+  entriesDropped: bigint
+}
+
+/**
+ * WAL manager configuration
+ */
+export interface WALConfig {
+  /** Enable WAL capture (default: true) */
+  enabled?: boolean
+  /** Capture mode (default: 'full') */
+  captureMode?: WALCaptureMode
+  /** Maximum entries to buffer in memory (default: 1000) */
+  maxBufferEntries?: number
+  /** Maximum buffer memory in bytes (default: 1MB) */
+  maxBufferBytes?: number
+  /** Include original SQL in entries (default: false, for debugging) */
+  includeSQL?: boolean
+  /** Auto-flush to storage when buffer reaches this percentage (default: 0.8) */
+  flushThreshold?: number
+  /** Tables to include (if empty, include all) */
+  includeTables?: string[]
+  /** Tables to exclude */
+  excludeTables?: string[]
+  /** Operations to capture (default: all) */
+  operations?: WALOperation[]
+}
+
+/**
+ * WAL flush result
+ */
+export interface WALFlushResult {
+  /** Number of entries flushed */
+  entriesFlushed: number
+  /** LSN range of flushed entries */
+  lsnRange: { start: bigint; end: bigint } | null
+  /** Duration of flush operation in ms */
+  durationMs: number
+}
+
+/**
+ * WAL storage interface for persisting entries
+ */
+export interface WALStorage {
+  /** Store WAL entries */
+  store(entries: SerializedWALEntry[]): Promise<void>
+  /** Get the last stored LSN */
+  getLastLsn(): Promise<bigint | null>
+  /** Read entries from storage (for replay) */
+  readEntries(fromLsn: bigint, limit?: number): Promise<SerializedWALEntry[]>
+}
+
+/**
+ * Default WAL configuration
+ */
+const DEFAULT_WAL_CONFIG: Required<WALConfig> = {
+  enabled: true,
+  captureMode: 'full',
+  maxBufferEntries: 1000,
+  maxBufferBytes: 1024 * 1024, // 1MB
+  includeSQL: false,
+  flushThreshold: 0.8,
+  includeTables: [],
+  excludeTables: [],
+  operations: ['INSERT', 'UPDATE', 'DELETE'],
+}
+
+/**
+ * Estimate the memory size of a WAL entry in bytes
+ */
+function estimateEntrySize(entry: WALEntry): number {
+  // Rough estimation based on JSON serialization
+  const json = JSON.stringify({
+    ...entry,
+    lsn: entry.lsn.toString(),
+  })
+  // Add ~20% overhead for object structure
+  return Math.ceil(json.length * 1.2)
+}
+
+/**
+ * Serialize a WAL entry for storage/transmission
+ */
+export function serializeWALEntry(entry: WALEntry): SerializedWALEntry {
+  return {
+    ...entry,
+    lsn: entry.lsn.toString(),
+  }
+}
+
+/**
+ * Deserialize a WAL entry from storage/transmission
+ */
+export function deserializeWALEntry(entry: SerializedWALEntry): WALEntry {
+  return {
+    ...entry,
+    lsn: BigInt(entry.lsn),
+  }
+}
+
+/**
+ * Parse SQL to detect write operations and extract table names
+ * Returns null if not a write operation
+ */
+export function parseWriteOperation(sql: string): {
+  operation: WALOperation
+  table: string
+  schema: string
+} | null {
+  // Normalize SQL: trim and collapse whitespace (preserve case for quoted identifiers)
+  const normalized = sql.trim().replace(/\s+/g, ' ')
+
+  // INSERT INTO [schema.]table
+  const insertMatch = normalized.match(/^INSERT\s+INTO\s+(?:([a-z_][a-z0-9_]*)\s*\.\s*)?([a-z_][a-z0-9_]*|"[^"]+")(?:\s|$|\()/i)
+  if (insertMatch) {
+    return {
+      operation: 'INSERT',
+      schema: insertMatch[1]?.toLowerCase() || 'public',
+      table: unquoteIdentifier(insertMatch[2]!),
+    }
+  }
+
+  // UPDATE [schema.]table SET
+  const updateMatch = normalized.match(/^UPDATE\s+(?:([a-z_][a-z0-9_]*)\s*\.\s*)?([a-z_][a-z0-9_]*|"[^"]+")(?:\s+SET|\s+AS)/i)
+  if (updateMatch) {
+    return {
+      operation: 'UPDATE',
+      schema: updateMatch[1]?.toLowerCase() || 'public',
+      table: unquoteIdentifier(updateMatch[2]!),
+    }
+  }
+
+  // DELETE FROM [schema.]table
+  const deleteMatch = normalized.match(/^DELETE\s+FROM\s+(?:([a-z_][a-z0-9_]*)\s*\.\s*)?([a-z_][a-z0-9_]*|"[^"]+")(?:\s|$)/i)
+  if (deleteMatch) {
+    return {
+      operation: 'DELETE',
+      schema: deleteMatch[1]?.toLowerCase() || 'public',
+      table: unquoteIdentifier(deleteMatch[2]!),
+    }
+  }
+
+  return null
+}
+
+/**
+ * Remove quotes from an identifier
+ */
+function unquoteIdentifier(identifier: string): string {
+  if (identifier.startsWith('"') && identifier.endsWith('"')) {
+    return identifier.slice(1, -1).replace(/""/g, '"')
+  }
+  return identifier.toLowerCase()
+}
+
+/**
+ * WAL event listener callback type
+ */
+export type WALEventListener = (entry: WALEntry) => void
+
+/**
+ * WAL flush event listener callback type
+ */
+export type WALFlushListener = (result: WALFlushResult) => void
+
+/**
+ * WAL Manager for capturing and buffering WAL entries
+ *
+ * Designed to work within DO memory constraints while providing
+ * reliable change capture for streaming to Iceberg.
+ */
+export class WALManager {
+  private config: Required<WALConfig>
+  private buffer: WALEntry[] = []
+  private bufferMemory = 0
+  private currentLsn: bigint = 0n
+  private totalCaptured: bigint = 0n
+  private totalFlushed: bigint = 0n
+  private totalDropped: bigint = 0n
+  private storage: WALStorage | null = null
+  private eventListeners: WALEventListener[] = []
+  private flushListeners: WALFlushListener[] = []
+  private currentTransactionId: string | null = null
+
+  constructor(config: WALConfig = {}, storage?: WALStorage) {
+    this.config = { ...DEFAULT_WAL_CONFIG, ...config }
+    this.storage = storage ?? null
+  }
+
+  /**
+   * Initialize WAL manager, restoring LSN from storage if available
+   */
+  async initialize(): Promise<void> {
+    if (this.storage) {
+      const lastLsn = await this.storage.getLastLsn()
+      if (lastLsn !== null) {
+        this.currentLsn = lastLsn
+      }
+    }
+  }
+
+  /**
+   * Get current configuration
+   */
+  getConfig(): Required<WALConfig> {
+    return { ...this.config }
+  }
+
+  /**
+   * Update configuration
+   */
+  setConfig(config: Partial<WALConfig>): void {
+    this.config = { ...this.config, ...config }
+  }
+
+  /**
+   * Check if WAL capture is enabled
+   */
+  isEnabled(): boolean {
+    return this.config.enabled
+  }
+
+  /**
+   * Enable or disable WAL capture
+   */
+  setEnabled(enabled: boolean): void {
+    this.config.enabled = enabled
+  }
+
+  /**
+   * Get current buffer statistics
+   */
+  getStats(): WALBufferStats {
+    return {
+      entryCount: this.buffer.length,
+      memoryBytes: this.bufferMemory,
+      oldestLsn: this.buffer.length > 0 ? this.buffer[0]!.lsn : null,
+      newestLsn: this.buffer.length > 0 ? this.buffer[this.buffer.length - 1]!.lsn : null,
+      totalEntriesCaptured: this.totalCaptured,
+      entriesFlushed: this.totalFlushed,
+      entriesDropped: this.totalDropped,
+    }
+  }
+
+  /**
+   * Get current LSN
+   */
+  getCurrentLsn(): bigint {
+    return this.currentLsn
+  }
+
+  /**
+   * Start a transaction context
+   */
+  beginTransaction(transactionId: string): void {
+    this.currentTransactionId = transactionId
+  }
+
+  /**
+   * End the current transaction context
+   */
+  endTransaction(): void {
+    this.currentTransactionId = null
+  }
+
+  /**
+   * Add an event listener for new WAL entries
+   */
+  onEntry(listener: WALEventListener): () => void {
+    this.eventListeners.push(listener)
+    return () => {
+      const index = this.eventListeners.indexOf(listener)
+      if (index >= 0) {
+        this.eventListeners.splice(index, 1)
+      }
+    }
+  }
+
+  /**
+   * Add a listener for flush events
+   */
+  onFlush(listener: WALFlushListener): () => void {
+    this.flushListeners.push(listener)
+    return () => {
+      const index = this.flushListeners.indexOf(listener)
+      if (index >= 0) {
+        this.flushListeners.splice(index, 1)
+      }
+    }
+  }
+
+  /**
+   * Capture a write operation
+   *
+   * @param sql - The SQL statement that was executed
+   * @param result - Query result containing affected rows
+   * @returns The captured WAL entry, or null if not captured
+   */
+  capture(
+    sql: string,
+    result: {
+      rows?: Record<string, unknown>[]
+      affectedRows?: number
+      operation?: WALOperation
+      table?: string
+      schema?: string
+      oldRows?: Record<string, unknown>[]
+    }
+  ): WALEntry | null {
+    if (!this.config.enabled) {
+      return null
+    }
+
+    // Parse SQL to detect write operation
+    let parsed = parseWriteOperation(sql)
+
+    // Use provided operation info if SQL parsing fails
+    if (!parsed && result.operation && result.table) {
+      parsed = {
+        operation: result.operation,
+        table: result.table,
+        schema: result.schema || 'public',
+      }
+    }
+
+    if (!parsed) {
+      return null
+    }
+
+    // Check if operation is in configured operations
+    if (!this.config.operations.includes(parsed.operation)) {
+      return null
+    }
+
+    // Check table filters
+    if (this.config.includeTables.length > 0 && !this.config.includeTables.includes(parsed.table)) {
+      return null
+    }
+    if (this.config.excludeTables.includes(parsed.table)) {
+      return null
+    }
+
+    // Generate LSN
+    this.currentLsn++
+    const lsn = this.currentLsn
+
+    // Build WAL entry based on operation type
+    const entry: WALEntry = {
+      lsn,
+      operation: parsed.operation,
+      schema: parsed.schema,
+      table: parsed.table,
+      timestamp: Date.now(),
+      transactionId: this.currentTransactionId ?? undefined,
+      sql: this.config.includeSQL ? sql : undefined,
+    }
+
+    // Populate row data based on operation type and capture mode
+    switch (parsed.operation) {
+      case 'INSERT':
+        if (result.rows && result.rows.length > 0) {
+          entry.newRow = result.rows[0]
+        }
+        break
+
+      case 'UPDATE':
+        if (result.rows && result.rows.length > 0) {
+          entry.newRow = result.rows[0]
+        }
+        if (result.oldRows && result.oldRows.length > 0) {
+          entry.oldRow = result.oldRows[0]
+
+          // Calculate changed columns for delta mode
+          if (this.config.captureMode === 'delta' && entry.newRow && entry.oldRow) {
+            entry.changedColumns = this.calculateChangedColumns(entry.oldRow, entry.newRow)
+          }
+        }
+        break
+
+      case 'DELETE':
+        if (result.oldRows && result.oldRows.length > 0) {
+          entry.oldRow = result.oldRows[0]
+        } else if (result.rows && result.rows.length > 0) {
+          // Some queries return deleted rows via RETURNING
+          entry.oldRow = result.rows[0]
+        }
+        break
+    }
+
+    // Add to buffer
+    this.addToBuffer(entry)
+
+    // Notify listeners
+    for (const listener of this.eventListeners) {
+      try {
+        listener(entry)
+      } catch (error) {
+        console.error('WAL event listener error:', error)
+      }
+    }
+
+    return entry
+  }
+
+  /**
+   * Add an entry to the buffer, managing memory constraints
+   */
+  private addToBuffer(entry: WALEntry): void {
+    const entrySize = estimateEntrySize(entry)
+
+    // Check if we need to flush before adding
+    const wouldExceedMemory = this.bufferMemory + entrySize > this.config.maxBufferBytes
+    const wouldExceedCount = this.buffer.length >= this.config.maxBufferEntries
+    const atFlushThreshold =
+      this.buffer.length >= this.config.maxBufferEntries * this.config.flushThreshold ||
+      this.bufferMemory >= this.config.maxBufferBytes * this.config.flushThreshold
+
+    // If we have storage and hit threshold, trigger async flush
+    if (atFlushThreshold && this.storage) {
+      // Fire and forget - don't block capture
+      void this.flush().catch((error) => {
+        console.error('WAL auto-flush error:', error)
+      })
+    }
+
+    // If buffer is full and no storage, drop oldest entries
+    if ((wouldExceedMemory || wouldExceedCount) && !this.storage) {
+      while (
+        this.buffer.length > 0 &&
+        (this.bufferMemory + entrySize > this.config.maxBufferBytes || this.buffer.length >= this.config.maxBufferEntries)
+      ) {
+        const dropped = this.buffer.shift()
+        if (dropped) {
+          this.bufferMemory -= estimateEntrySize(dropped)
+          this.totalDropped++
+        }
+      }
+    }
+
+    // Add entry to buffer
+    this.buffer.push(entry)
+    this.bufferMemory += entrySize
+    this.totalCaptured++
+  }
+
+  /**
+   * Calculate which columns changed between old and new row
+   */
+  private calculateChangedColumns(oldRow: Record<string, unknown>, newRow: Record<string, unknown>): string[] {
+    const changed: string[] = []
+    const allKeys = new Set([...Object.keys(oldRow), ...Object.keys(newRow)])
+
+    for (const key of allKeys) {
+      if (JSON.stringify(oldRow[key]) !== JSON.stringify(newRow[key])) {
+        changed.push(key)
+      }
+    }
+
+    return changed
+  }
+
+  /**
+   * Flush buffered entries to storage
+   */
+  async flush(): Promise<WALFlushResult> {
+    const startTime = performance.now()
+
+    if (this.buffer.length === 0) {
+      return {
+        entriesFlushed: 0,
+        lsnRange: null,
+        durationMs: performance.now() - startTime,
+      }
+    }
+
+    const entriesToFlush = [...this.buffer]
+    const serialized = entriesToFlush.map(serializeWALEntry)
+
+    const lsnRange = {
+      start: entriesToFlush[0]!.lsn,
+      end: entriesToFlush[entriesToFlush.length - 1]!.lsn,
+    }
+
+    if (this.storage) {
+      await this.storage.store(serialized)
+    }
+
+    // Clear buffer
+    this.buffer = []
+    this.bufferMemory = 0
+    this.totalFlushed += BigInt(entriesToFlush.length)
+
+    const result: WALFlushResult = {
+      entriesFlushed: entriesToFlush.length,
+      lsnRange,
+      durationMs: performance.now() - startTime,
+    }
+
+    // Notify flush listeners
+    for (const listener of this.flushListeners) {
+      try {
+        listener(result)
+      } catch (error) {
+        console.error('WAL flush listener error:', error)
+      }
+    }
+
+    return result
+  }
+
+  /**
+   * Get buffered entries without flushing
+   */
+  getBufferedEntries(): WALEntry[] {
+    return [...this.buffer]
+  }
+
+  /**
+   * Get serialized buffered entries
+   */
+  getSerializedBuffer(): SerializedWALEntry[] {
+    return this.buffer.map(serializeWALEntry)
+  }
+
+  /**
+   * Clear the buffer without flushing
+   */
+  clearBuffer(): number {
+    const count = this.buffer.length
+    this.buffer = []
+    this.bufferMemory = 0
+    this.totalDropped += BigInt(count)
+    return count
+  }
+
+  /**
+   * Read entries from storage (for replay)
+   */
+  async readFromStorage(fromLsn: bigint, limit?: number): Promise<WALEntry[]> {
+    if (!this.storage) {
+      return []
+    }
+
+    const serialized = await this.storage.readEntries(fromLsn, limit)
+    return serialized.map(deserializeWALEntry)
+  }
+}
+
+/**
+ * Durable Object Storage implementation of WALStorage
+ */
+export class DOWALStorage implements WALStorage {
+  private state: DurableObjectState
+  private prefix: string
+
+  constructor(state: DurableObjectState, prefix = 'wal:') {
+    this.state = state
+    this.prefix = prefix
+  }
+
+  async store(entries: SerializedWALEntry[]): Promise<void> {
+    if (entries.length === 0) {
+      return
+    }
+
+    // Store each entry with its LSN as the key
+    const puts: Record<string, SerializedWALEntry | { lsn: string }> = {}
+    for (const entry of entries) {
+      puts[`${this.prefix}${entry.lsn.padStart(20, '0')}`] = entry
+    }
+
+    // Also update the last LSN marker
+    puts[`${this.prefix}lastLsn`] = { lsn: entries[entries.length - 1]!.lsn }
+
+    await this.state.storage.put(puts)
+  }
+
+  async getLastLsn(): Promise<bigint | null> {
+    const marker = await this.state.storage.get<{ lsn: string }>(`${this.prefix}lastLsn`)
+    if (marker && marker.lsn) {
+      return BigInt(marker.lsn)
+    }
+    return null
+  }
+
+  async readEntries(fromLsn: bigint, limit = 100): Promise<SerializedWALEntry[]> {
+    const startKey = `${this.prefix}${fromLsn.toString().padStart(20, '0')}`
+    const endKey = `${this.prefix}9`.repeat(20) // Max LSN key
+
+    const entries = await this.state.storage.list<SerializedWALEntry>({
+      start: startKey,
+      end: endKey,
+      limit,
+    })
+
+    const result: SerializedWALEntry[] = []
+    for (const [key, value] of entries) {
+      if (!key.endsWith('lastLsn')) {
+        result.push(value)
+      }
+    }
+
+    return result
+  }
+}
+
+/**
+ * Create a WAL manager with DO storage
+ */
+export function createWALManager(state: DurableObjectState, config?: WALConfig): WALManager {
+  const storage = new DOWALStorage(state)
+  return new WALManager(config, storage)
+}