From 2b777fc2105b9645ca84d3f9069235a2cebedf46 Mon Sep 17 00:00:00 2001 From: Abhijeet Prasad Date: Wed, 31 Jul 2024 14:22:28 -0400 Subject: [PATCH] feat(cloudflare): Add Cloudflare D1 instrumentation --- packages/cloudflare/README.md | 12 ++ packages/cloudflare/src/d1.ts | 154 +++++++++++++++++ packages/cloudflare/src/index.ts | 2 + packages/cloudflare/test/d1.test.ts | 253 ++++++++++++++++++++++++++++ 4 files changed, 421 insertions(+) create mode 100644 packages/cloudflare/src/d1.ts create mode 100644 packages/cloudflare/test/d1.test.ts diff --git a/packages/cloudflare/README.md b/packages/cloudflare/README.md index dc0d6de01274..7c7512e2ed1d 100644 --- a/packages/cloudflare/README.md +++ b/packages/cloudflare/README.md @@ -136,3 +136,15 @@ Sentry.captureEvent({ ], }); ``` + +## Cloudflare D1 Instrumentation + +You can use the `instrumentD1WithSentry` method to instrument [Cloudflare D1](https://developers.cloudflare.com/d1/), +Cloudflare's serverless SQL database with Sentry. + +```javascript +// env.DB is the D1 DB binding configured in your `wrangler.toml` +const db = instrumentD1WithSentry(env.DB); +// Now you can use the database as usual +await db.prepare('SELECT * FROM table WHERE id = ?').bind(1).run(); +``` diff --git a/packages/cloudflare/src/d1.ts b/packages/cloudflare/src/d1.ts new file mode 100644 index 000000000000..a67574fabb9e --- /dev/null +++ b/packages/cloudflare/src/d1.ts @@ -0,0 +1,154 @@ +import type { D1Database, D1PreparedStatement, D1Response } from '@cloudflare/workers-types'; +import { SEMANTIC_ATTRIBUTE_SENTRY_ORIGIN, SPAN_STATUS_ERROR, addBreadcrumb, startSpan } from '@sentry/core'; +import type { Span, SpanAttributes, StartSpanOptions } from '@sentry/types'; + +// Patching is based on internal Cloudflare D1 API +// https://github.com/cloudflare/workerd/blob/cd5279e7b305003f1d9c851e73efa9d67e4b68b2/src/cloudflare/internal/d1-api.ts + +const patchedStatement = new WeakSet(); + +/** + * Patches the query methods of a Cloudflare D1 prepared statement with Sentry. + */ +function instrumentD1PreparedStatementQueries(statement: D1PreparedStatement, query: string): D1PreparedStatement { + if (patchedStatement.has(statement)) { + return statement; + } + + // eslint-disable-next-line @typescript-eslint/unbound-method + statement.first = new Proxy(statement.first, { + apply(target, thisArg, args: Parameters) { + return startSpan(createStartSpanOptions(query, 'first'), async () => { + const res = await Reflect.apply(target, thisArg, args); + createD1Breadcrumb(query, 'first'); + return res; + }); + }, + }); + + // eslint-disable-next-line @typescript-eslint/unbound-method + statement.run = new Proxy(statement.run, { + apply(target, thisArg, args: Parameters) { + return startSpan(createStartSpanOptions(query, 'run'), async span => { + const d1Response = await Reflect.apply(target, thisArg, args); + applyD1ReturnObjectToSpan(span, d1Response); + createD1Breadcrumb(query, 'run', d1Response); + return d1Response; + }); + }, + }); + + // eslint-disable-next-line @typescript-eslint/unbound-method + statement.all = new Proxy(statement.all, { + apply(target, thisArg, args: Parameters) { + return startSpan(createStartSpanOptions(query, 'all'), async span => { + const d1Result = await Reflect.apply(target, thisArg, args); + applyD1ReturnObjectToSpan(span, d1Result); + createD1Breadcrumb(query, 'all', d1Result); + return d1Result; + }); + }, + }); + + // eslint-disable-next-line @typescript-eslint/unbound-method + statement.raw = new Proxy(statement.raw, { + apply(target, thisArg, args: Parameters) { + return startSpan(createStartSpanOptions(query, 'raw'), async () => { + const res = await Reflect.apply(target, thisArg, args); + createD1Breadcrumb(query, 'raw'); + return res; + }); + }, + }); + + patchedStatement.add(statement); + + return statement; +} + +/** + * Instruments a Cloudflare D1 prepared statement with Sentry. + * + * This is meant to be used as a top-level call, under the hood it calls `instrumentD1PreparedStatementQueries` + * to patch the query methods. The reason for this abstraction is to ensure that the `bind` method is also patched. + */ +function instrumentD1PreparedStatement(statement: D1PreparedStatement, query: string): D1PreparedStatement { + // statement.bind() returns a new instance of D1PreparedStatement, so we have to patch it as well. + // eslint-disable-next-line @typescript-eslint/unbound-method + statement.bind = new Proxy(statement.bind, { + apply(target, thisArg, args: Parameters) { + return instrumentD1PreparedStatementQueries(Reflect.apply(target, thisArg, args), query); + }, + }); + + return instrumentD1PreparedStatementQueries(statement, query); +} + +/** + * Add D1Response meta information to a span. + * + * See: https://developers.cloudflare.com/d1/build-with-d1/d1-client-api/#return-object + */ +function applyD1ReturnObjectToSpan(span: Span, d1Result: D1Response): void { + if (!d1Result.success) { + span.setStatus({ code: SPAN_STATUS_ERROR }); + } + + span.setAttributes(getAttributesFromD1Response(d1Result)); +} + +function getAttributesFromD1Response(d1Result: D1Response): SpanAttributes { + return { + 'cloudflare.d1.duration': d1Result.meta.duration, + 'cloudflare.d1.rows_read': d1Result.meta.rows_read, + 'cloudflare.d1.rows_written': d1Result.meta.rows_written, + }; +} + +function createD1Breadcrumb(query: string, type: 'first' | 'run' | 'all' | 'raw', d1Result?: D1Response): void { + addBreadcrumb({ + category: 'query', + message: query, + data: { + ...(d1Result ? getAttributesFromD1Response(d1Result) : {}), + 'cloudflare.d1.query_type': type, + }, + }); +} + +function createStartSpanOptions(query: string, type: 'first' | 'run' | 'all' | 'raw'): StartSpanOptions { + return { + op: 'db.query', + name: query, + attributes: { + 'cloudflare.d1.query_type': type, + [SEMANTIC_ATTRIBUTE_SENTRY_ORIGIN]: 'auto.db.cloudflare.d1', + }, + }; +} + +/** + * Instruments Cloudflare D1 bindings with Sentry. + * + * Currently, only prepared statements are instrumented. `db.exec` and `db.batch` are not instrumented. + * + * @example + * + * ```js + * // env.DB is the D1 DB binding configured in your `wrangler.toml` + * const db = instrumentD1WithSentry(env.DB); + * // Now you can use the database as usual + * await db.prepare('SELECT * FROM table WHERE id = ?').bind(1).run(); + * ``` + */ +export function instrumentD1WithSentry(db: D1Database): D1Database { + // eslint-disable-next-line @typescript-eslint/unbound-method + db.prepare = new Proxy(db.prepare, { + apply(target, thisArg, args: Parameters) { + const [query] = args; + return instrumentD1PreparedStatement(Reflect.apply(target, thisArg, args), query); + }, + }); + + return db; +} diff --git a/packages/cloudflare/src/index.ts b/packages/cloudflare/src/index.ts index 3708d3ae9382..a4a466fa5bb5 100644 --- a/packages/cloudflare/src/index.ts +++ b/packages/cloudflare/src/index.ts @@ -91,3 +91,5 @@ export { CloudflareClient } from './client'; export { getDefaultIntegrations } from './sdk'; export { fetchIntegration } from './integrations/fetch'; + +export { instrumentD1WithSentry } from './d1'; diff --git a/packages/cloudflare/test/d1.test.ts b/packages/cloudflare/test/d1.test.ts new file mode 100644 index 000000000000..c86538b96208 --- /dev/null +++ b/packages/cloudflare/test/d1.test.ts @@ -0,0 +1,253 @@ +import { beforeEach, describe, expect, test, vi } from 'vitest'; + +import * as SentryCore from '@sentry/core'; + +import type { D1Database, D1PreparedStatement } from '@cloudflare/workers-types'; + +import { instrumentD1WithSentry } from '../src/d1'; + +const MOCK_FIRST_RETURN_VALUE = { id: 1, name: 'Foo' }; + +const MOCK_RAW_RETURN_VALUE = [ + { id: 1, name: 'Foo' }, + { id: 2, name: 'Bar' }, +]; + +const MOCK_D1_RESPONSE = { + success: true, + meta: { + duration: 1, + size_after: 2, + rows_read: 3, + rows_written: 4, + last_row_id: 5, + changed_db: false, + changes: 7, + }, +}; + +describe('instrumentD1WithSentry', () => { + beforeEach(() => { + vi.clearAllMocks(); + }); + + const startSpanSpy = vi.spyOn(SentryCore, 'startSpan'); + const addBreadcrumbSpy = vi.spyOn(SentryCore, 'addBreadcrumb'); + + function createMockD1Statement(): D1PreparedStatement { + return { + bind: vi.fn().mockImplementation(createMockD1Statement), + first: vi.fn().mockImplementation(() => Promise.resolve(MOCK_FIRST_RETURN_VALUE)), + run: vi.fn().mockImplementation(() => Promise.resolve(MOCK_D1_RESPONSE)), + all: vi.fn().mockImplementation(() => Promise.resolve(MOCK_D1_RESPONSE)), + raw: vi.fn().mockImplementation(() => Promise.resolve(MOCK_RAW_RETURN_VALUE)), + }; + } + + function createMockD1Database(): D1Database { + return { + prepare: vi.fn().mockImplementation(createMockD1Statement), + dump: vi.fn(), + batch: vi.fn(), + exec: vi.fn(), + }; + } + + describe('statement.first()', () => { + test('does not change return value', async () => { + const instrumentedDb = instrumentD1WithSentry(createMockD1Database()); + const response = await instrumentedDb.prepare('SELECT * FROM users').first(); + expect(response).toEqual(MOCK_FIRST_RETURN_VALUE); + }); + + test('instruments with spans', async () => { + const instrumentedDb = instrumentD1WithSentry(createMockD1Database()); + await instrumentedDb.prepare('SELECT * FROM users').first(); + + expect(startSpanSpy).toHaveBeenCalledTimes(1); + expect(startSpanSpy).toHaveBeenLastCalledWith( + { + attributes: { + 'cloudflare.d1.query_type': 'first', + 'sentry.origin': 'auto.db.cloudflare.d1', + }, + name: 'SELECT * FROM users', + op: 'db.query', + }, + expect.any(Function), + ); + }); + + test('instruments with breadcrumbs', async () => { + const instrumentedDb = instrumentD1WithSentry(createMockD1Database()); + await instrumentedDb.prepare('SELECT * FROM users').first(); + + expect(addBreadcrumbSpy).toHaveBeenCalledTimes(1); + expect(addBreadcrumbSpy).toHaveBeenLastCalledWith({ + category: 'query', + message: 'SELECT * FROM users', + data: { + 'cloudflare.d1.query_type': 'first', + }, + }); + }); + + test('works with statement.bind()', async () => { + const instrumentedDb = instrumentD1WithSentry(createMockD1Database()); + await instrumentedDb.prepare('SELECT * FROM users').bind().first(); + + expect(startSpanSpy).toHaveBeenCalledTimes(1); + expect(addBreadcrumbSpy).toHaveBeenCalledTimes(1); + }); + }); + + describe('statement.run()', () => { + test('does not change return value', async () => { + const instrumentedDb = instrumentD1WithSentry(createMockD1Database()); + const response = await instrumentedDb.prepare('INSERT INTO users (name) VALUES (?)').run(); + expect(response).toEqual(MOCK_D1_RESPONSE); + }); + + test('instruments with spans', async () => { + const instrumentedDb = instrumentD1WithSentry(createMockD1Database()); + await instrumentedDb.prepare('INSERT INTO users (name) VALUES (?)').run(); + + expect(startSpanSpy).toHaveBeenCalledTimes(1); + expect(startSpanSpy).toHaveBeenLastCalledWith( + { + attributes: { + 'cloudflare.d1.query_type': 'run', + 'sentry.origin': 'auto.db.cloudflare.d1', + }, + name: 'INSERT INTO users (name) VALUES (?)', + op: 'db.query', + }, + expect.any(Function), + ); + }); + + test('instruments with breadcrumbs', async () => { + const instrumentedDb = instrumentD1WithSentry(createMockD1Database()); + await instrumentedDb.prepare('INSERT INTO users (name) VALUES (?)').run(); + + expect(addBreadcrumbSpy).toHaveBeenCalledTimes(1); + expect(addBreadcrumbSpy).toHaveBeenLastCalledWith({ + category: 'query', + message: 'INSERT INTO users (name) VALUES (?)', + data: { + 'cloudflare.d1.query_type': 'run', + 'cloudflare.d1.duration': 1, + 'cloudflare.d1.rows_read': 3, + 'cloudflare.d1.rows_written': 4, + }, + }); + }); + + test('works with statement.bind()', async () => { + const instrumentedDb = instrumentD1WithSentry(createMockD1Database()); + await instrumentedDb.prepare('SELECT * FROM users').bind().run(); + + expect(startSpanSpy).toHaveBeenCalledTimes(1); + expect(addBreadcrumbSpy).toHaveBeenCalledTimes(1); + }); + }); + + describe('statement.all()', () => { + test('does not change return value', async () => { + const instrumentedDb = instrumentD1WithSentry(createMockD1Database()); + const response = await instrumentedDb.prepare('INSERT INTO users (name) VALUES (?)').run(); + expect(response).toEqual(MOCK_D1_RESPONSE); + }); + + test('instruments with spans', async () => { + const instrumentedDb = instrumentD1WithSentry(createMockD1Database()); + await instrumentedDb.prepare('INSERT INTO users (name) VALUES (?)').all(); + + expect(startSpanSpy).toHaveBeenCalledTimes(1); + expect(startSpanSpy).toHaveBeenLastCalledWith( + { + attributes: { + 'cloudflare.d1.query_type': 'all', + 'sentry.origin': 'auto.db.cloudflare.d1', + }, + name: 'INSERT INTO users (name) VALUES (?)', + op: 'db.query', + }, + expect.any(Function), + ); + }); + + test('instruments with breadcrumbs', async () => { + const instrumentedDb = instrumentD1WithSentry(createMockD1Database()); + await instrumentedDb.prepare('INSERT INTO users (name) VALUES (?)').all(); + + expect(addBreadcrumbSpy).toHaveBeenCalledTimes(1); + expect(addBreadcrumbSpy).toHaveBeenLastCalledWith({ + category: 'query', + message: 'INSERT INTO users (name) VALUES (?)', + data: { + 'cloudflare.d1.query_type': 'all', + 'cloudflare.d1.duration': 1, + 'cloudflare.d1.rows_read': 3, + 'cloudflare.d1.rows_written': 4, + }, + }); + }); + + test('works with statement.bind()', async () => { + const instrumentedDb = instrumentD1WithSentry(createMockD1Database()); + await instrumentedDb.prepare('SELECT * FROM users').bind().all(); + + expect(startSpanSpy).toHaveBeenCalledTimes(1); + expect(addBreadcrumbSpy).toHaveBeenCalledTimes(1); + }); + }); + + describe('statement.raw()', () => { + test('does not change return value', async () => { + const instrumentedDb = instrumentD1WithSentry(createMockD1Database()); + const response = await instrumentedDb.prepare('SELECT * FROM users').raw(); + expect(response).toEqual(MOCK_RAW_RETURN_VALUE); + }); + + test('instruments with spans', async () => { + const instrumentedDb = instrumentD1WithSentry(createMockD1Database()); + await instrumentedDb.prepare('SELECT * FROM users').raw(); + + expect(startSpanSpy).toHaveBeenCalledTimes(1); + expect(startSpanSpy).toHaveBeenLastCalledWith( + { + attributes: { + 'cloudflare.d1.query_type': 'raw', + 'sentry.origin': 'auto.db.cloudflare.d1', + }, + name: 'SELECT * FROM users', + op: 'db.query', + }, + expect.any(Function), + ); + }); + + test('instruments with breadcrumbs', async () => { + const instrumentedDb = instrumentD1WithSentry(createMockD1Database()); + await instrumentedDb.prepare('SELECT * FROM users').raw(); + + expect(addBreadcrumbSpy).toHaveBeenCalledTimes(1); + expect(addBreadcrumbSpy).toHaveBeenLastCalledWith({ + category: 'query', + message: 'SELECT * FROM users', + data: { + 'cloudflare.d1.query_type': 'raw', + }, + }); + }); + + test('works with statement.bind()', async () => { + const instrumentedDb = instrumentD1WithSentry(createMockD1Database()); + await instrumentedDb.prepare('SELECT * FROM users').bind().raw(); + + expect(startSpanSpy).toHaveBeenCalledTimes(1); + expect(addBreadcrumbSpy).toHaveBeenCalledTimes(1); + }); + }); +});