@dotdo/postgres 0.1.0

package/src/worker/proxy.ts

@@ -0,0 +1,694 @@
+/**
+ * Capnweb Proxy for PostgresDO
+ *
+ * Provides WebSocket-based RPC session pooling from external clients to DOs.
+ * Enables promise pipelining and magic map support for external callers.
+ *
+ * ## Architecture
+ *
+ * External Client -> Worker (capnweb proxy) -> DO (via WebSocket RPC)
+ *
+ * The proxy maintains a pool of RPC sessions to DOs, allowing:
+ * - Connection pooling (no duplicate WebSocket connections per dbId)
+ * - Session reuse across multiple requests
+ * - Automatic cleanup of stale sessions
+ *
+ * ## Usage
+ *
+ * ```typescript
+ * const proxy = createPostgresProxy(env)
+ * const session = await proxy.getSession('my-database')
+ * const result = await session.rpcQuery('SELECT * FROM users')
+ * ```
+ */
+
+import { RpcTarget } from 'cloudflare:workers'
+import type {
+  QueryResultRow,
+} from './types'
+import type {
+  RpcQueryResult,
+  RpcBatchQuery,
+  RpcBatchResult,
+} from './rpc'
+
+/**
+ * Environment with Postgres DO binding
+ */
+export interface PostgresProxyEnv {
+  POSTGRES_DO: DurableObjectNamespace
+}
+
+/**
+ * Environment with DocumentDB DO binding
+ */
+export interface DocumentDBProxyEnv {
+  DOCUMENTDB_DO: DurableObjectNamespace
+}
+
+/**
+ * RPC session configuration
+ */
+export interface ProxySessionConfig {
+  /** Maximum age of a session before it's considered stale (ms) */
+  maxAgeMs?: number
+  /** Interval to check for stale sessions (ms) */
+  cleanupIntervalMs?: number
+  /** Enable session keepalive pings */
+  keepalive?: boolean
+  /** Keepalive interval (ms) */
+  keepaliveIntervalMs?: number
+}
+
+/**
+ * Session metadata for tracking
+ */
+interface SessionMetadata {
+  createdAt: number
+  lastUsedAt: number
+  requestCount: number
+}
+
+/**
+ * Proxy session that wraps a DO stub with RPC methods
+ *
+ * This extends RpcTarget to enable capnweb features like:
+ * - Promise pipelining from external clients
+ * - Magic map for N+1 elimination
+ */
+export class PostgresProxySession extends RpcTarget {
+  private stub: DurableObjectStub
+  private metadata: SessionMetadata
+
+  constructor(stub: DurableObjectStub) {
+    super()
+    this.stub = stub
+    this.metadata = {
+      createdAt: Date.now(),
+      lastUsedAt: Date.now(),
+      requestCount: 0,
+    }
+  }
+
+  /**
+   * Execute a SQL query via RPC
+   */
+  async rpcQuery<T extends QueryResultRow = QueryResultRow>(
+    sql: string,
+    params?: unknown[]
+  ): Promise<RpcQueryResult<T>> {
+    this.metadata.lastUsedAt = Date.now()
+    this.metadata.requestCount++
+    // Call the DO's rpcQuery method directly via RPC
+    return (this.stub as unknown as { rpcQuery: (sql: string, params?: unknown[]) => Promise<RpcQueryResult<T>> }).rpcQuery(sql, params)
+  }
+
+  /**
+   * Execute a SQL query and return only the first row
+   */
+  async rpcQueryOne<T extends QueryResultRow = QueryResultRow>(
+    sql: string,
+    params?: unknown[]
+  ): Promise<T | null> {
+    this.metadata.lastUsedAt = Date.now()
+    this.metadata.requestCount++
+    return (this.stub as unknown as { rpcQueryOne: (sql: string, params?: unknown[]) => Promise<T | null> }).rpcQueryOne(sql, params)
+  }
+
+  /**
+   * Execute a query and return a scalar value
+   */
+  async rpcQueryScalar<T = unknown>(sql: string, params?: unknown[]): Promise<T | null> {
+    this.metadata.lastUsedAt = Date.now()
+    this.metadata.requestCount++
+    return (this.stub as unknown as { rpcQueryScalar: <U>(sql: string, params?: unknown[]) => Promise<U | null> }).rpcQueryScalar<T>(sql, params)
+  }
+
+  /**
+   * Execute a batch of queries
+   */
+  async rpcBatch(queries: RpcBatchQuery[]): Promise<RpcBatchResult> {
+    this.metadata.lastUsedAt = Date.now()
+    this.metadata.requestCount++
+    return (this.stub as unknown as { rpcBatch: (queries: RpcBatchQuery[]) => Promise<RpcBatchResult> }).rpcBatch(queries)
+  }
+
+  /**
+   * Execute a batch of queries within a transaction
+   */
+  async rpcBatchTransaction(queries: RpcBatchQuery[]): Promise<RpcBatchResult> {
+    this.metadata.lastUsedAt = Date.now()
+    this.metadata.requestCount++
+    return (this.stub as unknown as { rpcBatchTransaction: (queries: RpcBatchQuery[]) => Promise<RpcBatchResult> }).rpcBatchTransaction(queries)
+  }
+
+  /**
+   * Execute a SQL statement that doesn't return rows
+   */
+  async rpcExecute(sql: string, params?: unknown[]): Promise<{ rowCount: number; durationMs: number }> {
+    this.metadata.lastUsedAt = Date.now()
+    this.metadata.requestCount++
+    return (this.stub as unknown as { rpcExecute: (sql: string, params?: unknown[]) => Promise<{ rowCount: number; durationMs: number }> }).rpcExecute(sql, params)
+  }
+
+  /**
+   * Health check - verify the database is responsive
+   */
+  async rpcPing(): Promise<{ ok: true; durationMs: number }> {
+    this.metadata.lastUsedAt = Date.now()
+    this.metadata.requestCount++
+    return (this.stub as unknown as { rpcPing: () => Promise<{ ok: true; durationMs: number }> }).rpcPing()
+  }
+
+  /**
+   * Get database version
+   */
+  async rpcVersion(): Promise<string> {
+    this.metadata.lastUsedAt = Date.now()
+    this.metadata.requestCount++
+    return (this.stub as unknown as { rpcVersion: () => Promise<string> }).rpcVersion()
+  }
+
+  /**
+   * List all tables in the public schema
+   */
+  async rpcListTables(): Promise<string[]> {
+    this.metadata.lastUsedAt = Date.now()
+    this.metadata.requestCount++
+    return (this.stub as unknown as { rpcListTables: () => Promise<string[]> }).rpcListTables()
+  }
+
+  /**
+   * Get table schema information
+   */
+  async rpcDescribeTable(tableName: string): Promise<Array<{
+    column_name: string
+    data_type: string
+    is_nullable: boolean
+    column_default: string | null
+  }>> {
+    this.metadata.lastUsedAt = Date.now()
+    this.metadata.requestCount++
+    return (this.stub as unknown as { rpcDescribeTable: (tableName: string) => Promise<Array<{
+      column_name: string
+      data_type: string
+      is_nullable: boolean
+      column_default: string | null
+    }>> }).rpcDescribeTable(tableName)
+  }
+
+  /**
+   * Get database statistics
+   */
+  async rpcGetStats(): Promise<{
+    queryCount: number
+    totalDurationMs: number
+    avgDurationMs: number
+    lastQueryAt: Date | null
+    uptime: number
+    shutdownStatus: string
+  }> {
+    this.metadata.lastUsedAt = Date.now()
+    this.metadata.requestCount++
+    return (this.stub as unknown as { rpcGetStats: () => Promise<{
+      queryCount: number
+      totalDurationMs: number
+      avgDurationMs: number
+      lastQueryAt: Date | null
+      uptime: number
+      shutdownStatus: string
+    }> }).rpcGetStats()
+  }
+
+  /**
+   * Get session metadata
+   */
+  getMetadata(): Readonly<SessionMetadata> {
+    return { ...this.metadata }
+  }
+
+  /**
+   * Check if session is stale
+   */
+  isStale(maxAgeMs: number): boolean {
+    return Date.now() - this.metadata.lastUsedAt > maxAgeMs
+  }
+}
+
+/**
+ * PostgreSQL Proxy Pool
+ *
+ * Maintains a pool of RPC sessions to PostgresDO instances.
+ * Each database ID maps to a single session for efficient connection reuse.
+ */
+export class PostgresProxyPool {
+  private sessions = new Map<string, PostgresProxySession>()
+  private env: PostgresProxyEnv
+  private config: Required<ProxySessionConfig>
+
+  constructor(env: PostgresProxyEnv, config?: ProxySessionConfig) {
+    this.env = env
+    this.config = {
+      maxAgeMs: config?.maxAgeMs ?? 5 * 60 * 1000, // 5 minutes
+      cleanupIntervalMs: config?.cleanupIntervalMs ?? 60 * 1000, // 1 minute
+      keepalive: config?.keepalive ?? false,
+      keepaliveIntervalMs: config?.keepaliveIntervalMs ?? 30 * 1000, // 30 seconds
+    }
+  }
+
+  /**
+   * Get or create an RPC session for a database
+   */
+  async getSession(dbId: string): Promise<PostgresProxySession> {
+    let session = this.sessions.get(dbId)
+
+    if (!session || session.isStale(this.config.maxAgeMs)) {
+      // Create new session
+      const doId = this.env.POSTGRES_DO.idFromName(dbId)
+      const stub = this.env.POSTGRES_DO.get(doId)
+      session = new PostgresProxySession(stub)
+      this.sessions.set(dbId, session)
+    }
+
+    return session
+  }
+
+  /**
+   * Remove a specific session
+   */
+  removeSession(dbId: string): boolean {
+    return this.sessions.delete(dbId)
+  }
+
+  /**
+   * Clean up stale sessions
+   */
+  cleanup(): number {
+    let removed = 0
+    for (const [dbId, session] of this.sessions) {
+      if (session.isStale(this.config.maxAgeMs)) {
+        this.sessions.delete(dbId)
+        removed++
+      }
+    }
+    return removed
+  }
+
+  /**
+   * Get pool statistics
+   */
+  getStats(): {
+    activeSessions: number
+    sessionIds: string[]
+    config: Required<ProxySessionConfig>
+  } {
+    return {
+      activeSessions: this.sessions.size,
+      sessionIds: Array.from(this.sessions.keys()),
+      config: this.config,
+    }
+  }
+
+  /**
+   * Clear all sessions
+   */
+  clear(): void {
+    this.sessions.clear()
+  }
+}
+
+/**
+ * DocumentDB (MongoDB) Proxy Session
+ *
+ * Similar to PostgresProxySession but for MongoDB-compatible operations
+ */
+export class DocumentDBProxySession extends RpcTarget {
+  private stub: DurableObjectStub
+  private metadata: SessionMetadata
+  private dbName: string
+
+  constructor(stub: DurableObjectStub, dbName: string) {
+    super()
+    this.stub = stub
+    this.dbName = dbName
+    this.metadata = {
+      createdAt: Date.now(),
+      lastUsedAt: Date.now(),
+      requestCount: 0,
+    }
+  }
+
+  /**
+   * Execute a MongoDB command
+   */
+  async command(cmd: Record<string, unknown>): Promise<{ ok: 1 | 0; [key: string]: unknown }> {
+    this.metadata.lastUsedAt = Date.now()
+    this.metadata.requestCount++
+
+    // Use fetch to call the DO's command endpoint
+    const response = await this.stub.fetch(new Request('https://internal/command', {
+      method: 'POST',
+      headers: { 'Content-Type': 'application/json' },
+      body: JSON.stringify({ db: this.dbName, cmd }),
+    }))
+
+    return response.json() as Promise<{ ok: 1 | 0; [key: string]: unknown }>
+  }
+
+  /**
+   * Find documents in a collection
+   */
+  async find<T = Record<string, unknown>>(
+    collection: string,
+    filter: Record<string, unknown> = {},
+    options?: { limit?: number; skip?: number; sort?: Record<string, 1 | -1>; projection?: Record<string, 0 | 1> }
+  ): Promise<T[]> {
+    const result = await this.command({
+      find: collection,
+      filter,
+      ...options,
+    })
+
+    if (result.ok === 1 && result.cursor) {
+      return (result.cursor as { firstBatch: T[] }).firstBatch
+    }
+
+    return []
+  }
+
+  /**
+   * Insert documents into a collection
+   */
+  async insertMany<T extends Record<string, unknown>>(
+    collection: string,
+    documents: T[]
+  ): Promise<{ insertedCount: number }> {
+    const result = await this.command({
+      insert: collection,
+      documents,
+    })
+
+    return { insertedCount: (result.n as number) ?? 0 }
+  }
+
+  /**
+   * Insert a single document
+   */
+  async insertOne<T extends Record<string, unknown>>(
+    collection: string,
+    document: T
+  ): Promise<{ insertedCount: number }> {
+    return this.insertMany(collection, [document])
+  }
+
+  /**
+   * Update documents in a collection
+   */
+  async updateMany(
+    collection: string,
+    filter: Record<string, unknown>,
+    update: Record<string, unknown>
+  ): Promise<{ matchedCount: number; modifiedCount: number }> {
+    const result = await this.command({
+      update: collection,
+      updates: [{ q: filter, u: update, multi: true }],
+    })
+
+    return {
+      matchedCount: (result.n as number) ?? 0,
+      modifiedCount: (result.nModified as number) ?? 0,
+    }
+  }
+
+  /**
+   * Update a single document
+   */
+  async updateOne(
+    collection: string,
+    filter: Record<string, unknown>,
+    update: Record<string, unknown>
+  ): Promise<{ matchedCount: number; modifiedCount: number }> {
+    const result = await this.command({
+      update: collection,
+      updates: [{ q: filter, u: update, multi: false }],
+    })
+
+    return {
+      matchedCount: (result.n as number) ?? 0,
+      modifiedCount: (result.nModified as number) ?? 0,
+    }
+  }
+
+  /**
+   * Delete documents from a collection
+   */
+  async deleteMany(
+    collection: string,
+    filter: Record<string, unknown>
+  ): Promise<{ deletedCount: number }> {
+    const result = await this.command({
+      delete: collection,
+      deletes: [{ q: filter, limit: 0 }],
+    })
+
+    return { deletedCount: (result.n as number) ?? 0 }
+  }
+
+  /**
+   * Delete a single document
+   */
+  async deleteOne(
+    collection: string,
+    filter: Record<string, unknown>
+  ): Promise<{ deletedCount: number }> {
+    const result = await this.command({
+      delete: collection,
+      deletes: [{ q: filter, limit: 1 }],
+    })
+
+    return { deletedCount: (result.n as number) ?? 0 }
+  }
+
+  /**
+   * Count documents in a collection
+   */
+  async countDocuments(
+    collection: string,
+    filter: Record<string, unknown> = {}
+  ): Promise<number> {
+    const result = await this.command({
+      count: collection,
+      query: filter,
+    })
+
+    return (result.n as number) ?? 0
+  }
+
+  /**
+   * Run an aggregation pipeline
+   */
+  async aggregate<T = Record<string, unknown>>(
+    collection: string,
+    pipeline: Record<string, unknown>[]
+  ): Promise<T[]> {
+    const result = await this.command({
+      aggregate: collection,
+      pipeline,
+      cursor: {},
+    })
+
+    if (result.ok === 1 && result.cursor) {
+      return (result.cursor as { firstBatch: T[] }).firstBatch
+    }
+
+    return []
+  }
+
+  /**
+   * List collections in the database
+   */
+  async listCollections(): Promise<Array<{ name: string; type: string }>> {
+    const result = await this.command({
+      listCollections: 1,
+    })
+
+    if (result.ok === 1 && result.cursor) {
+      return (result.cursor as { firstBatch: Array<{ name: string; type: string }> }).firstBatch
+    }
+
+    return []
+  }
+
+  /**
+   * Get session metadata
+   */
+  getMetadata(): Readonly<SessionMetadata> {
+    return { ...this.metadata }
+  }
+
+  /**
+   * Check if session is stale
+   */
+  isStale(maxAgeMs: number): boolean {
+    return Date.now() - this.metadata.lastUsedAt > maxAgeMs
+  }
+}
+
+/**
+ * DocumentDB Proxy Pool
+ *
+ * Maintains a pool of RPC sessions to DocumentDBDO instances.
+ */
+export class DocumentDBProxyPool {
+  private sessions = new Map<string, DocumentDBProxySession>()
+  private env: DocumentDBProxyEnv
+  private config: Required<ProxySessionConfig>
+
+  constructor(env: DocumentDBProxyEnv, config?: ProxySessionConfig) {
+    this.env = env
+    this.config = {
+      maxAgeMs: config?.maxAgeMs ?? 5 * 60 * 1000,
+      cleanupIntervalMs: config?.cleanupIntervalMs ?? 60 * 1000,
+      keepalive: config?.keepalive ?? false,
+      keepaliveIntervalMs: config?.keepaliveIntervalMs ?? 30 * 1000,
+    }
+  }
+
+  /**
+   * Get or create an RPC session for a database
+   */
+  async getSession(dbName: string): Promise<DocumentDBProxySession> {
+    let session = this.sessions.get(dbName)
+
+    if (!session || session.isStale(this.config.maxAgeMs)) {
+      const doId = this.env.DOCUMENTDB_DO.idFromName(dbName)
+      const stub = this.env.DOCUMENTDB_DO.get(doId)
+      session = new DocumentDBProxySession(stub, dbName)
+      this.sessions.set(dbName, session)
+    }
+
+    return session
+  }
+
+  /**
+   * Remove a specific session
+   */
+  removeSession(dbName: string): boolean {
+    return this.sessions.delete(dbName)
+  }
+
+  /**
+   * Clean up stale sessions
+   */
+  cleanup(): number {
+    let removed = 0
+    for (const [dbName, session] of this.sessions) {
+      if (session.isStale(this.config.maxAgeMs)) {
+        this.sessions.delete(dbName)
+        removed++
+      }
+    }
+    return removed
+  }
+
+  /**
+   * Get pool statistics
+   */
+  getStats(): {
+    activeSessions: number
+    sessionIds: string[]
+    config: Required<ProxySessionConfig>
+  } {
+    return {
+      activeSessions: this.sessions.size,
+      sessionIds: Array.from(this.sessions.keys()),
+      config: this.config,
+    }
+  }
+
+  /**
+   * Clear all sessions
+   */
+  clear(): void {
+    this.sessions.clear()
+  }
+}
+
+/**
+ * Create a PostgreSQL proxy pool
+ */
+export function createPostgresProxy(env: PostgresProxyEnv, config?: ProxySessionConfig): PostgresProxyPool {
+  return new PostgresProxyPool(env, config)
+}
+
+/**
+ * Create a DocumentDB proxy pool
+ */
+export function createDocumentDBProxy(env: DocumentDBProxyEnv, config?: ProxySessionConfig): DocumentDBProxyPool {
+  return new DocumentDBProxyPool(env, config)
+}
+
+/**
+ * Generic proxy pool creator for any DO type
+ *
+ * This is a higher-order function that creates a proxy pool for any DO namespace.
+ * Useful for creating custom proxy pools with different session types.
+ *
+ * Uses Record<string, DurableObjectNamespace> per project convention for pure dictionary types.
+ * @see CLAUDE.md "Record vs Index Signature Usage"
+ */
+export function createProxyPool<TSession, TEnv extends Record<string, DurableObjectNamespace>>(
+  getDO: (env: TEnv, id: string) => DurableObjectStub,
+  createSession: (stub: DurableObjectStub, id: string) => TSession,
+  isStale: (session: TSession, maxAgeMs: number) => boolean,
+  config?: ProxySessionConfig
+) {
+  const sessions = new Map<string, TSession>()
+  const resolvedConfig: Required<ProxySessionConfig> = {
+    maxAgeMs: config?.maxAgeMs ?? 5 * 60 * 1000,
+    cleanupIntervalMs: config?.cleanupIntervalMs ?? 60 * 1000,
+    keepalive: config?.keepalive ?? false,
+    keepaliveIntervalMs: config?.keepaliveIntervalMs ?? 30 * 1000,
+  }
+
+  return {
+    async getSession(env: TEnv, id: string): Promise<TSession> {
+      let session = sessions.get(id)
+
+      if (!session || isStale(session, resolvedConfig.maxAgeMs)) {
+        const stub = getDO(env, id)
+        session = createSession(stub, id)
+        sessions.set(id, session)
+      }
+
+      return session
+    },
+
+    removeSession(id: string): boolean {
+      return sessions.delete(id)
+    },
+
+    cleanup(): number {
+      let removed = 0
+      for (const [id, session] of sessions) {
+        if (isStale(session, resolvedConfig.maxAgeMs)) {
+          sessions.delete(id)
+          removed++
+        }
+      }
+      return removed
+    },
+
+    getStats() {
+      return {
+        activeSessions: sessions.size,
+        sessionIds: Array.from(sessions.keys()),
+        config: resolvedConfig,
+      }
+    },
+
+    clear(): void {
+      sessions.clear()
+    },
+  }
+}