@dotdo/postgres 0.1.0

package/src/iceberg/transformer.test.ts

@@ -0,0 +1,689 @@
+/**
+ * WAL to Iceberg Transformer Tests
+ * Task: postgres-5iv - WAL to Iceberg transformation
+ */
+
+import { describe, it, expect, beforeEach, vi } from 'vitest'
+import {
+  WALIcebergTransformer,
+  createWALIcebergTransformer,
+  DefaultWALSchemaMapper,
+  TableSchemaMapper,
+  MultiTableTransformer,
+  type TransformCheckpoint,
+  type TransformBatchResult,
+} from './transformer'
+import type { SerializedWALEntry } from '../worker/wal'
+import { createWALSchema, createWALPartitionSpec, WAL_FIELD_IDS } from './schema'
+import { walEntryToParquetRecord } from './parquet'
+import type { IcebergSchema, IcebergPartitionSpec, IcebergField } from './types'
+
+/**
+ * Mock R2 bucket for testing
+ */
+class MockR2Bucket implements Partial<R2Bucket> {
+  private objects: Map<string, { data: Uint8Array; metadata: Record<string, string> }> = new Map()
+
+  async put(
+    key: string,
+    value: ArrayBuffer | ArrayBufferView | string | ReadableStream | Blob,
+    options?: R2PutOptions
+  ): Promise<R2Object> {
+    let data: Uint8Array
+    if (typeof value === 'string') {
+      data = new TextEncoder().encode(value)
+    } else if (value instanceof Uint8Array) {
+      data = value
+    } else if (value instanceof ArrayBuffer) {
+      data = new Uint8Array(value)
+    } else {
+      throw new Error('Unsupported value type')
+    }
+
+    this.objects.set(key, {
+      data,
+      metadata: options?.customMetadata ?? {},
+    })
+
+    return {
+      key,
+      size: data.length,
+      uploaded: new Date(),
+      httpMetadata: {},
+      customMetadata: options?.customMetadata ?? {},
+    } as R2Object
+  }
+
+  async get(key: string): Promise<R2ObjectBody | null> {
+    const obj = this.objects.get(key)
+    if (!obj) return null
+
+    return {
+      key,
+      size: obj.data.length,
+      uploaded: new Date(),
+      httpMetadata: {},
+      customMetadata: obj.metadata,
+      bytes: async () => obj.data,
+      text: async () => new TextDecoder().decode(obj.data),
+      json: async () => JSON.parse(new TextDecoder().decode(obj.data)),
+      body: new ReadableStream({
+        start(controller) {
+          controller.enqueue(obj.data)
+          controller.close()
+        },
+      }),
+      arrayBuffer: async () => obj.data.buffer.slice(obj.data.byteOffset, obj.data.byteOffset + obj.data.byteLength),
+      blob: async () => new Blob([obj.data]),
+      bodyUsed: false,
+    } as R2ObjectBody
+  }
+
+  async list(options?: R2ListOptions): Promise<R2Objects> {
+    const prefix = options?.prefix ?? ''
+    const objects: R2Object[] = []
+
+    for (const [key, obj] of this.objects) {
+      if (key.startsWith(prefix)) {
+        objects.push({
+          key,
+          size: obj.data.length,
+          uploaded: new Date(),
+          httpMetadata: {},
+          customMetadata: obj.metadata,
+        } as R2Object)
+      }
+    }
+
+    // Sort by key
+    objects.sort((a, b) => a.key.localeCompare(b.key))
+
+    const limit = options?.limit ?? 1000
+    const truncated = objects.length > limit
+
+    return {
+      objects: objects.slice(0, limit),
+      truncated,
+      cursor: truncated ? 'cursor' : undefined,
+      delimitedPrefixes: [],
+    }
+  }
+
+  async delete(keys: string | string[]): Promise<void> {
+    const keyArray = Array.isArray(keys) ? keys : [keys]
+    for (const key of keyArray) {
+      this.objects.delete(key)
+    }
+  }
+
+  // Helper methods for testing
+  getObject(key: string): { data: Uint8Array; metadata: Record<string, string> } | undefined {
+    return this.objects.get(key)
+  }
+
+  clear(): void {
+    this.objects.clear()
+  }
+
+  getKeys(): string[] {
+    return Array.from(this.objects.keys())
+  }
+}
+
+/**
+ * Mock DO storage for testing
+ */
+class MockDOStorage implements Partial<DurableObjectStorage> {
+  private data: Map<string, unknown> = new Map()
+
+  async get<T = unknown>(key: string): Promise<T | undefined> {
+    return this.data.get(key) as T | undefined
+  }
+
+  async put<T>(key: string, value: T): Promise<void> {
+    this.data.set(key, value)
+  }
+
+  async delete(keys: string | string[]): Promise<boolean | number> {
+    const keyArray = Array.isArray(keys) ? keys : [keys]
+    let deleted = 0
+    for (const key of keyArray) {
+      if (this.data.delete(key)) deleted++
+    }
+    return Array.isArray(keys) ? deleted : deleted > 0
+  }
+
+  // Helper for testing
+  clear(): void {
+    this.data.clear()
+  }
+}
+
+/**
+ * Mock DO state for testing
+ */
+class MockDOState implements Partial<DurableObjectState> {
+  storage = new MockDOStorage() as unknown as DurableObjectStorage
+  id = { toString: () => 'test-do-id' } as DurableObjectId
+
+  async blockConcurrencyWhile<T>(callback: () => Promise<T>): Promise<T> {
+    return callback()
+  }
+}
+
+/**
+ * Create test WAL entries
+ */
+function createTestWALEntries(count: number, startLsn = 1, startTime = Date.now()): SerializedWALEntry[] {
+  const entries: SerializedWALEntry[] = []
+  for (let i = 0; i < count; i++) {
+    const lsn = startLsn + i
+    entries.push({
+      lsn: String(lsn),
+      operation: i % 3 === 0 ? 'INSERT' : i % 3 === 1 ? 'UPDATE' : 'DELETE',
+      schema: 'public',
+      table: 'users',
+      newRow: i % 3 !== 2 ? { id: lsn, name: `User ${lsn}` } : undefined,
+      oldRow: i % 3 !== 0 ? { id: lsn, name: `Old User ${lsn}` } : undefined,
+      timestamp: startTime + i * 1000,
+    })
+  }
+  return entries
+}
+
+/**
+ * Compress data using gzip
+ */
+async function compressData(data: Uint8Array): Promise<Uint8Array> {
+  const stream = new CompressionStream('gzip')
+  const writer = stream.writable.getWriter()
+  const reader = stream.readable.getReader()
+
+  writer.write(data)
+  writer.close()
+
+  const chunks: Uint8Array[] = []
+  while (true) {
+    const { done, value } = await reader.read()
+    if (done) break
+    chunks.push(value)
+  }
+
+  const totalLength = chunks.reduce((sum, chunk) => sum + chunk.length, 0)
+  const result = new Uint8Array(totalLength)
+  let offset = 0
+  for (const chunk of chunks) {
+    result.set(chunk, offset)
+    offset += chunk.length
+  }
+
+  return result
+}
+
+/**
+ * Store WAL segment in mock bucket
+ */
+async function storeWALSegment(
+  bucket: MockR2Bucket,
+  doId: string,
+  entries: SerializedWALEntry[]
+): Promise<string> {
+  const firstEntry = entries[0]!
+  const lastEntry = entries[entries.length - 1]!
+
+  const now = new Date()
+  const year = now.getUTCFullYear()
+  const month = String(now.getUTCMonth() + 1).padStart(2, '0')
+  const day = String(now.getUTCDate()).padStart(2, '0')
+  const hour = String(now.getUTCHours()).padStart(2, '0')
+
+  const startLsn = firstEntry.lsn.padStart(20, '0')
+  const endLsn = lastEntry.lsn.padStart(20, '0')
+
+  const key = `wal/${doId}/${year}/${month}/${day}/${hour}/${startLsn}-${endLsn}.wal`
+
+  // Serialize as NDJSON
+  const jsonData = entries.map((e) => JSON.stringify(e)).join('\n')
+  const uncompressedData = new TextEncoder().encode(jsonData)
+
+  // Compress
+  const compressedData = await compressData(uncompressedData)
+
+  await bucket.put(key, compressedData, {
+    customMetadata: {
+      compression: 'gzip',
+      entryCount: String(entries.length),
+      startLsn: firstEntry.lsn,
+      endLsn: lastEntry.lsn,
+    },
+  })
+
+  return key
+}
+
+describe('WALIcebergTransformer', () => {
+  let walBucket: MockR2Bucket
+  let icebergBucket: MockR2Bucket
+  let state: MockDOState
+  const doId = 'test-do-123'
+
+  beforeEach(() => {
+    walBucket = new MockR2Bucket()
+    icebergBucket = new MockR2Bucket()
+    state = new MockDOState()
+  })
+
+  describe('initialization', () => {
+    it('should initialize with default configuration', async () => {
+      const transformer = createWALIcebergTransformer({
+        walBucket: walBucket as unknown as R2Bucket,
+        icebergBucket: icebergBucket as unknown as R2Bucket,
+        doId,
+        icebergTablePath: 'iceberg/wal-table',
+      })
+
+      await transformer.initialize()
+
+      expect(transformer.getStats()).toBeDefined()
+      expect(transformer.getCheckpoint()).toBeDefined()
+    })
+
+    it('should load existing checkpoint from DO storage', async () => {
+      // Store existing checkpoint
+      const existingCheckpoint: TransformCheckpoint = {
+        doId,
+        lastProcessedLsn: '100',
+        lastProcessedTimestamp: Date.now() - 10000,
+        entriesProcessed: 100,
+        snapshotsCreated: 2,
+        checkpointedAt: Date.now() - 5000,
+        processedSegmentKeys: ['wal/test/segment1.wal'],
+      }
+
+      await state.storage.put(`__wal_iceberg_checkpoint__${doId}`, existingCheckpoint)
+
+      const transformer = createWALIcebergTransformer({
+        walBucket: walBucket as unknown as R2Bucket,
+        icebergBucket: icebergBucket as unknown as R2Bucket,
+        doId,
+        icebergTablePath: 'iceberg/wal-table',
+        state: state as unknown as DurableObjectState,
+      })
+
+      await transformer.initialize()
+
+      const checkpoint = transformer.getCheckpoint()
+      expect(checkpoint?.lastProcessedLsn).toBe('100')
+      expect(checkpoint?.entriesProcessed).toBe(100)
+    })
+  })
+
+  describe('transformation', () => {
+    it('should transform WAL entries to Iceberg', async () => {
+      // Store WAL segment
+      const entries = createTestWALEntries(10)
+      await storeWALSegment(walBucket, doId, entries)
+
+      const transformer = createWALIcebergTransformer({
+        walBucket: walBucket as unknown as R2Bucket,
+        icebergBucket: icebergBucket as unknown as R2Bucket,
+        doId,
+        icebergTablePath: 'iceberg/wal-table',
+        state: state as unknown as DurableObjectState,
+      })
+
+      await transformer.initialize()
+      const result = await transformer.transformBatch()
+
+      expect(result.entriesProcessed).toBe(10)
+      expect(result.entriesSkipped).toBe(0)
+      expect(result.lsnRange).toBeDefined()
+      expect(result.lsnRange?.start).toBe('1')
+      expect(result.lsnRange?.end).toBe('10')
+    })
+
+    it('should skip already processed entries (deduplication)', async () => {
+      // Store WAL segment
+      const entries = createTestWALEntries(10)
+      await storeWALSegment(walBucket, doId, entries)
+
+      const transformer = createWALIcebergTransformer({
+        walBucket: walBucket as unknown as R2Bucket,
+        icebergBucket: icebergBucket as unknown as R2Bucket,
+        doId,
+        icebergTablePath: 'iceberg/wal-table',
+        state: state as unknown as DurableObjectState,
+        deduplication: true,
+      })
+
+      await transformer.initialize()
+
+      // First transformation
+      const result1 = await transformer.transformBatch()
+      expect(result1.entriesProcessed).toBe(10)
+
+      // Second transformation (should skip all)
+      const result2 = await transformer.transformBatch()
+      expect(result2.entriesProcessed).toBe(0)
+    })
+
+    it('should respect maxEntriesPerCycle', async () => {
+      // Store multiple WAL segments
+      for (let i = 0; i < 5; i++) {
+        const entries = createTestWALEntries(100, i * 100 + 1)
+        await storeWALSegment(walBucket, doId, entries)
+      }
+
+      const transformer = createWALIcebergTransformer({
+        walBucket: walBucket as unknown as R2Bucket,
+        icebergBucket: icebergBucket as unknown as R2Bucket,
+        doId,
+        icebergTablePath: 'iceberg/wal-table',
+        state: state as unknown as DurableObjectState,
+        maxEntriesPerCycle: 150,
+      })
+
+      await transformer.initialize()
+      const result = await transformer.transformBatch()
+
+      // Should process at most 150 entries
+      expect(result.entriesProcessed).toBeLessThanOrEqual(150)
+    })
+
+    it('should update checkpoint after transformation', async () => {
+      const entries = createTestWALEntries(10)
+      await storeWALSegment(walBucket, doId, entries)
+
+      const transformer = createWALIcebergTransformer({
+        walBucket: walBucket as unknown as R2Bucket,
+        icebergBucket: icebergBucket as unknown as R2Bucket,
+        doId,
+        icebergTablePath: 'iceberg/wal-table',
+        state: state as unknown as DurableObjectState,
+      })
+
+      await transformer.initialize()
+      await transformer.transformBatch()
+
+      const checkpoint = transformer.getCheckpoint()
+      expect(checkpoint?.lastProcessedLsn).toBe('10')
+      expect(checkpoint?.entriesProcessed).toBe(10)
+    })
+  })
+
+  describe('transformUntilCaughtUp', () => {
+    it('should transform all pending entries', async () => {
+      // Store multiple WAL segments with older timestamps to ensure lag detection works
+      const oldTimestamp = Date.now() - 10 * 60 * 1000 // 10 minutes ago
+      for (let i = 0; i < 5; i++) {
+        const entries = createTestWALEntries(100, i * 100 + 1, oldTimestamp + i * 1000)
+        await storeWALSegment(walBucket, doId, entries)
+      }
+
+      const transformer = createWALIcebergTransformer({
+        walBucket: walBucket as unknown as R2Bucket,
+        icebergBucket: icebergBucket as unknown as R2Bucket,
+        doId,
+        icebergTablePath: 'iceberg/wal-table',
+        state: state as unknown as DurableObjectState,
+        maxEntriesPerCycle: 100,
+        targetLagMs: 1, // Set very low target lag to ensure all batches are processed
+      })
+
+      await transformer.initialize()
+      const result = await transformer.transformUntilCaughtUp()
+
+      expect(result.totalEntriesProcessed).toBe(500)
+      expect(result.iterations).toBeGreaterThanOrEqual(5)
+    })
+
+    it('should stop when caught up', async () => {
+      // Store one small segment
+      const entries = createTestWALEntries(10)
+      await storeWALSegment(walBucket, doId, entries)
+
+      const transformer = createWALIcebergTransformer({
+        walBucket: walBucket as unknown as R2Bucket,
+        icebergBucket: icebergBucket as unknown as R2Bucket,
+        doId,
+        icebergTablePath: 'iceberg/wal-table',
+        state: state as unknown as DurableObjectState,
+      })
+
+      await transformer.initialize()
+      const result = await transformer.transformUntilCaughtUp(100)
+
+      // Should stop after processing all entries, not iterate 100 times
+      expect(result.iterations).toBeLessThanOrEqual(2)
+      expect(result.totalEntriesProcessed).toBe(10)
+    })
+  })
+
+  describe('resetCheckpoint', () => {
+    it('should reset all transformation state', async () => {
+      const entries = createTestWALEntries(10)
+      await storeWALSegment(walBucket, doId, entries)
+
+      const transformer = createWALIcebergTransformer({
+        walBucket: walBucket as unknown as R2Bucket,
+        icebergBucket: icebergBucket as unknown as R2Bucket,
+        doId,
+        icebergTablePath: 'iceberg/wal-table',
+        state: state as unknown as DurableObjectState,
+      })
+
+      await transformer.initialize()
+      await transformer.transformBatch()
+
+      // Verify some entries were processed
+      expect(transformer.getCheckpoint()?.entriesProcessed).toBe(10)
+
+      // Reset
+      await transformer.resetCheckpoint()
+
+      // Verify reset
+      const checkpoint = transformer.getCheckpoint()
+      expect(checkpoint?.lastProcessedLsn).toBe('0')
+      expect(checkpoint?.entriesProcessed).toBe(0)
+
+      const stats = transformer.getStats()
+      expect(stats.totalEntriesProcessed).toBe(0n)
+    })
+  })
+
+  describe('statistics', () => {
+    it('should track transformation statistics', async () => {
+      const entries = createTestWALEntries(50)
+      await storeWALSegment(walBucket, doId, entries)
+
+      const transformer = createWALIcebergTransformer({
+        walBucket: walBucket as unknown as R2Bucket,
+        icebergBucket: icebergBucket as unknown as R2Bucket,
+        doId,
+        icebergTablePath: 'iceberg/wal-table',
+        state: state as unknown as DurableObjectState,
+      })
+
+      await transformer.initialize()
+      await transformer.transformBatch()
+
+      const stats = transformer.getStats()
+      expect(stats.totalEntriesProcessed).toBe(50n)
+      expect(stats.lastTransformAt).not.toBeNull()
+    })
+  })
+})
+
+describe('DefaultWALSchemaMapper', () => {
+  it('should map WAL entries to default schema', () => {
+    const mapper = new DefaultWALSchemaMapper()
+
+    const entry: SerializedWALEntry = {
+      lsn: '123',
+      operation: 'INSERT',
+      schema: 'public',
+      table: 'users',
+      newRow: { id: 1, name: 'Alice' },
+      timestamp: Date.now(),
+    }
+
+    const record = mapper.mapEntry(entry, 'test-do')
+
+    expect(record.do_id).toBe('test-do')
+    expect(record.lsn).toBe('123')
+    expect(record.operation).toBe('INSERT')
+    expect(record.table_name).toBe('users')
+    expect(record.new_row).toBe('{"id":1,"name":"Alice"}')
+  })
+
+  it('should return valid Iceberg schema', () => {
+    const mapper = new DefaultWALSchemaMapper()
+    const schema = mapper.getSchema()
+
+    expect(schema.type).toBe('struct')
+    expect(schema.fields.length).toBeGreaterThan(0)
+  })
+
+  it('should return valid partition spec', () => {
+    const mapper = new DefaultWALSchemaMapper()
+    const spec = mapper.getPartitionSpec()
+
+    expect(spec['spec-id']).toBeDefined()
+    expect(spec.fields.length).toBeGreaterThan(0)
+  })
+})
+
+describe('TableSchemaMapper', () => {
+  it('should map WAL entries to table-specific schema', () => {
+    const schema: IcebergSchema = {
+      type: 'struct',
+      'schema-id': 0,
+      fields: [
+        { id: 1, name: 'id', required: true, type: 'long' },
+        { id: 2, name: 'name', required: false, type: 'string' },
+        { id: 3, name: 'email', required: false, type: 'string' },
+      ],
+    }
+
+    const partitionSpec: IcebergPartitionSpec = {
+      'spec-id': 0,
+      fields: [
+        { 'source-id': 1, 'field-id': 1000, name: 'id_bucket', transform: 'bucket[16]' },
+      ],
+    }
+
+    const mapper = new TableSchemaMapper('users', schema, partitionSpec)
+
+    const entry: SerializedWALEntry = {
+      lsn: '123',
+      operation: 'INSERT',
+      schema: 'public',
+      table: 'users',
+      newRow: { id: 1, name: 'Alice', email: 'alice@example.com' },
+      timestamp: Date.now(),
+    }
+
+    const record = mapper.mapEntry(entry, 'test-do')
+
+    expect(record.__do_id).toBe('test-do')
+    expect(record.__lsn).toBe('123')
+    expect(record.__operation).toBe('INSERT')
+    expect(record.id).toBe(1)
+    expect(record.name).toBe('Alice')
+    expect(record.email).toBe('alice@example.com')
+  })
+
+  it('should throw for wrong table entries', () => {
+    const schema: IcebergSchema = {
+      type: 'struct',
+      'schema-id': 0,
+      fields: [],
+    }
+
+    const partitionSpec: IcebergPartitionSpec = {
+      'spec-id': 0,
+      fields: [],
+    }
+
+    const mapper = new TableSchemaMapper('users', schema, partitionSpec)
+
+    const entry: SerializedWALEntry = {
+      lsn: '123',
+      operation: 'INSERT',
+      schema: 'public',
+      table: 'orders', // Wrong table
+      timestamp: Date.now(),
+    }
+
+    expect(() => mapper.mapEntry(entry, 'test-do')).toThrow()
+  })
+
+  it('should extract oldRow for DELETE operations', () => {
+    const schema: IcebergSchema = {
+      type: 'struct',
+      'schema-id': 0,
+      fields: [
+        { id: 1, name: 'id', required: true, type: 'long' },
+      ],
+    }
+
+    const partitionSpec: IcebergPartitionSpec = {
+      'spec-id': 0,
+      fields: [],
+    }
+
+    const mapper = new TableSchemaMapper('users', schema, partitionSpec)
+
+    const entry: SerializedWALEntry = {
+      lsn: '123',
+      operation: 'DELETE',
+      schema: 'public',
+      table: 'users',
+      oldRow: { id: 1 },
+      timestamp: Date.now(),
+    }
+
+    const record = mapper.mapEntry(entry, 'test-do')
+
+    expect(record.__operation).toBe('DELETE')
+    expect(record.id).toBe(1)
+  })
+})
+
+describe('MultiTableTransformer', () => {
+  let walBucket: MockR2Bucket
+  let icebergBucket: MockR2Bucket
+  let state: MockDOState
+
+  beforeEach(() => {
+    walBucket = new MockR2Bucket()
+    icebergBucket = new MockR2Bucket()
+    state = new MockDOState()
+  })
+
+  it('should register and retrieve table-specific transformers', async () => {
+    const multiTransformer = new MultiTableTransformer({
+      walBucket: walBucket as unknown as R2Bucket,
+      icebergBucket: icebergBucket as unknown as R2Bucket,
+      doId: 'test-do',
+      defaultTablePath: 'iceberg/default',
+      state: state as unknown as DurableObjectState,
+    })
+
+    multiTransformer.registerTable('users', 'iceberg/users')
+    multiTransformer.registerTable('orders', 'iceberg/orders')
+
+    await multiTransformer.initialize()
+
+    const usersTransformer = multiTransformer.getTransformer('users')
+    const ordersTransformer = multiTransformer.getTransformer('orders')
+    const otherTransformer = multiTransformer.getTransformer('other')
+
+    expect(usersTransformer).toBeDefined()
+    expect(ordersTransformer).toBeDefined()
+    expect(otherTransformer).toBeDefined()
+
+    // Other should use default transformer
+    expect(otherTransformer).toBe(multiTransformer.getTransformer('unknown'))
+  })
+})