@dotdo/postgres 0.1.0

package/src/pglite/query-executor.ts

@@ -0,0 +1,614 @@
+/**
+ * PGLite Query Executor
+ *
+ * A wrapper around PGLite that provides:
+ * - Query streaming via AsyncIterable
+ * - Batch query execution
+ * - Prepared statement management
+ * - Query statistics tracking
+ * - Query lifecycle hooks
+ * - Query cancellation via AbortSignal
+ * - Query timeout support
+ * - Type coercion for PostgreSQL types
+ *
+ * @example
+ * ```typescript
+ * import { createPGLiteQueryExecutor } from '@dotdo/postgres/pglite'
+ *
+ * const executor = await createPGLiteQueryExecutor()
+ * await executor.waitReady
+ *
+ * // Query with metrics
+ * const result = await executor.query('SELECT * FROM users')
+ * console.log(result.metrics.executionTimeMs)
+ *
+ * // Stream results
+ * for await (const row of executor.queryStream('SELECT * FROM large_table')) {
+ *   console.log(row)
+ * }
+ *
+ * // Batch queries
+ * const results = await executor.batch([
+ *   { sql: 'SELECT 1' },
+ *   { sql: 'SELECT 2' }
+ * ])
+ *
+ * // Prepared statements
+ * const stmt = await executor.prepare('SELECT * FROM users WHERE id = $1')
+ * const result1 = await stmt.execute([1])
+ * const result2 = await stmt.execute([2])
+ * await stmt.close()
+ * ```
+ */
+
+import { PGlite, type PGliteOptions, type Results, type QueryOptions, type Transaction } from '@dotdo/pglite'
+
+/**
+ * Query execution metrics
+ */
+export interface QueryMetrics {
+  executionTimeMs: number
+  planningTimeMs?: number
+  rowsReturned: number
+}
+
+/**
+ * Field information in query results
+ */
+export interface FieldInfo {
+  name: string
+  dataTypeID: number
+}
+
+/**
+ * Extended query result with metrics
+ */
+export interface QueryResultWithMetrics<T = unknown> {
+  rows: T[]
+  fields: FieldInfo[]
+  affectedRows: number
+  metrics: QueryMetrics
+}
+
+/**
+ * Query statistics
+ */
+export interface QueryStats {
+  queryCount: number
+  totalQueryTimeMs: number
+  averageQueryTimeMs: number
+  errorCount: number
+}
+
+/**
+ * Hook types for query lifecycle
+ */
+export type HookType = 'before' | 'after' | 'error'
+
+/**
+ * Query hook function signature
+ */
+export type QueryHook = (
+  query: { sql: string; params?: unknown[] },
+  result?: unknown,
+  error?: Error
+) => Promise<void> | void
+
+/**
+ * Prepared statement interface
+ */
+export interface PreparedStatement<T = unknown> {
+  execute(params?: unknown[]): Promise<QueryResultWithMetrics<T>>
+  close(): Promise<void>
+}
+
+/**
+ * Batch query item
+ */
+export interface BatchQueryItem {
+  sql: string
+  params?: unknown[]
+}
+
+/**
+ * Extended query options
+ */
+export interface ExtendedQueryOptions extends QueryOptions {
+  signal?: AbortSignal
+  timeout?: number
+}
+
+/**
+ * Interval object for structured interval representation
+ */
+export interface IntervalObject {
+  years: number
+  months: number
+  days: number
+  hours: number
+  minutes: number
+  seconds: number
+}
+
+/**
+ * Point object for geometric types
+ */
+export interface PointObject {
+  x: number
+  y: number
+}
+
+/**
+ * Options for PGLiteQueryExecutor
+ */
+export interface PGLiteQueryExecutorOptions extends PGliteOptions {
+  /** Enable type coercion for PostgreSQL types */
+  enableTypeCoercion?: boolean
+}
+
+/**
+ * PGLite Query Executor with enhanced functionality
+ */
+export class PGLiteQueryExecutor {
+  private _pglite: PGlite
+  private _queryCount = 0
+  private _totalQueryTimeMs = 0
+  private _errorCount = 0
+  private _beforeHooks: QueryHook[] = []
+  private _afterHooks: QueryHook[] = []
+  private _errorHooks: QueryHook[] = []
+  private _preparedStatements: Map<string, { stmtName: string; sql: string }> = new Map()
+  private _preparedCounter = 0
+  private _enableTypeCoercion: boolean
+
+  /**
+   * Promise that resolves when the database is ready
+   */
+  public readonly waitReady: Promise<void>
+
+  constructor(options?: PGLiteQueryExecutorOptions | string) {
+    let pgliteOptions: PGliteOptions | string | undefined
+    this._enableTypeCoercion = false
+
+    if (typeof options === 'object' && options !== null) {
+      const { enableTypeCoercion, ...rest } = options
+      this._enableTypeCoercion = enableTypeCoercion ?? false
+      pgliteOptions = rest
+    } else {
+      pgliteOptions = options
+    }
+
+    if (typeof pgliteOptions === 'string') {
+      this._pglite = new PGlite(pgliteOptions)
+    } else {
+      this._pglite = new PGlite(pgliteOptions)
+    }
+    this.waitReady = this._pglite.waitReady
+  }
+
+  /**
+   * Whether the database is ready for queries
+   */
+  get ready(): boolean {
+    return this._pglite.ready
+  }
+
+  /**
+   * Whether the database is closed
+   */
+  get closed(): boolean {
+    return this._pglite.closed
+  }
+
+  /**
+   * Execute a SQL query with metrics and optional timeout/cancellation
+   */
+  async query<T = unknown>(
+    sql: string,
+    params?: unknown[],
+    options?: ExtendedQueryOptions
+  ): Promise<QueryResultWithMetrics<T>> {
+    const startTime = performance.now()
+
+    // Check for abort signal
+    if (options?.signal?.aborted) {
+      throw new Error('Query cancelled')
+    }
+
+    // Run before hooks
+    await this._runHooks('before', { sql, params })
+
+    let result: Results<T>
+    let error: Error | undefined
+
+    try {
+      // Create a promise that races query with timeout/abort
+      result = await this._executeWithOptions<T>(sql, params, options)
+    } catch (e) {
+      error = e as Error
+      this._errorCount++
+      await this._runHooks('error', { sql, params }, undefined, error)
+      throw error
+    }
+
+    const executionTimeMs = performance.now() - startTime
+    this._queryCount++
+    this._totalQueryTimeMs += executionTimeMs
+
+    const resultWithMetrics: QueryResultWithMetrics<T> = {
+      rows: this._enableTypeCoercion ? this._coerceTypes(result.rows, result.fields) : result.rows,
+      fields: result.fields as FieldInfo[],
+      affectedRows: result.affectedRows ?? 0,
+      metrics: {
+        executionTimeMs,
+        rowsReturned: result.rows.length,
+      },
+    }
+
+    // Run after hooks
+    await this._runHooks('after', { sql, params }, resultWithMetrics)
+
+    return resultWithMetrics
+  }
+
+  /**
+   * Execute query with timeout and abort signal support
+   *
+   * Note: PGLite runs synchronously in WASM, so true query cancellation mid-execution
+   * is not possible. For timeout support, we use PostgreSQL's statement_timeout setting.
+   */
+  private async _executeWithOptions<T>(
+    sql: string,
+    params?: unknown[],
+    options?: ExtendedQueryOptions
+  ): Promise<Results<T>> {
+    const { signal, timeout, ...queryOptions } = options || {}
+
+    // Check abort signal before query
+    if (signal?.aborted) {
+      throw new Error('Query cancelled')
+    }
+
+    // For timeout, use PostgreSQL's statement_timeout
+    // This is the only reliable way to timeout queries in PGLite
+    if (timeout !== undefined && timeout > 0) {
+      // Set statement_timeout for this query
+      await this._pglite.query(`SET LOCAL statement_timeout = ${timeout}`)
+
+      try {
+        // Execute within a transaction to ensure LOCAL setting applies
+        return await this._pglite.transaction(async (tx) => {
+          return await tx.query<T>(sql, params)
+        })
+      } catch (error) {
+        // Check if it's a timeout error from PostgreSQL
+        const errorMessage = (error as Error).message || ''
+        if (errorMessage.includes('statement timeout') ||
+            errorMessage.includes('canceling statement due to statement timeout')) {
+          throw new Error('Query timeout')
+        }
+        throw error
+      }
+    }
+
+    // Standard execution without timeout
+    return this._pglite.query<T>(sql, params, queryOptions)
+  }
+
+  /**
+   * Coerce PostgreSQL types to JavaScript types
+   */
+  private _coerceTypes<T>(rows: T[], fields: { name: string; dataTypeID: number }[]): T[] {
+    return rows.map((row) => {
+      const coercedRow = { ...row } as Record<string, unknown>
+
+      for (const field of fields) {
+        const value = coercedRow[field.name]
+        if (value === null || value === undefined) continue
+
+        coercedRow[field.name] = this._coerceValue(value, field.dataTypeID)
+      }
+
+      return coercedRow as T
+    })
+  }
+
+  /**
+   * Coerce a single value based on PostgreSQL type OID
+   */
+  private _coerceValue(value: unknown, dataTypeID: number): unknown {
+    // PostgreSQL type OIDs
+    const BIGINT_OID = 20
+    const NUMERIC_OID = 1700
+    const MONEY_OID = 790
+    const INTERVAL_OID = 1186
+    const POINT_OID = 600
+    const BIGINT_ARRAY_OID = 1016
+
+    switch (dataTypeID) {
+      case BIGINT_OID:
+        // Convert bigint string to BigInt
+        if (typeof value === 'string') {
+          return BigInt(value)
+        }
+        return value
+
+      case NUMERIC_OID:
+        // Convert numeric string to number
+        if (typeof value === 'string') {
+          return parseFloat(value)
+        }
+        return value
+
+      case MONEY_OID:
+        // Convert money string to number (strip currency symbol and parse)
+        if (typeof value === 'string') {
+          return parseFloat(value.replace(/[^0-9.-]/g, ''))
+        }
+        return value
+
+      case INTERVAL_OID:
+        // Parse interval string to object
+        if (typeof value === 'string') {
+          return this._parseInterval(value)
+        }
+        return value
+
+      case POINT_OID:
+        // Parse point string to object
+        if (typeof value === 'string') {
+          return this._parsePoint(value)
+        }
+        return value
+
+      case BIGINT_ARRAY_OID:
+        // Convert bigint array elements to BigInt
+        if (Array.isArray(value)) {
+          return value.map((v) => (typeof v === 'string' ? BigInt(v) : v))
+        }
+        return value
+
+      default:
+        return value
+    }
+  }
+
+  /**
+   * Parse PostgreSQL interval string to structured object
+   */
+  private _parseInterval(value: string): IntervalObject {
+    const result: IntervalObject = {
+      years: 0,
+      months: 0,
+      days: 0,
+      hours: 0,
+      minutes: 0,
+      seconds: 0,
+    }
+
+    // Parse various interval formats
+    // e.g., "1 year 2 mons 3 days 04:05:06"
+    const yearMatch = value.match(/(\d+)\s*years?/i)
+    const monthMatch = value.match(/(\d+)\s*(?:months?|mons?)/i)
+    const dayMatch = value.match(/(\d+)\s*days?/i)
+    const timeMatch = value.match(/(\d+):(\d+):(\d+(?:\.\d+)?)/i)
+
+    if (yearMatch) result.years = parseInt(yearMatch[1], 10)
+    if (monthMatch) result.months = parseInt(monthMatch[1], 10)
+    if (dayMatch) result.days = parseInt(dayMatch[1], 10)
+    if (timeMatch) {
+      result.hours = parseInt(timeMatch[1], 10)
+      result.minutes = parseInt(timeMatch[2], 10)
+      result.seconds = parseFloat(timeMatch[3])
+    }
+
+    return result
+  }
+
+  /**
+   * Parse PostgreSQL point string to coordinate object
+   */
+  private _parsePoint(value: string): PointObject {
+    // Parse "(x,y)" format
+    const match = value.match(/\(([^,]+),([^)]+)\)/)
+    if (match) {
+      return {
+        x: parseFloat(match[1]),
+        y: parseFloat(match[2]),
+      }
+    }
+    return { x: 0, y: 0 }
+  }
+
+  /**
+   * Execute SQL statements (without returning results)
+   */
+  async exec(sql: string, options?: QueryOptions): Promise<Array<Results>> {
+    return this._pglite.exec(sql, options)
+  }
+
+  /**
+   * Execute a transaction
+   */
+  async transaction<T>(callback: (tx: Transaction) => Promise<T>): Promise<T> {
+    return this._pglite.transaction(callback)
+  }
+
+  /**
+   * Stream query results as an async iterable
+   */
+  async *queryStream<T = unknown>(
+    sql: string,
+    params?: unknown[],
+    options?: { batchSize?: number }
+  ): AsyncIterable<T> {
+    const batchSize = options?.batchSize ?? 100
+
+    // For streaming, we use a cursor-based approach
+    // First, execute the full query (PGLite doesn't have native cursor support outside transactions)
+    const result = await this.query<T>(sql, params)
+
+    // Yield rows in batches
+    for (let i = 0; i < result.rows.length; i += batchSize) {
+      const batch = result.rows.slice(i, i + batchSize)
+      for (const row of batch) {
+        yield row
+      }
+    }
+  }
+
+  /**
+   * Execute multiple queries in a batch
+   */
+  async batch(queries: BatchQueryItem[]): Promise<QueryResultWithMetrics[]> {
+    const results: QueryResultWithMetrics[] = []
+
+    for (const { sql, params } of queries) {
+      const result = await this.query(sql, params)
+      results.push(result)
+    }
+
+    return results
+  }
+
+  /**
+   * Create a prepared statement
+   */
+  async prepare<T = unknown>(sql: string): Promise<PreparedStatement<T>> {
+    const stmtName = `stmt_${++this._preparedCounter}`
+
+    // Create the prepared statement in PostgreSQL
+    await this._pglite.exec(`PREPARE ${stmtName} AS ${sql}`)
+
+    this._preparedStatements.set(stmtName, { stmtName, sql })
+
+    const self = this
+
+    return {
+      async execute(params?: unknown[]): Promise<QueryResultWithMetrics<T>> {
+        // Build the EXECUTE statement with parameters
+        let executeSql = `EXECUTE ${stmtName}`
+        if (params && params.length > 0) {
+          const paramStrings = params.map((p) => {
+            if (p === null) return 'NULL'
+            if (typeof p === 'string') return `'${p.replace(/'/g, "''")}'`
+            if (typeof p === 'number') return String(p)
+            if (typeof p === 'boolean') return p ? 'TRUE' : 'FALSE'
+            return String(p)
+          })
+          executeSql += `(${paramStrings.join(', ')})`
+        }
+
+        return self.query<T>(executeSql)
+      },
+
+      async close(): Promise<void> {
+        await self._pglite.exec(`DEALLOCATE ${stmtName}`)
+        self._preparedStatements.delete(stmtName)
+      },
+    }
+  }
+
+  /**
+   * Get query statistics
+   */
+  getStats(): QueryStats {
+    return {
+      queryCount: this._queryCount,
+      totalQueryTimeMs: this._totalQueryTimeMs,
+      averageQueryTimeMs: this._queryCount > 0 ? this._totalQueryTimeMs / this._queryCount : 0,
+      errorCount: this._errorCount,
+    }
+  }
+
+  /**
+   * Add a query lifecycle hook
+   * @returns A function to remove the hook
+   */
+  addHook(type: HookType, hook: QueryHook): () => void {
+    let hooks: QueryHook[]
+
+    switch (type) {
+      case 'before':
+        hooks = this._beforeHooks
+        break
+      case 'after':
+        hooks = this._afterHooks
+        break
+      case 'error':
+        hooks = this._errorHooks
+        break
+    }
+
+    hooks.push(hook)
+
+    return () => {
+      const index = hooks.indexOf(hook)
+      if (index > -1) {
+        hooks.splice(index, 1)
+      }
+    }
+  }
+
+  /**
+   * Run hooks of a specific type
+   */
+  private async _runHooks(
+    type: HookType,
+    query: { sql: string; params?: unknown[] },
+    result?: unknown,
+    error?: Error
+  ): Promise<void> {
+    let hooks: QueryHook[]
+
+    switch (type) {
+      case 'before':
+        hooks = this._beforeHooks
+        break
+      case 'after':
+        hooks = this._afterHooks
+        break
+      case 'error':
+        hooks = this._errorHooks
+        break
+    }
+
+    for (const hook of hooks) {
+      await hook(query, result, error)
+    }
+  }
+
+  /**
+   * Close the database
+   */
+  async close(): Promise<void> {
+    // Deallocate all prepared statements
+    const stmtNames = Array.from(this._preparedStatements.keys())
+    for (const stmtName of stmtNames) {
+      try {
+        await this._pglite.exec(`DEALLOCATE ${stmtName}`)
+      } catch {
+        // Ignore errors during cleanup
+      }
+    }
+    this._preparedStatements.clear()
+
+    return this._pglite.close()
+  }
+
+  /**
+   * Access to underlying PGLite instance
+   */
+  get pglite(): PGlite {
+    return this._pglite
+  }
+}
+
+/**
+ * Create a new PGLiteQueryExecutor instance
+ */
+export function createPGLiteQueryExecutor(
+  options?: PGLiteQueryExecutorOptions | string
+): PGLiteQueryExecutor {
+  return new PGLiteQueryExecutor(options)
+}