OpenTelemetry for Drizzle ORM + Better SQLite 3

database.ts

import Database from "better-sqlite3";
import { type BetterSQLite3Database, drizzle } from "drizzle-orm/better-sqlite3";
import { fileURLToPath } from "url";
import { dirname, join } from "path";
import { instrumentSqliteDrizzle } from "./drizzleTracing.js";
import * as schema from "./schema.js";

const __filename = fileURLToPath(import.meta.url);
const __dirname = dirname(__filename);

const dbPath = join(__dirname, "database.db");

const raw = new Database(dbPath, { fileMustExist: true, timeout: 5000 });

const drizzleDb = drizzle({ client: raw, schema });
instrumentSqliteDrizzle(drizzleDb, {
  dbName: "mydatabase",
  captureQueryText: true,
  maxQueryTextLength: 2000,
});

export drizzleDb;

drizzleTracing.ts

import { SpanKind, SpanStatusCode, trace } from "@opentelemetry/api";
import type { Span } from "@opentelemetry/api";
import type { BetterSQLite3Database } from "drizzle-orm/better-sqlite3";

interface InstrumentationOptions {
  tracerName?: string;
  dbSystem?: string;
  dbName?: string;
  captureQueryText?: boolean;
  maxQueryTextLength?: number;
}

const SESSION_FLAG = Symbol.for("revations.drizzle.sessionInstrumented");
const PREPARED_FLAG = Symbol.for("revations.drizzle.preparedInstrumented");

/**
 * Instruments a Drizzle better-sqlite3 database instance by wrapping the prepared
 * queries produced by the internal session. This is enough to capture every
 * Drizzle query because all helpers eventually call `session.prepareQuery`.
 */
export function instrumentSqliteDrizzle(
  db: BetterSQLite3Database<any>,
  options: InstrumentationOptions = {},
) {
  const session: any = (db as any)?.session;
  if (!session || session[SESSION_FLAG]) {
    return db;
  }

  const config = {
    tracerName: options.tracerName ?? "mydatabase-sqlite",
    dbSystem: options.dbSystem ?? "sqlite",
    dbName: options.dbName,
    captureQueryText: options.captureQueryText ?? true,
    maxQueryLength: options.maxQueryTextLength ?? 2000,
  };

  const tracer = trace.getTracer(config.tracerName);
  const originalPrepareQuery = session.prepareQuery;

  if (typeof originalPrepareQuery !== "function") {
    return db;
  }

  session.prepareQuery = function (...args: any[]) {
    const prepared = originalPrepareQuery.apply(this, args);
    const queryObj = args[0];
    return wrapPreparedQuery(prepared, queryObj);
  };

  session[SESSION_FLAG] = true;
  return db;

  function wrapPreparedQuery(prepared: any, queryObj: any) {
    if (!prepared || prepared[PREPARED_FLAG]) {
      return prepared;
    }

    const queryText = extractQueryText(queryObj);
    const operation = queryText ? extractOperation(queryText) : undefined;

    const methods: Array<keyof typeof prepared> = [
      "execute",
      "run",
      "all",
      "get",
      "values",
    ];

    for (const method of methods) {
      if (typeof prepared[method] !== "function") continue;
      const originalMethod = prepared[method].bind(prepared);

      prepared[method] = (...methodArgs: any[]) => {
        const spanName = buildSpanName(operation, method);
        return tracer.startActiveSpan(
          spanName,
          { kind: SpanKind.CLIENT },
          (span) => {
            decorateSpan(span, queryText, operation, method);
            try {
              const result = originalMethod(...methodArgs);
              return endSpanOnResult(span, result);
            } catch (error) {
              recordError(span, error);
              span.end();
              throw error;
            }
          },
        );
      };
    }

    prepared[PREPARED_FLAG] = true;
    return prepared;
  }

  function decorateSpan(
    span: Span,
    queryText: string | undefined,
    operation: string | undefined,
    method: keyof typeof session,
  ) {
    span.setAttribute("db.system", config.dbSystem);
    span.setAttribute("db.method", String(method));
    if (config.dbName) {
      span.setAttribute("db.name", config.dbName);
    }
    if (operation) {
      span.setAttribute("db.operation", operation);
    }
    if (config.captureQueryText && queryText) {
      span.setAttribute("db.statement", sanitizeQueryText(queryText, config.maxQueryLength));
    }
  }
}

function endSpanOnResult(span: any, result: any) {
  if (isPromiseLike(result)) {
    return result
      .then((value: unknown) => {
        span.setStatus({ code: SpanStatusCode.OK });
        span.end();
        return value;
      })
      .catch((error: unknown) => {
        recordError(span, error);
        span.end();
        throw error;
      });
  }

  span.setStatus({ code: SpanStatusCode.OK });
  span.end();
  return result;
}

function recordError(span: any, error: unknown) {
  span.recordException(error instanceof Error ? error : new Error(String(error)));
  span.setStatus({
    code: SpanStatusCode.ERROR,
    message: error instanceof Error ? error.message : undefined,
  });
}

function sanitizeQueryText(queryText: string, maxLength: number) {
  if (queryText.length <= maxLength) return queryText;
  return `${queryText.slice(0, maxLength)}...`;
}

function extractQueryText(queryObj: any): string | undefined {
  if (!queryObj) return undefined;
  if (typeof queryObj === "string") return queryObj;
  if (typeof queryObj.sql === "string") return queryObj.sql;
  if (typeof queryObj.queryString === "string") return queryObj.queryString;
  if (typeof queryObj.text === "string") return queryObj.text;
  if (typeof queryObj.queryChunks === "object" && typeof queryObj.sql === "string") {
    return queryObj.sql;
  }
  return undefined;
}

function extractOperation(queryText: string | undefined) {
  if (!queryText) return undefined;
  const match = queryText.trimStart().match(/^(\w+)/u);
  return match?.[1]?.toUpperCase();
}

function buildSpanName(operation: string | undefined, method: PropertyKey) {
  if (operation) {
    return `sqlite.${operation.toLowerCase()}`;
  }
  return `sqlite.${String(method)}`;
}

function isPromiseLike(value: unknown): value is Promise<unknown> {
  return Boolean(
    value &&
    (typeof value === "object" || typeof value === "function") &&
    "then" in value &&
    typeof (value as any).then === "function",
  );
}