autotel-drizzle 0.0.1

package/src/drizzle/index.ts

@@ -0,0 +1,533 @@
+/* eslint-disable @typescript-eslint/no-explicit-any */
+import { SpanKind, trace } from '@opentelemetry/api';
+import {
+  SEMATTRS_DB_NAME,
+  SEMATTRS_DB_OPERATION,
+  SEMATTRS_DB_STATEMENT,
+  SEMATTRS_DB_SYSTEM,
+  SEMATTRS_NET_PEER_NAME,
+  SEMATTRS_NET_PEER_PORT,
+} from '../common/constants';
+import { finalizeSpan, runWithSpan } from 'autotel/trace-helpers';
+
+const DEFAULT_TRACER_NAME = 'autotel-plugins/drizzle';
+const DEFAULT_DB_SYSTEM = 'postgresql';
+const INSTRUMENTED_FLAG = '__autotelDrizzleInstrumented' as const;
+const PREPARED_QUERY_METHODS = [
+  'all',
+  'execute',
+  'get',
+  'run',
+  'values',
+] as const;
+
+type QueryCallback = (error: unknown, result: unknown) => void;
+type QueryFunction = (...args: any[]) => any;
+type AttributeValue = string | number | boolean;
+type AttributeMap = Record<string, AttributeValue>;
+
+interface InstrumentableObject {
+  [key: string]: any;
+  [INSTRUMENTED_FLAG]?: true;
+}
+
+interface DrizzleClientLike extends InstrumentableObject {
+  query?: QueryFunction;
+  execute?: QueryFunction;
+}
+
+interface DrizzleSessionLike extends InstrumentableObject {
+  query?: QueryFunction;
+  execute?: QueryFunction;
+  prepareQuery?: QueryFunction;
+  transaction?: QueryFunction;
+}
+
+interface DrizzleDbLike extends InstrumentableObject {
+  $client?: DrizzleClientLike;
+  session?: DrizzleSessionLike;
+  _?: {
+    session?: DrizzleSessionLike;
+    [key: string]: any;
+  };
+}
+
+export interface InstrumentDrizzleConfig {
+  tracerName?: string;
+  dbSystem?: string;
+  dbName?: string;
+  captureQueryText?: boolean;
+  maxQueryTextLength?: number;
+  peerName?: string;
+  peerPort?: number;
+}
+
+interface ResolvedConfig {
+  tracerName: string;
+  dbSystem: string;
+  dbName?: string;
+  captureQueryText: boolean;
+  maxQueryTextLength: number;
+  peerName?: string;
+  peerPort?: number;
+}
+
+interface InstrumentationState {
+  tracer: ReturnType<typeof trace.getTracer>;
+  config: ResolvedConfig;
+}
+
+interface MethodInstrumentationOptions {
+  flagSuffix: string;
+  queryText: (args: any[]) => string | undefined;
+  callbackStyle?: 'last-arg';
+  extraAttributes?: AttributeMap;
+}
+
+function resolveConfig(config?: InstrumentDrizzleConfig): ResolvedConfig {
+  return {
+    tracerName: config?.tracerName ?? DEFAULT_TRACER_NAME,
+    dbSystem: config?.dbSystem ?? DEFAULT_DB_SYSTEM,
+    dbName: config?.dbName,
+    captureQueryText: config?.captureQueryText ?? true,
+    maxQueryTextLength: config?.maxQueryTextLength ?? 1000,
+    peerName: config?.peerName,
+    peerPort: config?.peerPort,
+  };
+}
+
+function getState(config?: InstrumentDrizzleConfig): InstrumentationState {
+  const resolved = resolveConfig(config);
+  return {
+    config: resolved,
+    tracer: trace.getTracer(resolved.tracerName),
+  };
+}
+
+function getFlagKey(suffix: string): string {
+  return `${INSTRUMENTED_FLAG}:${suffix}`;
+}
+
+function isObject(value: unknown): value is InstrumentableObject {
+  return value !== null && typeof value === 'object';
+}
+
+function isPromiseLike<T>(value: T | PromiseLike<T>): value is PromiseLike<T> {
+  return (
+    typeof value === 'object' &&
+    value !== null &&
+    typeof (value as PromiseLike<T>).then === 'function'
+  );
+}
+
+function extractQueryText(queryArg: unknown): string | undefined {
+  if (typeof queryArg === 'string') {
+    return queryArg;
+  }
+
+  if (!isObject(queryArg)) {
+    return undefined;
+  }
+
+  if (typeof queryArg.sql === 'string') {
+    return queryArg.sql;
+  }
+
+  if (typeof queryArg.text === 'string') {
+    return queryArg.text;
+  }
+
+  if (typeof queryArg.queryString === 'string') {
+    return queryArg.queryString;
+  }
+
+  if (
+    isObject(queryArg.queryChunks) &&
+    typeof (queryArg as Record<string, unknown>).sql === 'string'
+  ) {
+    return queryArg.sql as string;
+  }
+
+  return undefined;
+}
+
+function sanitizeQueryText(queryText: string, maxLength: number): string {
+  if (queryText.length <= maxLength) {
+    return queryText;
+  }
+
+  return `${queryText.slice(0, Math.max(0, maxLength))}...`;
+}
+
+function extractOperation(queryText: string): string | undefined {
+  const trimmed = queryText.trimStart();
+  const match = /^(?<operation>\w+)/u.exec(trimmed);
+  return match?.groups?.operation?.toUpperCase();
+}
+
+function buildSpan(
+  state: InstrumentationState,
+  queryText: string | undefined,
+  extraAttributes?: AttributeMap,
+) {
+  const operation = queryText ? extractOperation(queryText) : undefined;
+  const spanName = operation
+    ? `drizzle.${operation.toLowerCase()}`
+    : 'drizzle.query';
+  const span = state.tracer.startSpan(spanName, { kind: SpanKind.CLIENT });
+
+  span.setAttribute(SEMATTRS_DB_SYSTEM, state.config.dbSystem);
+
+  if (operation) {
+    span.setAttribute(SEMATTRS_DB_OPERATION, operation);
+  }
+
+  if (state.config.dbName !== undefined) {
+    span.setAttribute(SEMATTRS_DB_NAME, state.config.dbName);
+  }
+
+  if (state.config.captureQueryText && queryText !== undefined) {
+    span.setAttribute(
+      SEMATTRS_DB_STATEMENT,
+      sanitizeQueryText(queryText, state.config.maxQueryTextLength),
+    );
+  }
+
+  if (state.config.peerName !== undefined) {
+    span.setAttribute(SEMATTRS_NET_PEER_NAME, state.config.peerName);
+  }
+
+  if (state.config.peerPort !== undefined) {
+    span.setAttribute(SEMATTRS_NET_PEER_PORT, state.config.peerPort);
+  }
+
+  if (extraAttributes) {
+    for (const [key, value] of Object.entries(extraAttributes)) {
+      span.setAttribute(key, value);
+    }
+  }
+
+  return span;
+}
+
+function executeWithSpan<T>(span: any, fn: () => T): T {
+  return runWithSpan(span, () => {
+    try {
+      const result = fn();
+
+      if (isPromiseLike(result)) {
+        return result.then(
+          (value) => {
+            finalizeSpan(span);
+            return value;
+          },
+          (error) => {
+            finalizeSpan(span, error);
+            throw error;
+          },
+        ) as T;
+      }
+
+      finalizeSpan(span);
+      return result;
+    } catch (error) {
+      finalizeSpan(span, error);
+      throw error;
+    }
+  });
+}
+
+function instrumentMethod(
+  target: InstrumentableObject,
+  methodName: string,
+  state: InstrumentationState,
+  options: MethodInstrumentationOptions,
+): boolean {
+  if (typeof target[methodName] !== 'function') {
+    return false;
+  }
+
+  const flagKey = getFlagKey(options.flagSuffix);
+  if (target[flagKey]) {
+    return false;
+  }
+
+  const originalMethod = target[methodName] as QueryFunction;
+
+  target[methodName] = function instrumentedMethod(
+    this: any,
+    ...incomingArgs: any[]
+  ) {
+    const args = [...incomingArgs];
+    const callback =
+      options.callbackStyle === 'last-arg' && typeof args.at(-1) === 'function'
+        ? (args.pop() as QueryCallback)
+        : undefined;
+    const span = buildSpan(
+      state,
+      options.queryText(args),
+      options.extraAttributes,
+    );
+
+    if (callback) {
+      return runWithSpan(span, () => {
+        const wrappedCallback: QueryCallback = (error, result) => {
+          finalizeSpan(span, error);
+          callback(error, result);
+        };
+
+        try {
+          return Reflect.apply(originalMethod, this, [
+            ...args,
+            wrappedCallback,
+          ]);
+        } catch (error) {
+          finalizeSpan(span, error);
+          throw error;
+        }
+      });
+    }
+
+    return executeWithSpan(span, () =>
+      Reflect.apply(originalMethod, this, args),
+    );
+  };
+
+  target[flagKey] = true;
+  return true;
+}
+
+function instrumentPreparedQuery(
+  prepared: unknown,
+  state: InstrumentationState,
+  querySource: unknown,
+  extraAttributes?: AttributeMap,
+): boolean {
+  if (!isObject(prepared)) {
+    return false;
+  }
+
+  let instrumented = false;
+  const queryText = extractQueryText(querySource);
+
+  for (const methodName of PREPARED_QUERY_METHODS) {
+    instrumented =
+      instrumentMethod(prepared, methodName, state, {
+        flagSuffix: `prepared:${methodName}`,
+        queryText: () => queryText,
+        extraAttributes,
+      }) || instrumented;
+  }
+
+  return instrumented;
+}
+
+function instrumentPrepareQuery(
+  target: DrizzleSessionLike,
+  state: InstrumentationState,
+  extraAttributes?: AttributeMap,
+): boolean {
+  if (typeof target.prepareQuery !== 'function') {
+    return false;
+  }
+
+  const flagKey = getFlagKey('prepareQuery');
+  if (target[flagKey]) {
+    return false;
+  }
+
+  const originalPrepareQuery = target.prepareQuery;
+
+  target.prepareQuery = function instrumentedPrepareQuery(
+    this: any,
+    ...prepareArgs: any[]
+  ) {
+    const prepared = Reflect.apply(originalPrepareQuery, this, prepareArgs);
+    instrumentPreparedQuery(prepared, state, prepareArgs[0], extraAttributes);
+    return prepared;
+  };
+
+  target[flagKey] = true;
+  return true;
+}
+
+function instrumentTransactionTarget(
+  target: unknown,
+  state: InstrumentationState,
+): boolean {
+  if (!isObject(target)) {
+    return false;
+  }
+
+  const transactionAttributes = { 'db.transaction': true };
+  let instrumented = false;
+
+  instrumented =
+    instrumentMethod(target, 'query', state, {
+      flagSuffix: 'transaction:query',
+      queryText: (args) => extractQueryText(args[0]),
+      callbackStyle: 'last-arg',
+      extraAttributes: transactionAttributes,
+    }) || instrumented;
+
+  instrumented =
+    instrumentMethod(target, 'execute', state, {
+      flagSuffix: 'transaction:execute',
+      queryText: (args) => extractQueryText(args[0]),
+      callbackStyle: 'last-arg',
+      extraAttributes: transactionAttributes,
+    }) || instrumented;
+
+  instrumented =
+    instrumentPrepareQuery(
+      target as DrizzleSessionLike,
+      state,
+      transactionAttributes,
+    ) || instrumented;
+
+  if (isObject(target.session)) {
+    instrumented =
+      instrumentTransactionTarget(target.session, state) || instrumented;
+  }
+
+  if (isObject(target._?.session)) {
+    instrumented =
+      instrumentTransactionTarget(target._.session, state) || instrumented;
+  }
+
+  return instrumented;
+}
+
+function instrumentSession(
+  session: DrizzleSessionLike,
+  state: InstrumentationState,
+): boolean {
+  let instrumented = false;
+
+  instrumented =
+    instrumentMethod(session, 'query', state, {
+      flagSuffix: 'session:query',
+      queryText: (args) => extractQueryText(args[0]),
+      callbackStyle: 'last-arg',
+    }) || instrumented;
+
+  instrumented =
+    instrumentMethod(session, 'execute', state, {
+      flagSuffix: 'session:execute',
+      queryText: (args) => extractQueryText(args[0]),
+      callbackStyle: 'last-arg',
+    }) || instrumented;
+
+  instrumented = instrumentPrepareQuery(session, state) || instrumented;
+
+  if (typeof session.transaction === 'function') {
+    const flagKey = getFlagKey('session:transaction');
+
+    if (!session[flagKey]) {
+      const originalTransaction = session.transaction;
+
+      session.transaction = function instrumentedTransaction(
+        this: any,
+        callback: QueryFunction,
+        ...restArgs: any[]
+      ) {
+        if (typeof callback !== 'function') {
+          return Reflect.apply(originalTransaction, this, [
+            callback,
+            ...restArgs,
+          ]);
+        }
+
+        const wrappedCallback = (tx: unknown, ...callbackArgs: any[]) => {
+          instrumentTransactionTarget(tx, state);
+          return Reflect.apply(callback, this, [tx, ...callbackArgs]);
+        };
+
+        return Reflect.apply(originalTransaction, this, [
+          wrappedCallback,
+          ...restArgs,
+        ]);
+      };
+
+      session[flagKey] = true;
+      instrumented = true;
+    }
+  }
+
+  if (instrumented) {
+    session[INSTRUMENTED_FLAG] = true;
+  }
+
+  return instrumented;
+}
+
+export function instrumentDrizzle<TClient extends DrizzleClientLike>(
+  client: TClient,
+  config?: InstrumentDrizzleConfig,
+): TClient {
+  if (!client || !isObject(client)) {
+    return client;
+  }
+
+  const state = getState(config);
+  let instrumented = false;
+
+  instrumented =
+    instrumentMethod(client, 'query', state, {
+      flagSuffix: 'client:query',
+      queryText: (args) => extractQueryText(args[0]),
+      callbackStyle: 'last-arg',
+    }) || instrumented;
+
+  instrumented =
+    instrumentMethod(client, 'execute', state, {
+      flagSuffix: 'client:execute',
+      queryText: (args) => extractQueryText(args[0]),
+      callbackStyle: 'last-arg',
+    }) || instrumented;
+
+  if (instrumented) {
+    client[INSTRUMENTED_FLAG] = true;
+  }
+
+  return client;
+}
+
+export function instrumentDrizzleClient<TDb extends DrizzleDbLike>(
+  db: TDb,
+  config?: InstrumentDrizzleConfig,
+): TDb {
+  if (!db || !isObject(db)) {
+    return db;
+  }
+
+  if (db[INSTRUMENTED_FLAG]) {
+    return db;
+  }
+
+  const state = getState(config);
+  let instrumented = false;
+
+  instrumented =
+    instrumentSession(db as unknown as DrizzleSessionLike, state) ||
+    instrumented;
+
+  if (isObject(db.session)) {
+    instrumented = instrumentSession(db.session, state) || instrumented;
+  }
+
+  if (isObject(db._?.session)) {
+    instrumented = instrumentSession(db._.session, state) || instrumented;
+  }
+
+  if (isObject(db.$client)) {
+    instrumentDrizzle(db.$client, config);
+    instrumented = Boolean(db.$client[INSTRUMENTED_FLAG]) || instrumented;
+  }
+
+  if (instrumented) {
+    db[INSTRUMENTED_FLAG] = true;
+  }
+
+  return db;
+}

package/src/index.ts

@@ -0,0 +1,64 @@
+/**
+ * Autotel Drizzle - OpenTelemetry instrumentation for Drizzle ORM
+ *
+ * This package provides instrumentation for Drizzle ORM.
+ *
+ * Philosophy:
+ * Only include plugins for libraries that either:
+ * 1. Have NO official instrumentation (e.g., Drizzle ORM)
+ * 2. Have BROKEN official instrumentation (e.g., Mongoose in ESM+tsx)
+ * 3. Add SIGNIFICANT value beyond official packages (e.g., Kafka processing spans)
+ *
+ * For databases/ORMs with working official instrumentation, use those directly with the --import pattern:
+ * - MongoDB: @opentelemetry/instrumentation-mongodb
+ * - PostgreSQL: @opentelemetry/instrumentation-pg
+ * - MySQL: @opentelemetry/instrumentation-mysql2
+ * - Redis: @opentelemetry/instrumentation-redis
+ * - Kafka: @opentelemetry/instrumentation-kafkajs (use with autotel-plugins/kafka for processing spans)
+ *
+ * See: https://github.com/open-telemetry/opentelemetry-js-contrib
+ *
+ * @example
+ * ```typescript
+ * // Drizzle manual instrumentation
+ * import { instrumentDrizzleClient } from 'autotel-drizzle';
+ * import { drizzle } from 'drizzle-orm/node-postgres';
+ *
+ * const db = instrumentDrizzleClient(drizzle(pool));
+ * ```
+ *
+ * @packageDocumentation
+ */
+
+// Re-export common semantic conventions
+export {
+  SEMATTRS_DB_SYSTEM,
+  SEMATTRS_DB_SYSTEM_NAME,
+  SEMATTRS_DB_OPERATION,
+  SEMATTRS_DB_OPERATION_NAME,
+  SEMATTRS_DB_STATEMENT,
+  SEMATTRS_DB_NAME,
+  SEMATTRS_DB_NAMESPACE,
+  SEMATTRS_DB_COLLECTION_NAME,
+  SEMATTRS_DB_QUERY_TEXT,
+  SEMATTRS_DB_QUERY_SUMMARY,
+  SEMATTRS_NET_PEER_NAME,
+  SEMATTRS_NET_PEER_PORT,
+  SEMATTRS_GCP_BIGQUERY_JOB_ID,
+  SEMATTRS_GCP_BIGQUERY_JOB_LOCATION,
+  SEMATTRS_GCP_BIGQUERY_PROJECT_ID,
+  SEMATTRS_GCP_BIGQUERY_DESTINATION_TABLE,
+  SEMATTRS_GCP_BIGQUERY_SOURCE_TABLES,
+  SEMATTRS_GCP_BIGQUERY_STATEMENT_TYPE,
+  SEMATTRS_GCP_BIGQUERY_QUERY_HASH,
+  SEMATTRS_GCP_BIGQUERY_ROWS_AFFECTED,
+  SEMATTRS_GCP_BIGQUERY_ROWS_RETURNED,
+  SEMATTRS_GCP_BIGQUERY_SCHEMA_FIELDS,
+} from './common/constants';
+
+// Re-export Drizzle plugin
+export {
+  instrumentDrizzle,
+  instrumentDrizzleClient,
+  type InstrumentDrizzleConfig,
+} from './drizzle';