From 47aef47db3e981d76e68abd0cd1b81798fb5d904 Mon Sep 17 00:00:00 2001 From: Abdelrahman Awad Date: Fri, 7 Nov 2025 13:33:58 +0200 Subject: [PATCH 1/2] feat: implement tracing for db0 --- src/index.ts | 6 + src/tracing.ts | 97 ++++++++++ test/tracing.test.ts | 451 +++++++++++++++++++++++++++++++++++++++++++ 3 files changed, 554 insertions(+) create mode 100644 src/tracing.ts create mode 100644 test/tracing.test.ts diff --git a/src/index.ts b/src/index.ts index 2ce54c1..1261663 100644 --- a/src/index.ts +++ b/src/index.ts @@ -13,3 +13,9 @@ export type { } from "./types.ts"; export type { ConnectorName, ConnectorOptions } from "./_connectors.ts"; + +export { + withTracing, + type TraceContext, + type TracedOperation, +} from "./tracing.ts"; diff --git a/src/tracing.ts b/src/tracing.ts new file mode 100644 index 0000000..e0f5203 --- /dev/null +++ b/src/tracing.ts @@ -0,0 +1,97 @@ +import { type TracingChannel, tracingChannel } from "node:diagnostics_channel"; +import type { Connector, Database } from "./types.ts"; +import { sqlTemplate } from "./template.ts"; + +export type TracedOperation = "query"; + +export interface TraceContext { + query: string; + method: "exec" | "sql" | "prepare.all" | "prepare.run" | "prepare.get"; +} + +const channels: Record = { + query: createChannel("query"), +}; + +/** + * Create a tracing channel for a given operation. + */ +function createChannel(operation: TracedOperation) { + return tracingChannel(`db0.${operation}`); +} + +/** + * Trace a promise with a given operation and data. + */ +async function tracePromise( + operation: TracedOperation, + exec: () => Promise, + data: TraceContext, +): Promise { + const channel = channels[operation]; + + // TODO: Remove this cast once the @types/node types are updated. + // The @types/node types incorrectly mark tracePromise as returning void, + // but according to the JSDoc and actual implementation, it returns the promise. + // This is fixed in later versions of Node.js. + // See: https://nodejs.org/api/diagnostics_channel.html#channelstracepromisefn-context-thisarg-args + return channel.tracePromise(exec, data) as unknown as Promise; +} + +type MaybeTracedDatabase = + Database & { + __traced?: boolean; + }; + +/** + * Wrap a database instance with tracing functionality. + */ +export function withTracing( + db: MaybeTracedDatabase, +): Database { + // Avoids double patching + if (db.__traced) { + return db; + } + + const tracedDb: MaybeTracedDatabase = { ...db, __traced: true }; + + tracedDb.exec = (query) => + tracePromise("query", () => db.exec(query), { query, method: "exec" }); + + tracedDb.sql = (strings, ...values) => + tracePromise("query", () => db.sql(strings, ...values), { + query: sqlTemplate(strings, ...values)[0], + method: "sql", + }); + + /** + * Prepare needs a special treatment because it returns a statement instance that needs to be patched. + */ + tracedDb.prepare = (query) => { + const statement = db.prepare(query); + const tracedStatement = { ...statement }; + + tracedStatement.all = (...params) => + tracePromise("query", () => statement.all(...params), { + query, + method: "prepare.all", + }); + + tracedStatement.run = (...params) => + tracePromise("query", () => statement.run(...params), { + query, + method: "prepare.run", + }); + + tracedStatement.get = (...params) => + tracePromise("query", () => statement.get(...params), { + query, + method: "prepare.get", + }); + + return tracedStatement; + }; + + return tracedDb; +} diff --git a/test/tracing.test.ts b/test/tracing.test.ts new file mode 100644 index 0000000..3d931de --- /dev/null +++ b/test/tracing.test.ts @@ -0,0 +1,451 @@ +import { describe, it, expect, beforeEach, vi } from "vitest"; +import { tracingChannel } from "node:diagnostics_channel"; +import { createDatabase } from "../src/database.ts"; +import { withTracing } from "../src/tracing.ts"; +import type { Database } from "../src/types.ts"; +import type { TracedOperation, TraceContext } from "../src/tracing.ts"; +import connector from "../src/connectors/better-sqlite3.ts"; + +type TracingEvent = { + start?: { data: TraceContext }; + end?: { data: TraceContext }; + asyncStart?: { data: TraceContext }; + asyncEnd?: { data: TraceContext; result?: any; error?: Error }; + error?: { data: TraceContext; error: Error }; +}; + +function createTracingListener(operationName: TracedOperation) { + const events: TracingEvent = {}; + + // Create tracing channel + const channel = tracingChannel(`db0.${operationName}`); + + // Create handlers + const startHandler = vi.fn((message: any) => { + events.start = { data: message }; + }); + + const endHandler = vi.fn((message: any) => { + events.end = { data: message }; + }); + + const asyncStartHandler = vi.fn((message: any) => { + events.asyncStart = { data: message }; + }); + + const asyncEndHandler = vi.fn((message: any) => { + events.asyncEnd = { + data: message, + result: message.result, + error: message.error, + }; + }); + + const errorHandler = vi.fn((message: any) => { + events.error = { data: message, error: message.error }; + }); + + // Subscribe using the subscribe method which listens to all events + channel.subscribe({ + start: startHandler, + end: endHandler, + asyncStart: asyncStartHandler, + asyncEnd: asyncEndHandler, + error: errorHandler, + }); + + return { + events, + handlers: { + start: startHandler, + end: endHandler, + asyncStart: asyncStartHandler, + asyncEnd: asyncEndHandler, + error: errorHandler, + }, + cleanup: () => { + channel.unsubscribe({ + start: startHandler, + end: endHandler, + asyncStart: asyncStartHandler, + asyncEnd: asyncEndHandler, + error: errorHandler, + }); + }, + }; +} + +describe("tracing", () => { + let db: Database; + + beforeEach(async () => { + const plainDb = createDatabase( + connector({ + name: ":memory:", + }), + ); + db = withTracing(plainDb); + + // Create a test table + await db.exec( + `CREATE TABLE users (id INTEGER PRIMARY KEY, name TEXT, email TEXT)`, + ); + }); + + describe("opt-in behavior", () => { + it("should not emit tracing events without withTracing wrapper", async () => { + const plainDb = createDatabase( + connector({ + name: ":memory:", + }), + ); + const listener = createTracingListener("query"); + + await plainDb.exec(`CREATE TABLE test (id INTEGER PRIMARY KEY)`); + await plainDb.sql`SELECT * FROM test`; + + // No tracing events should be emitted + expect(listener.handlers.start).not.toHaveBeenCalled(); + expect(listener.handlers.end).not.toHaveBeenCalled(); + expect(listener.handlers.asyncStart).not.toHaveBeenCalled(); + expect(listener.handlers.asyncEnd).not.toHaveBeenCalled(); + expect(listener.handlers.error).not.toHaveBeenCalled(); + + listener.cleanup(); + }); + + it("should prevent double tracing when wrapped multiple times", async () => { + const plainDb = createDatabase( + connector({ + name: ":memory:", + }), + ); + const tracedOnce = withTracing(plainDb); + const tracedTwice = withTracing(tracedOnce); + + const listener = createTracingListener("query"); + + await tracedTwice.exec(`CREATE TABLE test (id INTEGER PRIMARY KEY)`); + + // Should only be called once, not twice + expect(listener.handlers.start).toHaveBeenCalledTimes(1); + expect(listener.handlers.end).toHaveBeenCalledTimes(1); + expect(listener.handlers.asyncStart).toHaveBeenCalledTimes(1); + expect(listener.handlers.asyncEnd).toHaveBeenCalledTimes(1); + + listener.cleanup(); + }); + }); + + describe("exec", () => { + it("should emit correct tracing events on success", async () => { + const listener = createTracingListener("query"); + + const result = await db.exec( + `INSERT INTO users (id, name, email) VALUES (1, 'John Doe', 'john@example.com')`, + ); + + expect(result).toBeDefined(); + expect(listener.handlers.start).toHaveBeenCalledTimes(1); + expect(listener.handlers.end).toHaveBeenCalledTimes(1); + expect(listener.handlers.asyncStart).toHaveBeenCalledTimes(1); + expect(listener.handlers.asyncEnd).toHaveBeenCalledTimes(1); + expect(listener.handlers.error).not.toHaveBeenCalled(); + + expect(listener.events.start?.data.query).toContain("INSERT INTO users"); + expect(listener.events.start?.data.method).toBe("exec"); + + listener.cleanup(); + }); + + it("should emit error event on failure", async () => { + const listener = createTracingListener("query"); + + await expect( + db.exec(`INSERT INTO non_existing_table VALUES (1, 'test')`), + ).rejects.toThrow(); + + expect(listener.handlers.start).toHaveBeenCalledTimes(1); + // asyncStart might not be called if error is thrown synchronously + expect(listener.handlers.asyncStart).not.toHaveBeenCalledTimes(1); + expect(listener.handlers.error).toHaveBeenCalledTimes(1); + expect(listener.events.error?.error).toBeDefined(); + expect(listener.events.error?.data.query).toContain( + "INSERT INTO non_existing_table", + ); + expect(listener.events.error?.data.method).toBe("exec"); + + listener.cleanup(); + }); + }); + + describe("sql", () => { + it("should emit correct tracing events on SELECT success", async () => { + const listener = createTracingListener("query"); + + await db.exec( + `INSERT INTO users (id, name, email) VALUES (1, 'John Doe', 'john@example.com')`, + ); + + const result = await db.sql`SELECT * FROM users WHERE id = ${1}`; + + expect(result.rows).toHaveLength(1); + expect(listener.handlers.start).toHaveBeenCalled(); + expect(listener.handlers.end).toHaveBeenCalled(); + expect(listener.handlers.asyncStart).toHaveBeenCalled(); + expect(listener.handlers.asyncEnd).toHaveBeenCalled(); + expect(listener.handlers.error).not.toHaveBeenCalled(); + + // Find the SELECT query event + const selectCalls = listener.handlers.start.mock.calls.filter((call) => + call[0].query.includes("SELECT"), + ); + expect(selectCalls.length).toBeGreaterThan(0); + expect(selectCalls[0][0].method).toBe("sql"); + expect(selectCalls[0][0].query).toContain("SELECT * FROM users"); + + listener.cleanup(); + }); + + it("should emit correct tracing events on INSERT with RETURNING", async () => { + const listener = createTracingListener("query"); + + const result = + await db.sql`INSERT INTO users (id, name, email) VALUES (${2}, ${"Jane Doe"}, ${"jane@example.com"}) RETURNING *`; + + expect(result.rows).toHaveLength(1); + expect(listener.handlers.start).toHaveBeenCalled(); + expect(listener.handlers.end).toHaveBeenCalled(); + expect(listener.handlers.asyncStart).toHaveBeenCalled(); + expect(listener.handlers.asyncEnd).toHaveBeenCalled(); + expect(listener.handlers.error).not.toHaveBeenCalled(); + + // Find the INSERT query event + const insertCalls = listener.handlers.start.mock.calls.filter((call) => + call[0].query.includes("INSERT"), + ); + expect(insertCalls.length).toBeGreaterThan(0); + expect(insertCalls[0][0].method).toBe("sql"); + expect(insertCalls[0][0].query).toContain("INSERT INTO users"); + expect(insertCalls[0][0].query).toContain("RETURNING"); + + listener.cleanup(); + }); + + it("should emit error event on failure", async () => { + const listener = createTracingListener("query"); + + await expect( + db.sql`SELECT * FROM non_existing_table WHERE id = ${1}`, + ).rejects.toThrow(); + + expect(listener.handlers.start).toHaveBeenCalledTimes(1); + expect(listener.handlers.asyncStart).toHaveBeenCalledTimes(1); + expect(listener.handlers.error).toHaveBeenCalledTimes(1); + expect(listener.events.error?.error).toBeDefined(); + expect(listener.events.error?.data.query).toContain( + "SELECT * FROM non_existing_table", + ); + expect(listener.events.error?.data.method).toBe("sql"); + + listener.cleanup(); + }); + }); + + describe("prepare.all", () => { + it("should emit correct tracing events on success", async () => { + const listener = createTracingListener("query"); + + await db.exec( + `INSERT INTO users (id, name, email) VALUES (1, 'John Doe', 'john@example.com')`, + ); + await db.exec( + `INSERT INTO users (id, name, email) VALUES (2, 'Jane Doe', 'jane@example.com')`, + ); + + const stmt = db.prepare("SELECT * FROM users WHERE id > ?"); + const rows = await stmt.all(0); + + expect(rows).toHaveLength(2); + + // Find the prepare.all query event + const prepareCalls = listener.handlers.start.mock.calls.filter( + (call) => call[0].method === "prepare.all", + ); + expect(prepareCalls.length).toBeGreaterThan(0); + expect(prepareCalls[0][0].query).toContain("SELECT * FROM users"); + expect(prepareCalls[0][0].method).toBe("prepare.all"); + + expect(listener.handlers.error).not.toHaveBeenCalled(); + + listener.cleanup(); + }); + + it("should emit error event on failure", async () => { + const listener = createTracingListener("query"); + + const stmt = db.prepare("SELECT * FROM non_existing_table WHERE id = ?"); + + await expect(stmt.all(1)).rejects.toThrow(); + + const prepareCalls = listener.handlers.error.mock.calls.filter( + (call) => call[0].method === "prepare.all", + ); + expect(prepareCalls.length).toBeGreaterThan(0); + expect(prepareCalls[0][0].error).toBeDefined(); + expect(prepareCalls[0][0].query).toContain( + "SELECT * FROM non_existing_table", + ); + expect(prepareCalls[0][0].method).toBe("prepare.all"); + + listener.cleanup(); + }); + }); + + describe("prepare.run", () => { + it("should emit correct tracing events on success", async () => { + const listener = createTracingListener("query"); + + const stmt = db.prepare( + "INSERT INTO users (id, name, email) VALUES (?, ?, ?)", + ); + const result = await stmt.run(3, "Bob Smith", "bob@example.com"); + + expect(result).toBeDefined(); + + // Find the prepare.run query event + const prepareCalls = listener.handlers.start.mock.calls.filter( + (call) => call[0].method === "prepare.run", + ); + expect(prepareCalls.length).toBeGreaterThan(0); + expect(prepareCalls[0][0].query).toContain("INSERT INTO users"); + expect(prepareCalls[0][0].method).toBe("prepare.run"); + + expect(listener.handlers.error).not.toHaveBeenCalled(); + + listener.cleanup(); + }); + + it("should emit error event on failure", async () => { + const listener = createTracingListener("query"); + + const stmt = db.prepare( + "INSERT INTO non_existing_table (id, name) VALUES (?, ?)", + ); + + await expect(stmt.run(1, "test")).rejects.toThrow(); + + const prepareCalls = listener.handlers.error.mock.calls.filter( + (call) => call[0].method === "prepare.run", + ); + expect(prepareCalls.length).toBeGreaterThan(0); + expect(prepareCalls[0][0].error).toBeDefined(); + expect(prepareCalls[0][0].query).toContain( + "INSERT INTO non_existing_table", + ); + expect(prepareCalls[0][0].method).toBe("prepare.run"); + + listener.cleanup(); + }); + }); + + describe("prepare.get", () => { + it("should emit correct tracing events on success", async () => { + const listener = createTracingListener("query"); + + await db.exec( + `INSERT INTO users (id, name, email) VALUES (1, 'John Doe', 'john@example.com')`, + ); + + const stmt = db.prepare("SELECT * FROM users WHERE id = ?"); + const row = await stmt.get(1); + + expect(row).toBeDefined(); + expect((row as any).name).toBe("John Doe"); + + // Find the prepare.get query event + const prepareCalls = listener.handlers.start.mock.calls.filter( + (call) => call[0].method === "prepare.get", + ); + expect(prepareCalls.length).toBeGreaterThan(0); + expect(prepareCalls[0][0].query).toContain("SELECT * FROM users"); + expect(prepareCalls[0][0].method).toBe("prepare.get"); + + expect(listener.handlers.error).not.toHaveBeenCalled(); + + listener.cleanup(); + }); + + it("should emit error event on failure", async () => { + const listener = createTracingListener("query"); + + const stmt = db.prepare("SELECT * FROM non_existing_table WHERE id = ?"); + + await expect(stmt.get(1)).rejects.toThrow(); + + const prepareCalls = listener.handlers.error.mock.calls.filter( + (call) => call[0].method === "prepare.get", + ); + expect(prepareCalls.length).toBeGreaterThan(0); + expect(prepareCalls[0][0].error).toBeDefined(); + expect(prepareCalls[0][0].query).toContain( + "SELECT * FROM non_existing_table", + ); + expect(prepareCalls[0][0].method).toBe("prepare.get"); + + listener.cleanup(); + }); + }); + + describe("query reconstruction in sql method", () => { + it("should correctly reconstruct query with template literals", async () => { + const listener = createTracingListener("query"); + + const name = "John Doe"; + const email = "john@example.com"; + await db.sql`SELECT * FROM users WHERE name = ${name} AND email = ${email}`; + + const selectCalls = listener.handlers.start.mock.calls.filter( + (call) => call[0].method === "sql" && call[0].query.includes("SELECT"), + ); + expect(selectCalls.length).toBeGreaterThan(0); + + // Query should be reconstructed with placeholders + const query = selectCalls[0][0].query; + expect(query).toContain("SELECT * FROM users"); + expect(query).toContain("WHERE"); + expect(query).toContain("name"); + expect(query).toContain("email"); + + listener.cleanup(); + }); + }); + + describe("multiple operations", () => { + it("should emit separate events for each operation", async () => { + const listener = createTracingListener("query"); + + await db.exec( + `INSERT INTO users (id, name, email) VALUES (1, 'John Doe', 'john@example.com')`, + ); + await db.sql`SELECT * FROM users WHERE id = ${1}`; + const stmt = db.prepare("SELECT * FROM users WHERE id = ?"); + await stmt.all(1); + + // Should have events for: exec, sql (prepare.all internally), and prepare.all + expect(listener.handlers.start.mock.calls.length).toBeGreaterThanOrEqual( + 3, + ); + + // Check that different methods were called + const methods = listener.handlers.start.mock.calls.map( + (call) => call[0].method, + ); + expect(methods).toContain("exec"); + expect(methods).toContain("sql"); + expect(methods).toContain("prepare.all"); + + listener.cleanup(); + }); + }); +}); From 375b1875f24e7abffa1bf88c77b8144f35596415 Mon Sep 17 00:00:00 2001 From: Abdelrahman Awad Date: Sun, 9 Nov 2025 00:25:33 +0200 Subject: [PATCH 2/2] feat: add dialect to the context data --- src/tracing.ts | 20 +++++++++++++++----- test/tracing.test.ts | 12 ++++++++++++ 2 files changed, 27 insertions(+), 5 deletions(-) diff --git a/src/tracing.ts b/src/tracing.ts index e0f5203..a171d4c 100644 --- a/src/tracing.ts +++ b/src/tracing.ts @@ -1,5 +1,5 @@ import { type TracingChannel, tracingChannel } from "node:diagnostics_channel"; -import type { Connector, Database } from "./types.ts"; +import type { Connector, Database, SQLDialect } from "./types.ts"; import { sqlTemplate } from "./template.ts"; export type TracedOperation = "query"; @@ -7,6 +7,7 @@ export type TracedOperation = "query"; export interface TraceContext { query: string; method: "exec" | "sql" | "prepare.all" | "prepare.run" | "prepare.get"; + dialect: SQLDialect; } const channels: Record = { @@ -57,12 +58,17 @@ export function withTracing( const tracedDb: MaybeTracedDatabase = { ...db, __traced: true }; tracedDb.exec = (query) => - tracePromise("query", () => db.exec(query), { query, method: "exec" }); + tracePromise("query", () => db.exec(query), { + query, + method: "exec", + dialect: db.dialect, + }); tracedDb.sql = (strings, ...values) => tracePromise("query", () => db.sql(strings, ...values), { query: sqlTemplate(strings, ...values)[0], method: "sql", + dialect: db.dialect, }); /** @@ -71,23 +77,27 @@ export function withTracing( tracedDb.prepare = (query) => { const statement = db.prepare(query); const tracedStatement = { ...statement }; + const partialContext = { + query, + dialect: db.dialect, + }; tracedStatement.all = (...params) => tracePromise("query", () => statement.all(...params), { - query, method: "prepare.all", + ...partialContext, }); tracedStatement.run = (...params) => tracePromise("query", () => statement.run(...params), { - query, method: "prepare.run", + ...partialContext, }); tracedStatement.get = (...params) => tracePromise("query", () => statement.get(...params), { - query, method: "prepare.get", + ...partialContext, }); return tracedStatement; diff --git a/test/tracing.test.ts b/test/tracing.test.ts index 3d931de..6708e19 100644 --- a/test/tracing.test.ts +++ b/test/tracing.test.ts @@ -154,6 +154,7 @@ describe("tracing", () => { expect(listener.events.start?.data.query).toContain("INSERT INTO users"); expect(listener.events.start?.data.method).toBe("exec"); + expect(listener.events.start?.data.dialect).toBe("sqlite"); listener.cleanup(); }); @@ -174,6 +175,7 @@ describe("tracing", () => { "INSERT INTO non_existing_table", ); expect(listener.events.error?.data.method).toBe("exec"); + expect(listener.events.error?.data.dialect).toBe("sqlite"); listener.cleanup(); }); @@ -203,6 +205,7 @@ describe("tracing", () => { expect(selectCalls.length).toBeGreaterThan(0); expect(selectCalls[0][0].method).toBe("sql"); expect(selectCalls[0][0].query).toContain("SELECT * FROM users"); + expect(selectCalls[0][0].dialect).toBe("sqlite"); listener.cleanup(); }); @@ -228,6 +231,7 @@ describe("tracing", () => { expect(insertCalls[0][0].method).toBe("sql"); expect(insertCalls[0][0].query).toContain("INSERT INTO users"); expect(insertCalls[0][0].query).toContain("RETURNING"); + expect(insertCalls[0][0].dialect).toBe("sqlite"); listener.cleanup(); }); @@ -247,6 +251,7 @@ describe("tracing", () => { "SELECT * FROM non_existing_table", ); expect(listener.events.error?.data.method).toBe("sql"); + expect(listener.events.error?.data.dialect).toBe("sqlite"); listener.cleanup(); }); @@ -275,6 +280,7 @@ describe("tracing", () => { expect(prepareCalls.length).toBeGreaterThan(0); expect(prepareCalls[0][0].query).toContain("SELECT * FROM users"); expect(prepareCalls[0][0].method).toBe("prepare.all"); + expect(prepareCalls[0][0].dialect).toBe("sqlite"); expect(listener.handlers.error).not.toHaveBeenCalled(); @@ -297,6 +303,7 @@ describe("tracing", () => { "SELECT * FROM non_existing_table", ); expect(prepareCalls[0][0].method).toBe("prepare.all"); + expect(prepareCalls[0][0].dialect).toBe("sqlite"); listener.cleanup(); }); @@ -320,6 +327,7 @@ describe("tracing", () => { expect(prepareCalls.length).toBeGreaterThan(0); expect(prepareCalls[0][0].query).toContain("INSERT INTO users"); expect(prepareCalls[0][0].method).toBe("prepare.run"); + expect(prepareCalls[0][0].dialect).toBe("sqlite"); expect(listener.handlers.error).not.toHaveBeenCalled(); @@ -344,6 +352,7 @@ describe("tracing", () => { "INSERT INTO non_existing_table", ); expect(prepareCalls[0][0].method).toBe("prepare.run"); + expect(prepareCalls[0][0].dialect).toBe("sqlite"); listener.cleanup(); }); @@ -370,6 +379,7 @@ describe("tracing", () => { expect(prepareCalls.length).toBeGreaterThan(0); expect(prepareCalls[0][0].query).toContain("SELECT * FROM users"); expect(prepareCalls[0][0].method).toBe("prepare.get"); + expect(prepareCalls[0][0].dialect).toBe("sqlite"); expect(listener.handlers.error).not.toHaveBeenCalled(); @@ -392,6 +402,7 @@ describe("tracing", () => { "SELECT * FROM non_existing_table", ); expect(prepareCalls[0][0].method).toBe("prepare.get"); + expect(prepareCalls[0][0].dialect).toBe("sqlite"); listener.cleanup(); }); @@ -416,6 +427,7 @@ describe("tracing", () => { expect(query).toContain("WHERE"); expect(query).toContain("name"); expect(query).toContain("email"); + expect(selectCalls[0][0].dialect).toBe("sqlite"); listener.cleanup(); });