@dotdo/postgres 0.1.0

package/src/observability/instrumentation.ts

@@ -0,0 +1,788 @@
+/**
+ * PostgreSQL Instrumentation
+ *
+ * Provides instrumented wrappers for PostgreSQL operations in Cloudflare Workers.
+ * Implements tracing for queries, cache operations, DO lifecycle, and errors.
+ */
+
+import type {
+  Tracer,
+  SpanContext,
+  MetricsCollector,
+  QuerySpanAttributes,
+  CacheSpanAttributes,
+  DOLifecycleSpanAttributes,
+  ErrorSpanAttributes,
+  SpanOptions,
+} from './types'
+import { SpanKind, SpanStatusCode } from './types'
+import { MetricNames } from './metrics'
+import { extractSpanContext, injectSpanContext } from './propagation'
+
+// Re-export for external use
+export type { ActiveSpan } from './types'
+export { createChildContext } from './propagation'
+
+// ============================================================================
+// Query Instrumentation
+// ============================================================================
+
+/**
+ * Query execution options for instrumentation
+ */
+export interface QueryOptions {
+  /** The SQL statement */
+  sql: string
+  /** Query parameters */
+  params?: unknown[]
+  /** Database name */
+  database?: string
+  /** Whether result is from cache */
+  cached?: boolean
+  /** Parent span context */
+  parentContext?: SpanContext
+}
+
+/**
+ * Query result for instrumentation
+ */
+export interface QueryResult<T = unknown> {
+  /** Result rows */
+  rows: T[]
+  /** Number of rows returned */
+  rowCount: number
+  /** Number of rows affected (for mutations) */
+  affectedRows?: number
+  /** Execution duration in milliseconds */
+  durationMs: number
+}
+
+/**
+ * Instrumented query executor
+ *
+ * Wraps a query execution function with tracing and metrics.
+ *
+ * @example
+ * ```typescript
+ * const instrumentedQuery = createQueryInstrumentation(tracer, metrics)
+ *
+ * const result = await instrumentedQuery(
+ *   { sql: 'SELECT * FROM users WHERE id = $1', params: [1] },
+ *   async (sql, params) => pglite.query(sql, params)
+ * )
+ * ```
+ */
+export function createQueryInstrumentation(
+  tracer: Tracer,
+  metrics: MetricsCollector
+) {
+  return async function instrumentedQuery<T>(
+    options: QueryOptions,
+    executor: (sql: string, params?: unknown[]) => Promise<QueryResult<T>>
+  ): Promise<QueryResult<T>> {
+    const spanOptions: SpanOptions = {
+      kind: SpanKind.CLIENT,
+      attributes: {
+        'db.system': 'postgresql',
+        'db.statement': options.sql,
+        'db.name': options.database,
+      },
+      parent: options.parentContext,
+    }
+
+    return tracer.startActiveSpan('postgres.query', spanOptions, async (span) => {
+      const startTime = performance.now()
+      let result: QueryResult<T>
+
+      try {
+        // Detect operation type from SQL
+        const operation = detectSQLOperation(options.sql)
+        span.setAttribute('db.operation', operation)
+
+        result = await executor(options.sql, options.params)
+
+        const durationMs = performance.now() - startTime
+
+        // Set result attributes
+        const attrs: QuerySpanAttributes = {
+          'db.system': 'postgresql',
+          'db.operation': operation,
+          'db.duration_ms': durationMs,
+          'db.rows_returned': result.rowCount,
+          'db.cached': options.cached,
+        }
+
+        if (result.affectedRows !== undefined) {
+          attrs['db.rows_affected'] = result.affectedRows
+        }
+
+        span.setAttributes(attrs)
+
+        // Record metrics
+        metrics.increment(MetricNames.QUERY_COUNT, 1, { operation })
+        metrics.histogram(MetricNames.QUERY_DURATION, durationMs, { operation })
+        metrics.histogram(MetricNames.QUERY_ROWS_RETURNED, result.rowCount, {
+          operation,
+        })
+
+        if (result.affectedRows !== undefined) {
+          metrics.histogram(MetricNames.QUERY_ROWS_AFFECTED, result.affectedRows, {
+            operation,
+          })
+        }
+
+        span.setStatus({ code: SpanStatusCode.OK })
+        return result
+      } catch (error) {
+        const durationMs = performance.now() - startTime
+
+        // Record error
+        span.recordException(error as Error)
+        span.setAttribute('db.duration_ms', durationMs)
+
+        // Record error metrics
+        metrics.increment(MetricNames.QUERY_ERRORS, 1, {
+          operation: detectSQLOperation(options.sql),
+        })
+
+        throw error
+      }
+    })
+  }
+}
+
+/**
+ * Detect the SQL operation type from a statement
+ */
+function detectSQLOperation(sql: string): string {
+  const normalized = sql.trim().toUpperCase()
+  if (normalized.startsWith('SELECT')) return 'SELECT'
+  if (normalized.startsWith('INSERT')) return 'INSERT'
+  if (normalized.startsWith('UPDATE')) return 'UPDATE'
+  if (normalized.startsWith('DELETE')) return 'DELETE'
+  if (normalized.startsWith('CREATE')) return 'CREATE'
+  if (normalized.startsWith('ALTER')) return 'ALTER'
+  if (normalized.startsWith('DROP')) return 'DROP'
+  if (normalized.startsWith('BEGIN')) return 'BEGIN'
+  if (normalized.startsWith('COMMIT')) return 'COMMIT'
+  if (normalized.startsWith('ROLLBACK')) return 'ROLLBACK'
+  return 'OTHER'
+}
+
+// ============================================================================
+// Cache Instrumentation
+// ============================================================================
+
+/**
+ * Cache operation type
+ */
+export type CacheOperation = 'get' | 'put' | 'delete' | 'has'
+
+/**
+ * Cache tier
+ */
+export type CacheTier = 'cache' | 'do' | 'r2'
+
+/**
+ * Cache operation result
+ */
+export interface CacheOperationResult {
+  /** Whether the operation was a hit (for get/has) */
+  hit?: boolean
+  /** Bytes read or written */
+  bytes?: number
+  /** Duration in milliseconds */
+  durationMs: number
+}
+
+/**
+ * Instrumented cache operations
+ *
+ * Wraps cache operations with tracing and metrics.
+ *
+ * @example
+ * ```typescript
+ * const instrumentedCache = createCacheInstrumentation(tracer, metrics)
+ *
+ * const data = await instrumentedCache.get('page-123', 'cache', async () => {
+ *   return cacheLayer.get('page-123')
+ * })
+ * ```
+ */
+export function createCacheInstrumentation(
+  tracer: Tracer,
+  metrics: MetricsCollector
+) {
+  return {
+    async get<T>(
+      key: string,
+      tier: CacheTier,
+      executor: () => Promise<T | null>,
+      parentContext?: SpanContext
+    ): Promise<{ value: T | null; hit: boolean; durationMs: number }> {
+      const spanOptions: SpanOptions = {
+        kind: SpanKind.CLIENT,
+        attributes: {
+          'cache.operation': 'get',
+          'cache.key': key,
+          'cache.tier': tier,
+        },
+        parent: parentContext,
+      }
+
+      return tracer.startActiveSpan(`cache.${tier}.get`, spanOptions, async (span) => {
+        const startTime = performance.now()
+
+        try {
+          const value = await executor()
+          const durationMs = performance.now() - startTime
+          const hit = value !== null
+
+          const attrs: CacheSpanAttributes = {
+            'cache.operation': 'get',
+            'cache.hit': hit,
+            'cache.key': key,
+            'cache.tier': tier,
+            'cache.duration_ms': durationMs,
+          }
+          span.setAttributes(attrs)
+
+          // Record metrics
+          if (hit) {
+            metrics.increment(MetricNames.CACHE_HITS, 1, { tier })
+          } else {
+            metrics.increment(MetricNames.CACHE_MISSES, 1, { tier })
+          }
+          metrics.histogram(MetricNames.CACHE_LATENCY, durationMs, {
+            tier,
+            operation: 'get',
+          })
+
+          span.setStatus({ code: SpanStatusCode.OK })
+          return { value, hit, durationMs }
+        } catch (error) {
+          span.recordException(error as Error)
+          throw error
+        }
+      })
+    },
+
+    async put(
+      key: string,
+      tier: CacheTier,
+      bytes: number,
+      executor: () => Promise<void>,
+      parentContext?: SpanContext
+    ): Promise<{ durationMs: number }> {
+      const spanOptions: SpanOptions = {
+        kind: SpanKind.CLIENT,
+        attributes: {
+          'cache.operation': 'put',
+          'cache.key': key,
+          'cache.tier': tier,
+          'cache.bytes': bytes,
+        },
+        parent: parentContext,
+      }
+
+      return tracer.startActiveSpan(`cache.${tier}.put`, spanOptions, async (span) => {
+        const startTime = performance.now()
+
+        try {
+          await executor()
+          const durationMs = performance.now() - startTime
+
+          span.setAttribute('cache.duration_ms', durationMs)
+
+          // Record metrics
+          metrics.histogram(MetricNames.CACHE_BYTES_WRITTEN, bytes, { tier })
+          metrics.histogram(MetricNames.CACHE_LATENCY, durationMs, {
+            tier,
+            operation: 'put',
+          })
+
+          span.setStatus({ code: SpanStatusCode.OK })
+          return { durationMs }
+        } catch (error) {
+          span.recordException(error as Error)
+          throw error
+        }
+      })
+    },
+
+    async delete(
+      key: string,
+      tier: CacheTier,
+      executor: () => Promise<boolean>,
+      parentContext?: SpanContext
+    ): Promise<{ deleted: boolean; durationMs: number }> {
+      const spanOptions: SpanOptions = {
+        kind: SpanKind.CLIENT,
+        attributes: {
+          'cache.operation': 'delete',
+          'cache.key': key,
+          'cache.tier': tier,
+        },
+        parent: parentContext,
+      }
+
+      return tracer.startActiveSpan(`cache.${tier}.delete`, spanOptions, async (span) => {
+        const startTime = performance.now()
+
+        try {
+          const deleted = await executor()
+          const durationMs = performance.now() - startTime
+
+          span.setAttribute('cache.duration_ms', durationMs)
+
+          metrics.histogram(MetricNames.CACHE_LATENCY, durationMs, {
+            tier,
+            operation: 'delete',
+          })
+
+          span.setStatus({ code: SpanStatusCode.OK })
+          return { deleted, durationMs }
+        } catch (error) {
+          span.recordException(error as Error)
+          throw error
+        }
+      })
+    },
+  }
+}
+
+// ============================================================================
+// Durable Object Instrumentation
+// ============================================================================
+
+/**
+ * DO lifecycle operation
+ */
+export type DOOperation = 'activation' | 'hibernation' | 'fetch' | 'alarm'
+
+/**
+ * DO lifecycle options
+ */
+export interface DOLifecycleOptions {
+  /** DO ID */
+  id: string
+  /** DO name (optional) */
+  name?: string
+  /** Storage used in bytes (optional) */
+  storageUsed?: number
+  /** Parent span context */
+  parentContext?: SpanContext
+}
+
+/**
+ * Instrumented Durable Object lifecycle tracking
+ *
+ * @example
+ * ```typescript
+ * const doInstrumentation = createDOInstrumentation(tracer, metrics)
+ *
+ * await doInstrumentation.trackActivation({ id: doId.toString() }, async () => {
+ *   await this.initPGlite()
+ * })
+ * ```
+ */
+export function createDOInstrumentation(
+  tracer: Tracer,
+  metrics: MetricsCollector
+) {
+  return {
+    /**
+     * Track DO activation
+     */
+    async trackActivation<T>(
+      options: DOLifecycleOptions,
+      executor: () => Promise<T>
+    ): Promise<T> {
+      const spanOptions: SpanOptions = {
+        kind: SpanKind.SERVER,
+        attributes: {
+          'do.operation': 'activation',
+          'do.id': options.id,
+          'do.name': options.name,
+        },
+        parent: options.parentContext,
+      }
+
+      return tracer.startActiveSpan('do.activation', spanOptions, async (span) => {
+        const startTime = performance.now()
+
+        try {
+          metrics.increment(MetricNames.DO_ACTIVATIONS, 1)
+
+          const result = await executor()
+          const durationMs = performance.now() - startTime
+
+          const attrs: DOLifecycleSpanAttributes = {
+            'do.operation': 'activation',
+            'do.id': options.id,
+            'do.duration_ms': durationMs,
+          }
+
+          if (options.storageUsed !== undefined) {
+            attrs['do.storage_used'] = options.storageUsed
+          }
+
+          span.setAttributes(attrs)
+          metrics.histogram(MetricNames.DO_LATENCY, durationMs, {
+            operation: 'activation',
+          })
+
+          span.setStatus({ code: SpanStatusCode.OK })
+          return result
+        } catch (error) {
+          span.recordException(error as Error)
+          metrics.increment(MetricNames.ERROR_COUNT, 1, { type: 'do_activation' })
+          throw error
+        }
+      })
+    },
+
+    /**
+     * Track DO hibernation
+     */
+    async trackHibernation<T>(
+      options: DOLifecycleOptions,
+      executor: () => Promise<T>
+    ): Promise<T> {
+      const spanOptions: SpanOptions = {
+        kind: SpanKind.SERVER,
+        attributes: {
+          'do.operation': 'hibernation',
+          'do.id': options.id,
+          'do.name': options.name,
+        },
+        parent: options.parentContext,
+      }
+
+      return tracer.startActiveSpan('do.hibernation', spanOptions, async (span) => {
+        const startTime = performance.now()
+
+        try {
+          metrics.increment(MetricNames.DO_HIBERNATIONS, 1)
+
+          const result = await executor()
+          const durationMs = performance.now() - startTime
+
+          span.setAttribute('do.duration_ms', durationMs)
+          metrics.histogram(MetricNames.DO_LATENCY, durationMs, {
+            operation: 'hibernation',
+          })
+
+          span.setStatus({ code: SpanStatusCode.OK })
+          return result
+        } catch (error) {
+          span.recordException(error as Error)
+          metrics.increment(MetricNames.ERROR_COUNT, 1, { type: 'do_hibernation' })
+          throw error
+        }
+      })
+    },
+
+    /**
+     * Track DO fetch request
+     */
+    async trackFetch<T>(
+      options: DOLifecycleOptions,
+      executor: () => Promise<T>
+    ): Promise<T> {
+      const spanOptions: SpanOptions = {
+        kind: SpanKind.SERVER,
+        attributes: {
+          'do.operation': 'fetch',
+          'do.id': options.id,
+          'do.name': options.name,
+        },
+        parent: options.parentContext,
+      }
+
+      return tracer.startActiveSpan('do.fetch', spanOptions, async (span) => {
+        const startTime = performance.now()
+
+        try {
+          metrics.increment(MetricNames.DO_REQUESTS, 1)
+
+          const result = await executor()
+          const durationMs = performance.now() - startTime
+
+          span.setAttribute('do.duration_ms', durationMs)
+          metrics.histogram(MetricNames.DO_LATENCY, durationMs, {
+            operation: 'fetch',
+          })
+
+          span.setStatus({ code: SpanStatusCode.OK })
+          return result
+        } catch (error) {
+          span.recordException(error as Error)
+          metrics.increment(MetricNames.ERROR_COUNT, 1, { type: 'do_fetch' })
+          throw error
+        }
+      })
+    },
+
+    /**
+     * Track DO alarm
+     */
+    async trackAlarm<T>(
+      options: DOLifecycleOptions,
+      executor: () => Promise<T>
+    ): Promise<T> {
+      const spanOptions: SpanOptions = {
+        kind: SpanKind.SERVER,
+        attributes: {
+          'do.operation': 'alarm',
+          'do.id': options.id,
+          'do.name': options.name,
+        },
+        parent: options.parentContext,
+      }
+
+      return tracer.startActiveSpan('do.alarm', spanOptions, async (span) => {
+        const startTime = performance.now()
+
+        try {
+          const result = await executor()
+          const durationMs = performance.now() - startTime
+
+          span.setAttribute('do.duration_ms', durationMs)
+          metrics.histogram(MetricNames.DO_LATENCY, durationMs, {
+            operation: 'alarm',
+          })
+
+          span.setStatus({ code: SpanStatusCode.OK })
+          return result
+        } catch (error) {
+          span.recordException(error as Error)
+          metrics.increment(MetricNames.ERROR_COUNT, 1, { type: 'do_alarm' })
+          throw error
+        }
+      })
+    },
+
+    /**
+     * Update storage usage gauge
+     */
+    updateStorageUsage(bytes: number): void {
+      metrics.gauge(MetricNames.DO_STORAGE_USED, bytes)
+    },
+  }
+}
+
+// ============================================================================
+// Error Tracking
+// ============================================================================
+
+/**
+ * Error category for classification
+ */
+export type ErrorCategory =
+  | 'query'
+  | 'connection'
+  | 'cache'
+  | 'storage'
+  | 'timeout'
+  | 'validation'
+  | 'unknown'
+
+/**
+ * Error tracking options
+ */
+export interface ErrorTrackingOptions {
+  /** Error category */
+  category: ErrorCategory
+  /** Whether the error is recoverable */
+  recoverable?: boolean
+  /** Error code (e.g., SQL error code) */
+  code?: string
+  /** Additional context */
+  context?: Record<string, unknown>
+  /** Parent span context */
+  parentContext?: SpanContext
+}
+
+/**
+ * Instrumented error tracking
+ *
+ * @example
+ * ```typescript
+ * const errorTracker = createErrorInstrumentation(tracer, metrics)
+ *
+ * try {
+ *   await someOperation()
+ * } catch (error) {
+ *   errorTracker.recordError(error, { category: 'query' })
+ *   throw error
+ * }
+ * ```
+ */
+export function createErrorInstrumentation(
+  tracer: Tracer,
+  metrics: MetricsCollector
+) {
+  return {
+    /**
+     * Record an error with tracing and metrics
+     */
+    recordError(error: Error | string, options: ErrorTrackingOptions): void {
+      const errorMessage = typeof error === 'string' ? error : error.message
+      const errorType = typeof error === 'string' ? 'Error' : error.name
+      const errorStack = typeof error === 'string' ? undefined : error.stack
+
+      // Create an error span
+      const span = tracer.startSpan('error', {
+        kind: SpanKind.INTERNAL,
+        attributes: {
+          'error.type': errorType,
+          'error.message': errorMessage,
+          'error.stack': errorStack,
+          'error.code': options.code,
+          'error.category': options.category,
+          'error.recoverable': options.recoverable ?? false,
+        } satisfies ErrorSpanAttributes,
+        parent: options.parentContext,
+      })
+
+      span.recordException(error)
+      span.setStatus({
+        code: SpanStatusCode.ERROR,
+        message: errorMessage,
+      })
+      span.end()
+
+      // Record metrics
+      metrics.increment(MetricNames.ERROR_COUNT, 1, {
+        category: options.category,
+        type: errorType,
+        recoverable: String(options.recoverable ?? false),
+      })
+    },
+
+    /**
+     * Wrap an async function with error tracking
+     */
+    async withErrorTracking<T>(
+      category: ErrorCategory,
+      executor: () => Promise<T>,
+      parentContext?: SpanContext
+    ): Promise<T> {
+      try {
+        return await executor()
+      } catch (error) {
+        const options: ErrorTrackingOptions = { category }
+        if (parentContext !== undefined) {
+          options.parentContext = parentContext
+        }
+        this.recordError(error as Error, options)
+        throw error
+      }
+    },
+  }
+}
+
+// ============================================================================
+// Request Instrumentation (Fetch handler)
+// ============================================================================
+
+/**
+ * Options for request instrumentation
+ */
+export interface RequestInstrumentationOptions {
+  /** Service name */
+  serviceName: string
+  /** Extract user-defined attributes from request */
+  extractAttributes?: (request: Request) => Record<string, string>
+}
+
+/**
+ * Create an instrumented fetch handler wrapper
+ *
+ * @example
+ * ```typescript
+ * const instrumentedFetch = createRequestInstrumentation(tracer, metrics, {
+ *   serviceName: 'postgres-do',
+ * })
+ *
+ * export default {
+ *   fetch: instrumentedFetch(async (request, env, ctx) => {
+ *     // Your handler logic
+ *   }),
+ * }
+ * ```
+ */
+export function createRequestInstrumentation(
+  tracer: Tracer,
+  metrics: MetricsCollector,
+  options: RequestInstrumentationOptions
+) {
+  return function instrumentedFetch<Env>(
+    handler: (
+      request: Request,
+      env: Env,
+      ctx: ExecutionContext
+    ) => Promise<Response>
+  ): (request: Request, env: Env, ctx: ExecutionContext) => Promise<Response> {
+    return async (request, env, ctx) => {
+      // Extract trace context from incoming request
+      const parentContext = extractSpanContext(request.headers)
+
+      const url = new URL(request.url)
+      const spanOptions: SpanOptions = {
+        kind: SpanKind.SERVER,
+        attributes: {
+          'http.method': request.method,
+          'http.url': url.pathname,
+          'http.host': url.host,
+          'http.scheme': url.protocol.replace(':', ''),
+          'service.name': options.serviceName,
+          ...options.extractAttributes?.(request),
+        },
+        parent: parentContext ?? undefined,
+      }
+
+      return tracer.startActiveSpan(
+        `${request.method} ${url.pathname}`,
+        spanOptions,
+        async (span) => {
+          const startTime = performance.now()
+
+          try {
+            const response = await handler(request, env, ctx)
+            const durationMs = performance.now() - startTime
+
+            span.setAttributes({
+              'http.status_code': response.status,
+              'http.duration_ms': durationMs,
+            })
+
+            // Inject trace context into response headers
+            const responseHeaders = new Headers(response.headers)
+            injectSpanContext(span.getContext(), responseHeaders)
+
+            span.setStatus({ code: SpanStatusCode.OK })
+            return new Response(response.body, {
+              status: response.status,
+              statusText: response.statusText,
+              headers: responseHeaders,
+            })
+          } catch (error) {
+            const durationMs = performance.now() - startTime
+
+            span.recordException(error as Error)
+            span.setAttribute('http.duration_ms', durationMs)
+
+            metrics.increment(MetricNames.ERROR_COUNT, 1, {
+              category: 'request',
+              type: (error as Error).name,
+            })
+
+            throw error
+          }
+        }
+      )
+    }
+  }
+}