elasticdash-sdk 0.2.0

package/src/interceptors/db-auto.ts

@@ -0,0 +1,243 @@
+import { getCaptureContext } from '../capture/recorder.js'
+import { getCurrentTrace } from '../trace-adapter/context.js'
+import { rawDateNow } from './side-effects.js'
+import { getHttpRunContext, getHttpFrozenEvent, pushTelemetryEvent, tryAutoInitHttpContext, getObservabilityContext } from './telemetry-push.js'
+
+type AnyFn = (...args: unknown[]) => unknown
+
+interface MethodPatch {
+  proto: Record<string, unknown>
+  method: string
+  original: AnyFn
+}
+
+const appliedPatches: MethodPatch[] = []
+
+function toTraceArgs(input: unknown): Record<string, unknown> | undefined {
+  if (input && typeof input === 'object' && !Array.isArray(input)) {
+    return input as Record<string, unknown>
+  }
+  if (input === undefined) return undefined
+  return { value: input }
+}
+
+function wrapProtoMethod(proto: object, method: string, eventName: string): void {
+  const p = proto as Record<string, unknown>
+  if (typeof p[method] !== 'function') return
+
+  const original = p[method] as AnyFn
+  appliedPatches.push({ proto: p, method, original })
+
+  p[method] = function (this: unknown, ...args: unknown[]) {
+    // Skip callback-style calls to avoid breaking legacy APIs
+    if (args.length > 0 && typeof args[args.length - 1] === 'function') {
+      return original.apply(this, args)
+    }
+
+    const ctx = getCaptureContext()
+    const httpCtx = getHttpRunContext()
+    const obsCtx = getObservabilityContext()
+    const input = args.length === 1 ? args[0] : args
+
+    if (!ctx && !httpCtx && !obsCtx) return original.apply(this, args)
+
+    // Observability-only mode: record and push, no mocks/replay
+    if (!ctx && !httpCtx && obsCtx) {
+      const id = obsCtx.nextId()
+      const start = rawDateNow()
+
+      let result: unknown
+      try {
+        result = original.apply(this, args)
+      } catch (err) {
+        pushTelemetryEvent({ id, type: 'db', name: eventName, input, output: { error: String(err) }, timestamp: start, durationMs: rawDateNow() - start })
+        throw err
+      }
+
+      if (result != null && typeof (result as Promise<unknown>).then === 'function') {
+        return (result as Promise<unknown>)
+          .then((output: unknown) => {
+            pushTelemetryEvent({ id, type: 'db', name: eventName, input, output, timestamp: start, durationMs: rawDateNow() - start })
+            return output
+          })
+          .catch((err: unknown) => {
+            pushTelemetryEvent({ id, type: 'db', name: eventName, input, output: { error: String(err) }, timestamp: start, durationMs: rawDateNow() - start })
+            throw err
+          })
+      }
+
+      pushTelemetryEvent({ id, type: 'db', name: eventName, input, output: result, timestamp: start, durationMs: rawDateNow() - start })
+      return result
+    }
+
+    // HTTP mode (no capture context): replay frozen events or execute live
+    if (!ctx && httpCtx) {
+      const id = httpCtx.nextId()
+
+      // Replay frozen step
+      const frozen = getHttpFrozenEvent(id)
+      if (frozen && frozen.type === 'db') {
+        pushTelemetryEvent(frozen)
+        return Promise.resolve(frozen.output)
+      }
+
+      // Not frozen → execute live, push telemetry
+      const start = rawDateNow()
+
+      let result: unknown
+      try {
+        result = original.apply(this, args)
+      } catch (err) {
+        pushTelemetryEvent({ id, type: 'db', name: eventName, input, output: { error: String(err) }, timestamp: start, durationMs: rawDateNow() - start })
+        throw err
+      }
+
+      if (result != null && typeof (result as Promise<unknown>).then === 'function') {
+        return (result as Promise<unknown>)
+          .then((output: unknown) => {
+            pushTelemetryEvent({ id, type: 'db', name: eventName, input, output, timestamp: start, durationMs: rawDateNow() - start })
+            return output
+          })
+          .catch((err: unknown) => {
+            pushTelemetryEvent({ id, type: 'db', name: eventName, input, output: { error: String(err) }, timestamp: start, durationMs: rawDateNow() - start })
+            throw err
+          })
+      }
+
+      pushTelemetryEvent({ id, type: 'db', name: eventName, input, output: result, timestamp: start, durationMs: rawDateNow() - start })
+      return result
+    }
+
+    // Capture mode (enhanced with telemetry + TraceHandle)
+    const trace = getCurrentTrace()
+    const { recorder, replay } = ctx!
+    const id = recorder.nextId()
+
+    if (replay.shouldReplay(id)) {
+      const historicalEvent = replay.getRecordedEvent(id)
+      if (historicalEvent) recorder.record(historicalEvent)
+      if (httpCtx) pushTelemetryEvent(historicalEvent ?? { id, type: 'db', name: eventName, input, output: null, timestamp: rawDateNow(), durationMs: 0 })
+      const replayed = replay.getRecordedResult(id)
+      if (trace && typeof trace.recordToolCall === 'function') {
+        trace.recordToolCall({ name: eventName, args: toTraceArgs(input), result: replayed, workflowEventId: id })
+      }
+      return Promise.resolve(replayed)
+    }
+
+    const start = rawDateNow()
+
+    let result: unknown
+    try {
+      result = original.apply(this, args)
+    } catch (err) {
+      const durationMs = rawDateNow() - start
+      const event = { id, type: 'db' as const, name: eventName, input, output: { error: String(err) }, timestamp: start, durationMs }
+      recorder.record(event)
+      if (httpCtx) pushTelemetryEvent(event)
+      if (trace && typeof trace.recordToolCall === 'function') {
+        trace.recordToolCall({ name: eventName, args: toTraceArgs(input), result: { error: String(err) }, workflowEventId: id, durationMs })
+      }
+      throw err
+    }
+
+    if (result != null && typeof (result as Promise<unknown>).then === 'function') {
+      return (result as Promise<unknown>)
+        .then((output: unknown) => {
+          const durationMs = rawDateNow() - start
+          const event = { id, type: 'db' as const, name: eventName, input, output, timestamp: start, durationMs }
+          recorder.record(event)
+          if (httpCtx) pushTelemetryEvent(event)
+          if (trace && typeof trace.recordToolCall === 'function') {
+            trace.recordToolCall({ name: eventName, args: toTraceArgs(input), result: output, workflowEventId: id, durationMs })
+          }
+          return output
+        })
+        .catch((err: unknown) => {
+          const durationMs = rawDateNow() - start
+          const event = { id, type: 'db' as const, name: eventName, input, output: { error: String(err) }, timestamp: start, durationMs }
+          recorder.record(event)
+          if (httpCtx) pushTelemetryEvent(event)
+          if (trace && typeof trace.recordToolCall === 'function') {
+            trace.recordToolCall({ name: eventName, args: toTraceArgs(input), result: { error: String(err) }, workflowEventId: id, durationMs })
+          }
+          throw err
+        })
+    }
+
+    // Sync return (rare for DB calls)
+    const durationMs = rawDateNow() - start
+    const event = { id, type: 'db' as const, name: eventName, input, output: result, timestamp: start, durationMs }
+    recorder.record(event)
+    if (httpCtx) pushTelemetryEvent(event)
+    if (trace && typeof trace.recordToolCall === 'function') {
+      trace.recordToolCall({ name: eventName, args: toTraceArgs(input), result, workflowEventId: id, durationMs })
+    }
+    return result
+  }
+}
+
+async function tryPatchPg(): Promise<void> {
+  // @ts-ignore — optional peer dependency
+  const pgMod = await import('pg') as Record<string, unknown>
+  const pg = (pgMod.default as Record<string, unknown> | undefined) ?? pgMod
+  const Client = pg.Client as { prototype: object } | undefined
+  // Patch Client.prototype only — Pool.query delegates to Client internally
+  if (Client?.prototype) {
+    wrapProtoMethod(Client.prototype, 'query', 'pg.query')
+  }
+}
+
+async function tryPatchMysql2(): Promise<void> {
+  // @ts-ignore — optional peer dependency
+  const mod = await import('mysql2/promise') as Record<string, unknown>
+  const mysql2 = (mod.default as Record<string, unknown> | undefined) ?? mod
+  const Connection = mysql2.Connection as { prototype: object } | undefined
+  if (Connection?.prototype) {
+    wrapProtoMethod(Connection.prototype, 'query', 'mysql2.query')
+    wrapProtoMethod(Connection.prototype, 'execute', 'mysql2.execute')
+  }
+}
+
+async function tryPatchMongodb(): Promise<void> {
+  // @ts-ignore — optional peer dependency
+  const mongMod = await import('mongodb') as Record<string, unknown>
+  const Collection = (
+    mongMod.Collection ??
+    (mongMod.default as Record<string, unknown> | undefined)?.Collection
+  ) as { prototype: object } | undefined
+  if (Collection?.prototype) {
+    for (const method of ['find', 'findOne', 'insertOne', 'updateOne', 'deleteOne', 'aggregate']) {
+      wrapProtoMethod(Collection.prototype, method, `mongodb.${method}`)
+    }
+  }
+}
+
+async function tryPatchIoredis(): Promise<void> {
+  // @ts-ignore — optional peer dependency
+  const mod = await import('ioredis') as Record<string, unknown>
+  const Redis = (mod.default ?? mod) as { prototype: object } | undefined
+  if (Redis?.prototype) {
+    wrapProtoMethod(Redis.prototype, 'call', 'redis.call')
+  }
+}
+
+/**
+ * Auto-instruments common DB driver prototypes. Safe to call when drivers are
+ * not installed — missing modules are silently skipped.
+ */
+export async function installDBAutoInterceptor(): Promise<void> {
+  await Promise.allSettled([
+    tryPatchPg(),
+    tryPatchMysql2(),
+    tryPatchMongodb(),
+    tryPatchIoredis(),
+  ])
+}
+
+/** Restores all patched DB driver prototypes to their originals. */
+export function uninstallDBAutoInterceptor(): void {
+  for (const { proto, method, original } of appliedPatches) {
+    proto[method] = original
+  }
+  appliedPatches.length = 0
+}

package/src/interceptors/db.ts

@@ -0,0 +1,156 @@
+import { getCaptureContext } from '../capture/recorder.js'
+import { getCurrentTrace } from '../trace-adapter/context.js'
+import { rawDateNow } from './side-effects.js'
+import { getHttpRunContext, getHttpFrozenEvent, pushTelemetryEvent, tryAutoInitHttpContext, getObservabilityContext } from './telemetry-push.js'
+
+function toTraceArgs(input: unknown): Record<string, unknown> | undefined {
+  if (input && typeof input === 'object' && !Array.isArray(input)) {
+    return input as Record<string, unknown>
+  }
+  if (input === undefined) return undefined
+  return { value: input }
+}
+
+/**
+ * Wraps named methods on a DB client instance in-place so their calls are
+ * recorded as "db" events. Returns the same client object.
+ *
+ * Supports all three execution modes:
+ * - Capture mode: replay from TraceRecorder / ReplayController
+ * - HTTP mode: replay frozen events from dashboard, push telemetry
+ * - Observability mode: record and push telemetry
+ *
+ * @param client  Any DB client (pg.Client, redis client, mongoose Model, etc.)
+ * @param methodNames  Method names to wrap
+ * @param label  Optional label prefix for event names (defaults to constructor name)
+ */
+export function wrapDB<T extends object>(
+  client: T,
+  methodNames: (keyof T & string)[],
+  label?: string,
+): T {
+  const prefix = label ?? (client.constructor?.name ?? 'db')
+
+  for (const method of methodNames) {
+    const original = (client as Record<string, unknown>)[method]
+    if (typeof original !== 'function') continue
+
+    ;(client as Record<string, unknown>)[method] = async (...args: unknown[]) => {
+      await tryAutoInitHttpContext()
+      const ctx = getCaptureContext()
+      const httpCtx = getHttpRunContext()
+      const obsCtx = getObservabilityContext()
+      const name = `${prefix}.${method}`
+      const input = args.length === 1 ? args[0] : args
+
+      if (!ctx && !httpCtx && !obsCtx) return (original as (...a: unknown[]) => unknown).apply(client, args)
+
+      // Observability-only mode: record and push, no mocks/replay
+      if (!ctx && !httpCtx && obsCtx) {
+        const id = obsCtx.nextId()
+        const start = rawDateNow()
+        try {
+          const output = await (original as (...a: unknown[]) => unknown).apply(client, args)
+          pushTelemetryEvent({ id, type: 'db', name, input, output, timestamp: start, durationMs: rawDateNow() - start })
+          return output
+        } catch (e) {
+          pushTelemetryEvent({ id, type: 'db', name, input, output: { error: String(e) }, timestamp: start, durationMs: rawDateNow() - start })
+          throw e
+        }
+      }
+
+      // HTTP mode (no capture context): replay frozen events or execute live
+      if (!ctx && httpCtx) {
+        const id = httpCtx.nextId()
+
+        // Replay frozen step
+        const frozen = getHttpFrozenEvent(id)
+        if (frozen && frozen.type === 'db') {
+          pushTelemetryEvent(frozen)
+          return frozen.output
+        }
+
+        // Not frozen → execute live, push telemetry
+        const start = rawDateNow()
+        try {
+          const output = await (original as (...a: unknown[]) => unknown).apply(client, args)
+          pushTelemetryEvent({ id, type: 'db', name, input, output, timestamp: start, durationMs: rawDateNow() - start })
+          return output
+        } catch (e) {
+          pushTelemetryEvent({ id, type: 'db', name, input, output: { error: String(e) }, timestamp: start, durationMs: rawDateNow() - start })
+          throw e
+        }
+      }
+
+      // Capture mode
+      const trace = getCurrentTrace()
+      const { recorder, replay } = ctx!
+      const id = recorder.nextId()
+
+      if (replay.shouldReplay(id)) {
+        const historical = replay.getRecordedEvent(id)
+        if (historical) recorder.record(historical)
+        if (httpCtx) pushTelemetryEvent(historical ?? { id, type: 'db', name, input, output: null, timestamp: rawDateNow(), durationMs: 0 })
+        const replayed = replay.getRecordedResult(id)
+        if (trace && typeof trace.recordToolCall === 'function') {
+          trace.recordToolCall({ name, args: toTraceArgs(input), result: replayed, workflowEventId: id })
+        }
+        return replayed
+      }
+
+      const start = rawDateNow()
+      try {
+        const output = await (original as (...a: unknown[]) => unknown).apply(client, args)
+        const durationMs = rawDateNow() - start
+        const event = { id, type: 'db' as const, name, input, output, timestamp: start, durationMs }
+        recorder.record(event)
+        if (httpCtx) pushTelemetryEvent(event)
+        if (trace && typeof trace.recordToolCall === 'function') {
+          trace.recordToolCall({ name, args: toTraceArgs(input), result: output, workflowEventId: id, durationMs })
+        }
+        return output
+      } catch (e) {
+        const durationMs = rawDateNow() - start
+        const event = { id, type: 'db' as const, name, input, output: { error: String(e) }, timestamp: start, durationMs }
+        recorder.record(event)
+        if (httpCtx) pushTelemetryEvent(event)
+        if (trace && typeof trace.recordToolCall === 'function') {
+          trace.recordToolCall({ name, args: toTraceArgs(input), result: { error: String(e) }, workflowEventId: id, durationMs })
+        }
+        throw e
+      }
+    }
+  }
+
+  return client
+}
+
+// --- Driver-specific convenience helpers ---
+
+/** Wraps `query` on a pg.Client or pg.Pool instance. */
+export function wrapPgClient<T extends object>(client: T): T {
+  return wrapDB(client, ['query'] as (keyof T & string)[], 'pg')
+}
+
+/** Wraps `raw` on a knex instance. */
+export function wrapKnex<T extends object>(knex: T): T {
+  return wrapDB(knex, ['raw'] as (keyof T & string)[], 'knex')
+}
+
+/** Wraps common query methods on a MongoDB collection. */
+export function wrapMongoCollection<T extends object>(collection: T): T {
+  return wrapDB(
+    collection,
+    ['find', 'findOne', 'insertOne', 'updateOne', 'deleteOne'] as (keyof T & string)[],
+    'mongo',
+  )
+}
+
+/** Wraps common methods on a Redis client. */
+export function wrapRedisClient<T extends object>(client: T): T {
+  return wrapDB(
+    client,
+    ['get', 'set', 'del', 'hget', 'hset'] as (keyof T & string)[],
+    'redis',
+  )
+}