From 7b7dc83f7d8733d63a77f2a0d655513053f00bcd Mon Sep 17 00:00:00 2001 From: Alex Holovach Date: Sun, 5 Oct 2025 08:19:58 -0500 Subject: [PATCH] update docs --- packages/otel-drizzle/README.md | 177 +++++---- packages/otel-drizzle/src/index.ts | 563 ++++++++++++++++++----------- 2 files changed, 447 insertions(+), 293 deletions(-) diff --git a/packages/otel-drizzle/README.md b/packages/otel-drizzle/README.md index 1e8de12..aff11b8 100644 --- a/packages/otel-drizzle/README.md +++ b/packages/otel-drizzle/README.md @@ -42,118 +42,149 @@ Works with any observability platform that supports OpenTelemetry including: ## Usage -There are two ways to instrument Drizzle ORM with OpenTelemetry: +### Instrument Your Drizzle Database (Recommended) -### Option 1: Instrument the Connection Pool (Recommended) - -Wrap your database connection pool with `instrumentDrizzle()` before passing it to Drizzle: +Use `instrumentDrizzleClient()` to add tracing to your Drizzle database instance. This is the simplest and most straightforward approach: ```typescript -import { drizzle } from "drizzle-orm/node-postgres"; -import { Pool } from "pg"; -import { instrumentDrizzle } from "@kubiks/otel-drizzle"; +import { drizzle } from "drizzle-orm/postgres-js"; +import { instrumentDrizzleClient } from "@kubiks/otel-drizzle"; -const pool = new Pool({ connectionString: process.env.DATABASE_URL }); -const instrumentedPool = instrumentDrizzle(pool); -const db = drizzle(instrumentedPool); +// Create your Drizzle database instance as usual +const db = drizzle(process.env.DATABASE_URL!); + +// Add instrumentation with a single line +instrumentDrizzleClient(db); // That's it! All queries are now traced automatically const users = await db.select().from(usersTable); ``` -### Option 2: Instrument an Existing Drizzle Client +### Database-Specific Examples -If you already have a Drizzle database instance or don't have access to the underlying pool, use `instrumentDrizzleClient()`. This method instruments the database at the session level, capturing all query operations: +#### PostgreSQL ```typescript -// Works with postgres-js (Postgres.js) +// PostgreSQL with postgres.js (recommended for serverless) import { drizzle } from "drizzle-orm/postgres-js"; +import postgres from "postgres"; import { instrumentDrizzleClient } from "@kubiks/otel-drizzle"; -import * as schema from "./schema"; -const db = drizzle(process.env.DATABASE_URL!, { schema }); +// Using connection string directly +const db = drizzle(process.env.DATABASE_URL!); +instrumentDrizzleClient(db, { dbSystem: "postgresql" }); -// Instrument the existing database instance -instrumentDrizzleClient(db); - -// All queries are now traced automatically -const users = await db.select().from(schema.users); -// Direct execute calls are also traced -await db.execute("SELECT * FROM users"); -// Transactions are also traced -await db.transaction(async (tx) => { - await tx.insert(schema.users).values({ name: "John" }); -}); -``` - -### Optional Configuration - -Both instrumentation methods accept the same configuration options: - -```typescript -// Option 1: Instrument the pool -const pool = new Pool({ connectionString: process.env.DATABASE_URL }); -const instrumentedPool = instrumentDrizzle(pool, { - dbSystem: "postgresql", // Database type (default: 'postgresql') - dbName: "myapp", // Database name for spans - captureQueryText: true, // Include SQL in traces (default: true) - maxQueryTextLength: 1000, // Max SQL length (default: 1000) - peerName: "db.example.com", // Database server hostname - peerPort: 5432, // Database server port -}); -const db = drizzle(instrumentedPool); - -// Option 2: Instrument the Drizzle client -const db = drizzle(pool, { schema }); +// Or with a client instance +const queryClient = postgres(process.env.DATABASE_URL!); +const db = drizzle({ client: queryClient }); instrumentDrizzleClient(db, { dbSystem: "postgresql", dbName: "myapp", - captureQueryText: true, peerName: "db.example.com", peerPort: 5432, }); ``` -### Works with All Drizzle-Supported Databases - -This package automatically detects and instruments **all databases that Drizzle ORM supports**. It works by detecting whether your database driver uses a `query` or `execute` method and instrumenting it appropriately. This includes: - -- **PostgreSQL** (node-postgres, postgres.js, Neon, Vercel Postgres, etc.) -- **MySQL** (mysql2, PlanetScale, TiDB, etc.) -- **SQLite** (better-sqlite3, LibSQL/Turso, Cloudflare D1, etc.) -- **And any other Drizzle-supported database** - ```typescript -// PostgreSQL with postgres-js (Postgres.js) - use instrumentDrizzleClient -import { drizzle } from "drizzle-orm/postgres-js"; -const db = drizzle(process.env.DATABASE_URL!); -instrumentDrizzleClient(db); - -// PostgreSQL with node-postgres (pg) - use instrumentDrizzle on pool +// PostgreSQL with node-postgres (pg) import { drizzle } from "drizzle-orm/node-postgres"; import { Pool } from "pg"; -const pool = new Pool({ connectionString: process.env.DATABASE_URL }); -const db = drizzle(instrumentDrizzle(pool)); +import { instrumentDrizzleClient } from "@kubiks/otel-drizzle"; -// MySQL with mysql2 (uses 'execute' or 'query' method) +// Using connection string directly +const db = drizzle(process.env.DATABASE_URL!); +instrumentDrizzleClient(db, { dbSystem: "postgresql" }); + +// Or with a pool instance +const pool = new Pool({ connectionString: process.env.DATABASE_URL }); +const db = drizzle({ client: pool }); +instrumentDrizzleClient(db, { dbSystem: "postgresql" }); +``` + +#### MySQL + +```typescript +// MySQL with mysql2 import { drizzle } from "drizzle-orm/mysql2"; import mysql from "mysql2/promise"; -const connection = await mysql.createConnection(process.env.DATABASE_URL); -const db = drizzle(instrumentDrizzle(connection, { dbSystem: "mysql" })); +import { instrumentDrizzleClient } from "@kubiks/otel-drizzle"; +// Using connection string directly +const db = drizzle(process.env.DATABASE_URL!); +instrumentDrizzleClient(db, { dbSystem: "mysql" }); + +// Or with a connection instance +const connection = await mysql.createConnection({ + host: "localhost", + user: "root", + database: "mydb", + // ... other connection options +}); +const db = drizzle({ client: connection }); +instrumentDrizzleClient(db, { + dbSystem: "mysql", + dbName: "mydb", + peerName: "localhost", + peerPort: 3306, +}); +``` + +#### SQLite + +```typescript // SQLite with better-sqlite3 import { drizzle } from "drizzle-orm/better-sqlite3"; import Database from "better-sqlite3"; -const sqlite = new Database("database.db"); -const db = drizzle(instrumentDrizzle(sqlite, { dbSystem: "sqlite" })); +import { instrumentDrizzleClient } from "@kubiks/otel-drizzle"; -// LibSQL/Turso (uses 'execute' method) +// Using file path directly +const db = drizzle("sqlite.db"); +instrumentDrizzleClient(db, { dbSystem: "sqlite" }); + +// Or with a Database instance +const sqlite = new Database("sqlite.db"); +const db = drizzle({ client: sqlite }); +instrumentDrizzleClient(db, { dbSystem: "sqlite" }); +``` + +```typescript +// SQLite with LibSQL/Turso import { drizzle } from "drizzle-orm/libsql"; import { createClient } from "@libsql/client"; -const client = createClient({ url: "...", authToken: "..." }); -const db = drizzle(instrumentDrizzle(client, { dbSystem: "sqlite" })); +import { instrumentDrizzleClient } from "@kubiks/otel-drizzle"; + +// Using connection config directly +const db = drizzle({ + connection: { + url: process.env.DATABASE_URL!, + authToken: process.env.DATABASE_AUTH_TOKEN, + } +}); +instrumentDrizzleClient(db, { dbSystem: "sqlite" }); + +// Or with a client instance +const client = createClient({ + url: process.env.DATABASE_URL!, + authToken: process.env.DATABASE_AUTH_TOKEN, +}); +const db = drizzle({ client }); +instrumentDrizzleClient(db, { dbSystem: "sqlite" }); ``` +### Configuration Options + +```typescript +instrumentDrizzleClient(db, { + dbSystem: "postgresql", // Database type: 'postgresql' | 'mysql' | 'sqlite' (default: 'postgresql') + dbName: "myapp", // Database name for spans + captureQueryText: true, // Include SQL in traces (default: true) + maxQueryTextLength: 1000, // Max SQL length (default: 1000) + peerName: "db.example.com", // Database server hostname + peerPort: 5432, // Database server port +}); +``` + + ## What You Get Each database query automatically creates a span with rich telemetry data: diff --git a/packages/otel-drizzle/src/index.ts b/packages/otel-drizzle/src/index.ts index afe30a0..fa91c34 100644 --- a/packages/otel-drizzle/src/index.ts +++ b/packages/otel-drizzle/src/index.ts @@ -143,16 +143,15 @@ function finalizeSpan(span: Span, error?: unknown): void { } /** - * Instruments a database connection pool or client with OpenTelemetry tracing. + * Instruments a database connection pool/client with OpenTelemetry tracing. * - * This function wraps the connection's `query` or `execute` method to automatically create - * spans for each database operation. It automatically detects which method is available - * and instruments it appropriately for any database driver. - * The instrumentation is idempotent - calling it multiple times on the same - * connection will only instrument it once. + * This function wraps the connection's `query` and `execute` methods to create spans for each database + * operation. + * The instrumentation is idempotent - calling it multiple times on the same connection will only + * instrument it once. * * @typeParam TClient - The type of the database connection pool or client - * @param client - The database connection pool or client to instrument (e.g., pg Pool, mysql2 Connection, LibSQL Client) + * @param client - The database connection pool or client to instrument * @param config - Optional configuration for instrumentation behavior * @returns The instrumented pool/client (same instance, modified in place) * @@ -164,19 +163,13 @@ function finalizeSpan(span: Span, error?: unknown): void { * import { instrumentDrizzle } from '@kubiks/otel-drizzle'; * * const pool = new Pool({ connectionString: process.env.DATABASE_URL }); - * const instrumentedPool = instrumentDrizzle(pool); - * const db = drizzle(instrumentedPool); - * - * // With custom configuration * const instrumentedPool = instrumentDrizzle(pool, { * dbSystem: 'postgresql', * dbName: 'myapp', - * captureQueryText: true, - * maxQueryTextLength: 1000, * peerName: 'db.example.com', * peerPort: 5432, * }); - * const db = drizzle(instrumentedPool); + * const db = drizzle({ client: instrumentedPool }); * ``` * * @example @@ -186,8 +179,13 @@ function finalizeSpan(span: Span, error?: unknown): void { * import mysql from 'mysql2/promise'; * import { instrumentDrizzle } from '@kubiks/otel-drizzle'; * - * const connection = await mysql.createConnection(process.env.DATABASE_URL); - * const db = drizzle(instrumentDrizzle(connection, { dbSystem: 'mysql' })); + * const connection = await mysql.createConnection({ + * host: 'localhost', + * user: 'root', + * database: 'mydb', + * }); + * const instrumentedConnection = instrumentDrizzle(connection, { dbSystem: 'mysql' }); + * const db = drizzle({ client: instrumentedConnection }); * ``` * * @example @@ -197,19 +195,24 @@ function finalizeSpan(span: Span, error?: unknown): void { * import Database from 'better-sqlite3'; * import { instrumentDrizzle } from '@kubiks/otel-drizzle'; * - * const sqlite = new Database('database.db'); - * const db = drizzle(instrumentDrizzle(sqlite, { dbSystem: 'sqlite' })); + * const sqlite = new Database('sqlite.db'); + * const instrumentedSqlite = instrumentDrizzle(sqlite, { dbSystem: 'sqlite' }); + * const db = drizzle({ client: instrumentedSqlite }); * ``` * * @example * ```typescript - * // LibSQL/Turso (automatically detects 'execute' method) + * // LibSQL/Turso * import { drizzle } from 'drizzle-orm/libsql'; * import { createClient } from '@libsql/client'; * import { instrumentDrizzle } from '@kubiks/otel-drizzle'; * - * const client = createClient({ url: '...', authToken: '...' }); - * const db = drizzle(instrumentDrizzle(client, { dbSystem: 'sqlite' })); + * const client = createClient({ + * url: process.env.DATABASE_URL!, + * authToken: process.env.DATABASE_AUTH_TOKEN, + * }); + * const instrumentedClient = instrumentDrizzle(client, { dbSystem: 'sqlite' }); + * const db = drizzle({ client: instrumentedClient }); * ``` */ export function instrumentDrizzle( @@ -219,11 +222,11 @@ export function instrumentDrizzle( if (!client) { return client; } - + // Check if client has query or execute method const hasQuery = typeof client.query === "function"; const hasExecute = typeof client.execute === "function"; - + if (!hasQuery && !hasExecute) { return client; } @@ -243,7 +246,7 @@ export function instrumentDrizzle( } = config ?? {}; const tracer = trace.getTracer(tracerName); - + // Store the original method (query or execute) const methodName = hasQuery ? "query" : "execute"; const originalMethod = hasQuery ? client.query : client.execute; @@ -338,7 +341,7 @@ export function instrumentDrizzle( }; client[INSTRUMENTED_FLAG] = true; - + // Replace the original method with the instrumented one if (hasQuery) { client.query = instrumentedMethod; @@ -369,21 +372,10 @@ interface DrizzleDbLike { } /** - * Instruments an already created Drizzle database instance with OpenTelemetry tracing. + * Instruments a Drizzle database instance with OpenTelemetry tracing. * - * This function instruments the database at the session level, intercepting: - * - `session.prepareQuery` - Used by all query builders (select, insert, update, delete) - * - `session.query` - Used for direct SQL execution - * - `$client` methods - As a fallback for underlying connection - * - * This ensures all database operations are traced, whether you use: - * - Query builders: `db.select().from(table)` - * - Direct execution: `db.execute(sql)` - * - Transactions: `db.transaction()` - * - * This is useful when you have a Drizzle database instance created with - * `const db = drizzle(connectionString)` or `const db = drizzle(pool, { schema })` - * and want to add instrumentation to it without having direct access to the underlying pool. + * This function instruments the database at the session level, automatically tracing all database + * operations including query builders, direct SQL execution, and transactions. * * The instrumentation is idempotent - calling it multiple times on the same * database will only instrument it once. @@ -395,39 +387,109 @@ interface DrizzleDbLike { * * @example * ```typescript - * // When you have a pre-created Drizzle database instance - * import { drizzle } from 'drizzle-orm/node-postgres'; - * import { Pool } from 'pg'; + * // PostgreSQL with postgres.js + * import { drizzle } from 'drizzle-orm/postgres-js'; + * import postgres from 'postgres'; * import { instrumentDrizzleClient } from '@kubiks/otel-drizzle'; - * import * as schema from './schema'; * - * const pool = new Pool({ connectionString: process.env.DATABASE_URL }); - * const db = drizzle(pool, { schema }); + * // Using connection string + * const db = drizzle(process.env.DATABASE_URL!); + * instrumentDrizzleClient(db, { dbSystem: 'postgresql' }); * - * // Instrument the existing db instance - * instrumentDrizzleClient(db, { - * dbSystem: 'postgresql', - * dbName: 'myapp', - * captureQueryText: true, - * peerName: 'db.example.com', - * peerPort: 5432, - * }); - * - * // Now all queries through db are traced - * const users = await db.select().from(schema.users); + * // Or with a client instance + * const queryClient = postgres(process.env.DATABASE_URL!); + * const db = drizzle({ client: queryClient }); + * instrumentDrizzleClient(db, { dbSystem: 'postgresql' }); * ``` * * @example * ```typescript - * // Works with any Drizzle driver + * // PostgreSQL with node-postgres (pg) + * import { drizzle } from 'drizzle-orm/node-postgres'; + * import { Pool } from 'pg'; + * import { instrumentDrizzleClient } from '@kubiks/otel-drizzle'; + * + * // Using connection string + * const db = drizzle(process.env.DATABASE_URL!); + * instrumentDrizzleClient(db, { dbSystem: 'postgresql' }); + * + * // Or with a pool + * const pool = new Pool({ connectionString: process.env.DATABASE_URL }); + * const db = drizzle({ client: pool }); + * instrumentDrizzleClient(db, { + * dbSystem: 'postgresql', + * dbName: 'myapp', + * peerName: 'db.example.com', + * peerPort: 5432, + * }); + * ``` + * + * @example + * ```typescript + * // MySQL with mysql2 * import { drizzle } from 'drizzle-orm/mysql2'; * import mysql from 'mysql2/promise'; * import { instrumentDrizzleClient } from '@kubiks/otel-drizzle'; * - * const connection = await mysql.createConnection(process.env.DATABASE_URL); - * const db = drizzle(connection); - * + * // Using connection string + * const db = drizzle(process.env.DATABASE_URL!); * instrumentDrizzleClient(db, { dbSystem: 'mysql' }); + * + * // Or with a connection + * const connection = await mysql.createConnection({ + * host: 'localhost', + * user: 'root', + * database: 'mydb', + * }); + * const db = drizzle({ client: connection }); + * instrumentDrizzleClient(db, { + * dbSystem: 'mysql', + * dbName: 'mydb', + * peerName: 'localhost', + * peerPort: 3306, + * }); + * ``` + * + * @example + * ```typescript + * // SQLite with better-sqlite3 + * import { drizzle } from 'drizzle-orm/better-sqlite3'; + * import Database from 'better-sqlite3'; + * import { instrumentDrizzleClient } from '@kubiks/otel-drizzle'; + * + * // Using file path + * const db = drizzle('sqlite.db'); + * instrumentDrizzleClient(db, { dbSystem: 'sqlite' }); + * + * // Or with a Database instance + * const sqlite = new Database('sqlite.db'); + * const db = drizzle({ client: sqlite }); + * instrumentDrizzleClient(db, { dbSystem: 'sqlite' }); + * ``` + * + * @example + * ```typescript + * // SQLite with LibSQL/Turso + * import { drizzle } from 'drizzle-orm/libsql'; + * import { createClient } from '@libsql/client'; + * import { instrumentDrizzleClient } from '@kubiks/otel-drizzle'; + * + * // Using connection config + * const db = drizzle({ + * connection: { + * url: process.env.DATABASE_URL!, + * authToken: process.env.DATABASE_AUTH_TOKEN, + * } + * }); + * instrumentDrizzleClient(db, { dbSystem: 'sqlite' }); + * + * // Or with a client instance + * const client = createClient({ + * url: process.env.DATABASE_URL!, + * authToken: process.env.DATABASE_AUTH_TOKEN, + * }); + * const db = drizzle({ client }); + * instrumentDrizzleClient(db, { dbSystem: 'sqlite' }); * ``` */ export function instrumentDrizzleClient( @@ -460,23 +522,31 @@ export function instrumentDrizzleClient( // This is where all queries actually go through if ((db as any).session && !instrumented) { const session = (db as any).session; - + // Check if session has prepareQuery method (used by select/insert/update/delete) - if (typeof session.prepareQuery === "function" && !session[INSTRUMENTED_FLAG]) { + if ( + typeof session.prepareQuery === "function" && + !session[INSTRUMENTED_FLAG] + ) { const originalPrepareQuery = session.prepareQuery; - - session.prepareQuery = function(...args: any[]) { + + session.prepareQuery = function (...args: any[]) { const prepared = originalPrepareQuery.apply(this, args); - + // Wrap the prepared query's execute method if (prepared && typeof prepared.execute === "function") { const originalPreparedExecute = prepared.execute; - - prepared.execute = function(this: any, ...executeArgs: any[]) { + + prepared.execute = function (this: any, ...executeArgs: any[]) { // Extract query information from the query object const queryObj = args[0]; // The query object passed to prepareQuery - const queryText = queryObj?.sql || queryObj?.queryString || extractQueryText(queryObj); - const operation = queryText ? extractOperation(queryText) : undefined; + const queryText = + queryObj?.sql || + queryObj?.queryString || + extractQueryText(queryObj); + const operation = queryText + ? extractOperation(queryText) + : undefined; const spanName = operation ? `drizzle.${operation.toLowerCase()}` : "drizzle.query"; @@ -494,7 +564,10 @@ export function instrumentDrizzleClient( } if (captureQueryText && queryText !== undefined) { - const sanitized = sanitizeQueryText(queryText, maxQueryTextLength); + const sanitized = sanitizeQueryText( + queryText, + maxQueryTextLength, + ); span.setAttribute(SEMATTRS_DB_STATEMENT, sanitized); } @@ -528,20 +601,25 @@ export function instrumentDrizzleClient( }); }; } - + return prepared; }; - + session[INSTRUMENTED_FLAG] = true; instrumented = true; } // Also instrument direct query method if exists - if (typeof session.query === "function" && !session[INSTRUMENTED_FLAG + "_query"]) { + if ( + typeof session.query === "function" && + !session[INSTRUMENTED_FLAG + "_query"] + ) { const originalQuery = session.query; - - session.query = function(this: any, queryString: string, params: any[]) { - const operation = queryString ? extractOperation(queryString) : undefined; + + session.query = function (this: any, queryString: string, params: any[]) { + const operation = queryString + ? extractOperation(queryString) + : undefined; const spanName = operation ? `drizzle.${operation.toLowerCase()}` : "drizzle.query"; @@ -592,188 +670,233 @@ export function instrumentDrizzleClient( } }); }; - + session[INSTRUMENTED_FLAG + "_query"] = true; instrumented = true; } // Instrument transaction method to ensure transaction sessions are also instrumented - if (typeof session.transaction === "function" && !session[INSTRUMENTED_FLAG + "_transaction"]) { + if ( + typeof session.transaction === "function" && + !session[INSTRUMENTED_FLAG + "_transaction"] + ) { const originalTransaction = session.transaction; - - session.transaction = function(this: any, transactionCallback: any, ...restArgs: any[]) { + + session.transaction = function ( + this: any, + transactionCallback: any, + ...restArgs: any[] + ) { // Wrap the transaction callback to instrument the tx object - const wrappedCallback = async function(tx: any) { - // Instrument the transaction's session if it has one - if (tx && (tx.session || tx._?.session || tx)) { - const txSession = tx.session || tx._?.session || tx; - - // Instrument tx.execute if it exists - if (typeof tx.execute === "function" && !tx[INSTRUMENTED_FLAG + "_execute"]) { - const originalTxExecute = tx.execute; - - tx.execute = function(this: any, ...executeArgs: any[]) { - const queryText = extractQueryText(executeArgs[0]); - const operation = queryText ? extractOperation(queryText) : undefined; - const spanName = operation - ? `drizzle.${operation.toLowerCase()}` - : "drizzle.query"; + const wrappedCallback = async function (tx: any) { + // Instrument the transaction's session if it has one + if (tx && (tx.session || tx._?.session || tx)) { + const txSession = tx.session || tx._?.session || tx; - // Start span - const span = tracer.startSpan(spanName, { kind: SpanKind.CLIENT }); - span.setAttribute(SEMATTRS_DB_SYSTEM, dbSystem); - span.setAttribute("db.transaction", true); + // Instrument tx.execute if it exists + if ( + typeof tx.execute === "function" && + !tx[INSTRUMENTED_FLAG + "_execute"] + ) { + const originalTxExecute = tx.execute; - if (operation) { - span.setAttribute(SEMATTRS_DB_OPERATION, operation); - } + tx.execute = function (this: any, ...executeArgs: any[]) { + const queryText = extractQueryText(executeArgs[0]); + const operation = queryText + ? extractOperation(queryText) + : undefined; + const spanName = operation + ? `drizzle.${operation.toLowerCase()}` + : "drizzle.query"; - if (dbName) { - span.setAttribute(SEMATTRS_DB_NAME, dbName); - } + // Start span + const span = tracer.startSpan(spanName, { + kind: SpanKind.CLIENT, + }); + span.setAttribute(SEMATTRS_DB_SYSTEM, dbSystem); + span.setAttribute("db.transaction", true); - if (captureQueryText && queryText !== undefined) { - const sanitized = sanitizeQueryText(queryText, maxQueryTextLength); - span.setAttribute(SEMATTRS_DB_STATEMENT, sanitized); - } + if (operation) { + span.setAttribute(SEMATTRS_DB_OPERATION, operation); + } - if (peerName) { - span.setAttribute(SEMATTRS_NET_PEER_NAME, peerName); - } + if (dbName) { + span.setAttribute(SEMATTRS_DB_NAME, dbName); + } - if (peerPort) { - span.setAttribute(SEMATTRS_NET_PEER_PORT, peerPort); - } + if (captureQueryText && queryText !== undefined) { + const sanitized = sanitizeQueryText( + queryText, + maxQueryTextLength, + ); + span.setAttribute(SEMATTRS_DB_STATEMENT, sanitized); + } - const activeContext = trace.setSpan(context.active(), span); + if (peerName) { + span.setAttribute(SEMATTRS_NET_PEER_NAME, peerName); + } - // Execute the query - return context.with(activeContext, () => { - try { - const result = originalTxExecute.apply(this, executeArgs); - return Promise.resolve(result) - .then((value) => { - finalizeSpan(span); - return value; - }) - .catch((error) => { - finalizeSpan(span, error); - throw error; - }); - } catch (error) { - finalizeSpan(span, error); - throw error; - } - }); - }; - - tx[INSTRUMENTED_FLAG + "_execute"] = true; - } - - // Also instrument txSession.prepareQuery if it exists - if (typeof txSession.prepareQuery === "function" && !txSession[INSTRUMENTED_FLAG + "_tx"]) { - const originalTxPrepareQuery = txSession.prepareQuery; - - txSession.prepareQuery = function(...prepareArgs: any[]) { - const prepared = originalTxPrepareQuery.apply(this, prepareArgs); - - // Wrap the prepared query's execute method - if (prepared && typeof prepared.execute === "function") { - const originalPreparedExecute = prepared.execute; - - prepared.execute = function(this: any, ...executeArgs: any[]) { - // Extract query information from the query object - const queryObj = prepareArgs[0]; // The query object passed to prepareQuery - const queryText = queryObj?.sql || queryObj?.queryString || extractQueryText(queryObj); - const operation = queryText ? extractOperation(queryText) : undefined; - const spanName = operation - ? `drizzle.${operation.toLowerCase()}` - : "drizzle.query"; + if (peerPort) { + span.setAttribute(SEMATTRS_NET_PEER_PORT, peerPort); + } - // Start span - const span = tracer.startSpan(spanName, { kind: SpanKind.CLIENT }); - span.setAttribute(SEMATTRS_DB_SYSTEM, dbSystem); - span.setAttribute("db.transaction", true); + const activeContext = trace.setSpan(context.active(), span); - if (operation) { - span.setAttribute(SEMATTRS_DB_OPERATION, operation); - } - - if (dbName) { - span.setAttribute(SEMATTRS_DB_NAME, dbName); - } - - if (captureQueryText && queryText !== undefined) { - const sanitized = sanitizeQueryText(queryText, maxQueryTextLength); - span.setAttribute(SEMATTRS_DB_STATEMENT, sanitized); - } - - if (peerName) { - span.setAttribute(SEMATTRS_NET_PEER_NAME, peerName); - } - - if (peerPort) { - span.setAttribute(SEMATTRS_NET_PEER_PORT, peerPort); - } - - const activeContext = trace.setSpan(context.active(), span); - - // Execute the prepared query - return context.with(activeContext, () => { - try { - const result = originalPreparedExecute.apply(this, executeArgs); - return Promise.resolve(result) - .then((value) => { - finalizeSpan(span); - return value; - }) - .catch((error) => { - finalizeSpan(span, error); - throw error; - }); - } catch (error) { - finalizeSpan(span, error); - throw error; - } + // Execute the query + return context.with(activeContext, () => { + try { + const result = originalTxExecute.apply(this, executeArgs); + return Promise.resolve(result) + .then((value) => { + finalizeSpan(span); + return value; + }) + .catch((error) => { + finalizeSpan(span, error); + throw error; }); - }; + } catch (error) { + finalizeSpan(span, error); + throw error; } - - return prepared; - }; - - txSession[INSTRUMENTED_FLAG + "_tx"] = true; - } + }); + }; + + tx[INSTRUMENTED_FLAG + "_execute"] = true; } - - // Call the original callback with the instrumented tx - return transactionCallback(tx); - }; - + + // Also instrument txSession.prepareQuery if it exists + if ( + typeof txSession.prepareQuery === "function" && + !txSession[INSTRUMENTED_FLAG + "_tx"] + ) { + const originalTxPrepareQuery = txSession.prepareQuery; + + txSession.prepareQuery = function (...prepareArgs: any[]) { + const prepared = originalTxPrepareQuery.apply( + this, + prepareArgs, + ); + + // Wrap the prepared query's execute method + if (prepared && typeof prepared.execute === "function") { + const originalPreparedExecute = prepared.execute; + + prepared.execute = function ( + this: any, + ...executeArgs: any[] + ) { + // Extract query information from the query object + const queryObj = prepareArgs[0]; // The query object passed to prepareQuery + const queryText = + queryObj?.sql || + queryObj?.queryString || + extractQueryText(queryObj); + const operation = queryText + ? extractOperation(queryText) + : undefined; + const spanName = operation + ? `drizzle.${operation.toLowerCase()}` + : "drizzle.query"; + + // Start span + const span = tracer.startSpan(spanName, { + kind: SpanKind.CLIENT, + }); + span.setAttribute(SEMATTRS_DB_SYSTEM, dbSystem); + span.setAttribute("db.transaction", true); + + if (operation) { + span.setAttribute(SEMATTRS_DB_OPERATION, operation); + } + + if (dbName) { + span.setAttribute(SEMATTRS_DB_NAME, dbName); + } + + if (captureQueryText && queryText !== undefined) { + const sanitized = sanitizeQueryText( + queryText, + maxQueryTextLength, + ); + span.setAttribute(SEMATTRS_DB_STATEMENT, sanitized); + } + + if (peerName) { + span.setAttribute(SEMATTRS_NET_PEER_NAME, peerName); + } + + if (peerPort) { + span.setAttribute(SEMATTRS_NET_PEER_PORT, peerPort); + } + + const activeContext = trace.setSpan(context.active(), span); + + // Execute the prepared query + return context.with(activeContext, () => { + try { + const result = originalPreparedExecute.apply( + this, + executeArgs, + ); + return Promise.resolve(result) + .then((value) => { + finalizeSpan(span); + return value; + }) + .catch((error) => { + finalizeSpan(span, error); + throw error; + }); + } catch (error) { + finalizeSpan(span, error); + throw error; + } + }); + }; + } + + return prepared; + }; + + txSession[INSTRUMENTED_FLAG + "_tx"] = true; + } + } + + // Call the original callback with the instrumented tx + return transactionCallback(tx); + }; + // Call the original transaction with the wrapped callback return originalTransaction.apply(this, [wrappedCallback, ...restArgs]); }; - + session[INSTRUMENTED_FLAG + "_transaction"] = true; instrumented = true; } } - // Second priority: Try to instrument via $client - // This handles the underlying connection pool if (db.$client && !instrumented) { const client = db.$client; // Check if client has query or execute function - if (typeof client.query === "function" || typeof client.execute === "function") { + if ( + typeof client.query === "function" || + typeof client.execute === "function" + ) { instrumentDrizzle(client, config); instrumented = true; } } // Third priority: Try to instrument via session.execute as fallback - if (db._ && db._.session && typeof db._.session.execute === "function" && !instrumented) { + if ( + db._ && + db._.session && + typeof db._.session.execute === "function" && + !instrumented + ) { const session = db._.session; - + // Check if already instrumented if (session[INSTRUMENTED_FLAG]) { return db;