Skip to content

Commit dc4ec0c

Browse files
JPeer264claude
JPeer264
and
claude
committed
feat(cloudflare): Add batch, exec, and withSession D1 instrumentation
Instrument `db.batch()`, `db.exec()`, and `db.withSession()` methods that were previously not covered by D1 instrumentation. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
1 parent 59c5ec0 commit dc4ec0c

5 files changed

Lines changed: 427 additions & 27 deletions

File tree

‎dev-packages/cloudflare-integration-tests/suites/d1/index.ts‎

Lines changed: 22 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,11 @@ export default Sentry.withSentry(
2222
return new Response('ok');
2323
}
2424

25+
if (url.pathname === '/exec') {
26+
await env.DB.exec('CREATE TABLE IF NOT EXISTS users (id INTEGER PRIMARY KEY, name TEXT)');
27+
return new Response('ok');
28+
}
29+
2530
if (url.pathname === '/double-instrument') {
2631
const prepareBeforeManual = env.DB.prepare;
2732
const db = Sentry.instrumentD1WithSentry(env.DB);
@@ -38,6 +43,23 @@ export default Sentry.withSentry(
3843
return new Response('ok');
3944
}
4045

46+
if (url.pathname === '/batch') {
47+
await env.DB.batch([
48+
env.DB.prepare('INSERT INTO users (name) VALUES (?)').bind('Alice'),
49+
env.DB.prepare('INSERT INTO users (name) VALUES (?)').bind('Bob'),
50+
]);
51+
return new Response('ok');
52+
}
53+
54+
if (url.pathname === '/with-session/batch') {
55+
const session = env.DB.withSession();
56+
await session.batch([
57+
session.prepare('INSERT INTO users (name) VALUES (?)').bind('Alice'),
58+
session.prepare('INSERT INTO users (name) VALUES (?)').bind('Bob'),
59+
]);
60+
return new Response('ok');
61+
}
62+
4163
return new Response('not found', { status: 404 });
4264
},
4365
} satisfies ExportedHandler<Env>,

‎dev-packages/cloudflare-integration-tests/suites/d1/test.ts‎

Lines changed: 112 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -29,8 +29,10 @@ it('instruments D1 prepare().all() automatically via env', async ({ signal }) =>
2929
expect(querySpan).toBeDefined();
3030
expect(querySpan).toEqual({
3131
data: {
32+
'db.system.name': 'cloudflare-d1',
33+
'db.operation.name': 'all',
34+
'db.query.text': 'SELECT * FROM users WHERE id = ?',
3235
'cloudflare.d1.duration': expect.any(Number),
33-
'cloudflare.d1.query_type': 'all',
3436
'cloudflare.d1.rows_read': expect.any(Number),
3537
'cloudflare.d1.rows_written': expect.any(Number),
3638
'sentry.op': 'db.query',
@@ -94,6 +96,41 @@ it('captures error event when a D1 query references a non-existent table', async
9496
await runner.completed();
9597
});
9698

99+
it('instruments D1 exec() automatically via env', async ({ signal }) => {
100+
const runner = createRunner(__dirname)
101+
.ignore('event')
102+
.expect((envelope: Envelope) => {
103+
if (envelopeItemType(envelope) !== 'transaction') return;
104+
const d1Spans = findD1Spans(envelope);
105+
106+
const execSpan = d1Spans.find(
107+
s => s.description === 'CREATE TABLE IF NOT EXISTS users (id INTEGER PRIMARY KEY, name TEXT)',
108+
);
109+
expect(execSpan).toBeDefined();
110+
expect(execSpan).toEqual({
111+
data: {
112+
'db.system.name': 'cloudflare-d1',
113+
'db.operation.name': 'exec',
114+
'db.query.text': 'CREATE TABLE IF NOT EXISTS users (id INTEGER PRIMARY KEY, name TEXT)',
115+
'sentry.op': 'db.query',
116+
'sentry.origin': 'auto.db.cloudflare.d1',
117+
},
118+
description: 'CREATE TABLE IF NOT EXISTS users (id INTEGER PRIMARY KEY, name TEXT)',
119+
op: 'db.query',
120+
origin: 'auto.db.cloudflare.d1',
121+
parent_span_id: expect.any(String),
122+
span_id: expect.any(String),
123+
start_timestamp: expect.any(Number),
124+
timestamp: expect.any(Number),
125+
trace_id: expect.any(String),
126+
});
127+
})
128+
.start(signal);
129+
130+
await runner.makeRequest('get', '/exec');
131+
await runner.completed();
132+
});
133+
97134
it('does not double-instrument when instrumentD1WithSentry is used on top of env instrumentation', async ({
98135
signal,
99136
}) => {
@@ -112,3 +149,77 @@ it('does not double-instrument when instrumentD1WithSentry is used on top of env
112149
expect(response).toBe('true');
113150
await runner.completed();
114151
});
152+
153+
it('instruments D1 withSession().batch() identically to db.batch()', async ({ signal }) => {
154+
let directBatchSpan: Record<string, unknown> | undefined;
155+
let sessionBatchSpan: Record<string, unknown> | undefined;
156+
157+
const runner = createRunner(__dirname)
158+
.ignore('event')
159+
.expect((envelope: Envelope) => {
160+
expect(envelopeItem(envelope).transaction).toBe('GET /batch');
161+
162+
directBatchSpan = findD1Spans(envelope).find(s => s.description === 'D1 batch');
163+
})
164+
.expect((envelope: Envelope) => {
165+
expect(envelopeItem(envelope).transaction).toBe('GET /with-session/batch');
166+
167+
sessionBatchSpan = findD1Spans(envelope).find(s => s.description === 'D1 batch');
168+
})
169+
.start(signal);
170+
171+
await runner.makeRequest('get', '/batch');
172+
await runner.makeRequest('get', '/with-session/batch');
173+
await runner.completed();
174+
175+
expect(directBatchSpan).toBeDefined();
176+
expect(sessionBatchSpan).toBeDefined();
177+
178+
const normalize = (span: Record<string, unknown>): Record<string, unknown> => {
179+
const {
180+
span_id: _spanId,
181+
parent_span_id: _parentSpanId,
182+
start_timestamp: _start,
183+
timestamp: _end,
184+
trace_id: _traceId,
185+
...rest
186+
} = span;
187+
return rest;
188+
};
189+
190+
expect(normalize(sessionBatchSpan!)).toEqual(normalize(directBatchSpan!));
191+
});
192+
193+
it('instruments D1 batch() automatically via env', async ({ signal }) => {
194+
const runner = createRunner(__dirname)
195+
.ignore('event')
196+
.expect((envelope: Envelope) => {
197+
if (envelopeItemType(envelope) !== 'transaction') return;
198+
const d1Spans = findD1Spans(envelope);
199+
200+
const batchSpan = d1Spans.find(s => s.description === 'D1 batch');
201+
expect(batchSpan).toBeDefined();
202+
expect(batchSpan).toEqual({
203+
data: {
204+
'db.system.name': 'cloudflare-d1',
205+
'db.operation.name': 'batch',
206+
'db.query.text': 'INSERT INTO users (name) VALUES (?)\nINSERT INTO users (name) VALUES (?)',
207+
'db.operation.batch.size': 2,
208+
'sentry.op': 'db.query',
209+
'sentry.origin': 'auto.db.cloudflare.d1',
210+
},
211+
description: 'D1 batch',
212+
op: 'db.query',
213+
origin: 'auto.db.cloudflare.d1',
214+
parent_span_id: expect.any(String),
215+
span_id: expect.any(String),
216+
start_timestamp: expect.any(Number),
217+
timestamp: expect.any(Number),
218+
trace_id: expect.any(String),
219+
});
220+
})
221+
.start(signal);
222+
223+
await runner.makeRequest('get', '/batch');
224+
await runner.completed();
225+
});

‎dev-packages/cloudflare-integration-tests/suites/tracing/d1/test.ts‎

Lines changed: 23 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -15,7 +15,26 @@ it('D1 database queries create spans with correct attributes', async ({ signal }
1515
data: {
1616
[SEMANTIC_ATTRIBUTE_SENTRY_OP]: 'db.query',
1717
[SEMANTIC_ATTRIBUTE_SENTRY_ORIGIN]: 'auto.db.cloudflare.d1',
18-
'cloudflare.d1.query_type': 'run',
18+
'db.system.name': 'cloudflare-d1',
19+
'db.operation.name': 'exec',
20+
'db.query.text': 'CREATE TABLE IF NOT EXISTS users (id INTEGER PRIMARY KEY, name TEXT)',
21+
},
22+
description: 'CREATE TABLE IF NOT EXISTS users (id INTEGER PRIMARY KEY, name TEXT)',
23+
op: 'db.query',
24+
origin: 'auto.db.cloudflare.d1',
25+
parent_span_id: expect.any(String),
26+
span_id: expect.any(String),
27+
start_timestamp: expect.any(Number),
28+
timestamp: expect.any(Number),
29+
trace_id: expect.any(String),
30+
},
31+
{
32+
data: {
33+
[SEMANTIC_ATTRIBUTE_SENTRY_OP]: 'db.query',
34+
[SEMANTIC_ATTRIBUTE_SENTRY_ORIGIN]: 'auto.db.cloudflare.d1',
35+
'db.system.name': 'cloudflare-d1',
36+
'db.operation.name': 'run',
37+
'db.query.text': 'INSERT INTO users (name) VALUES (?)',
1938
'cloudflare.d1.duration': expect.any(Number),
2039
'cloudflare.d1.rows_read': expect.any(Number),
2140
'cloudflare.d1.rows_written': expect.any(Number),
@@ -44,7 +63,9 @@ it('D1 database queries create spans with correct attributes', async ({ signal }
4463
data: {
4564
[SEMANTIC_ATTRIBUTE_SENTRY_OP]: 'db.query',
4665
[SEMANTIC_ATTRIBUTE_SENTRY_ORIGIN]: 'auto.db.cloudflare.d1',
47-
'cloudflare.d1.query_type': 'first',
66+
'db.system.name': 'cloudflare-d1',
67+
'db.operation.name': 'first',
68+
'db.query.text': 'SELECT * FROM users WHERE name = ?',
4869
},
4970
description: 'SELECT * FROM users WHERE name = ?',
5071
op: 'db.query',

‎packages/cloudflare/src/instrumentations/worker/instrumentD1.ts‎

Lines changed: 76 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
/* eslint-disable @typescript-eslint/unbound-method */
2-
import type { D1Database, D1PreparedStatement, D1Response } from '@cloudflare/workers-types';
2+
import type { D1Database, D1DatabaseSession, D1PreparedStatement, D1Response } from '@cloudflare/workers-types';
33
import type { Span, SpanAttributes, StartSpanOptions } from '@sentry/core';
44
import { addBreadcrumb, SEMANTIC_ATTRIBUTE_SENTRY_ORIGIN, SPAN_STATUS_ERROR, startSpan } from '@sentry/core';
55
import { ensureInstrumented } from '../../instrument';
@@ -107,35 +107,103 @@ function getAttributesFromD1Response(d1Result: D1Response): SpanAttributes {
107107
};
108108
}
109109

110-
function createD1Breadcrumb(query: string, type: 'first' | 'run' | 'all' | 'raw', d1Result?: D1Response): void {
110+
type D1QueryType = 'first' | 'run' | 'all' | 'raw' | 'batch' | 'exec';
111+
112+
function createD1Breadcrumb(query: string, type: D1QueryType, d1Result?: D1Response): void {
111113
addBreadcrumb({
112114
category: 'query',
113115
message: query,
114116
data: {
115117
...(d1Result ? getAttributesFromD1Response(d1Result) : {}),
116-
'cloudflare.d1.query_type': type,
118+
'db.operation.name': type,
117119
},
118120
});
119121
}
120122

121-
function createStartSpanOptions(query: string, type: 'first' | 'run' | 'all' | 'raw'): StartSpanOptions {
123+
function createStartSpanOptions(query: string, type: D1QueryType): StartSpanOptions {
122124
return {
123125
op: 'db.query',
124126
name: query,
125127
attributes: {
126-
'cloudflare.d1.query_type': type,
128+
'db.system.name': 'cloudflare-d1',
129+
'db.operation.name': type,
130+
'db.query.text': query,
127131
[SEMANTIC_ATTRIBUTE_SENTRY_ORIGIN]: 'auto.db.cloudflare.d1',
128132
},
129133
};
130134
}
131135

132-
function _instrumentD1(db: D1Database): D1Database {
133-
db.prepare = new Proxy(db.prepare, {
134-
apply(target, thisArg, args: Parameters<typeof db.prepare>) {
136+
function instrumentPrepare(
137+
prepare: D1Database['prepare'] | D1DatabaseSession['prepare'],
138+
): D1Database['prepare'] | D1DatabaseSession['prepare'] {
139+
return new Proxy(prepare, {
140+
apply(target, thisArg, args: Parameters<typeof prepare>) {
135141
const [query] = args;
136142
return instrumentD1PreparedStatement(Reflect.apply(target, thisArg, args), query);
137143
},
138144
});
145+
}
146+
147+
function instrumentBatch(
148+
batch: D1Database['batch'] | D1DatabaseSession['batch'],
149+
): D1Database['batch'] | D1DatabaseSession['batch'] {
150+
return new Proxy(batch, {
151+
apply(target, thisArg, args: Parameters<typeof batch>) {
152+
const statements = args[0];
153+
// D1PreparedStatement exposes a `statement` property at runtime, but it's not in @cloudflare/workers-types.
154+
// https://github.com/cloudflare/workerd/blob/dc12d7650b4f5d4f9ba6a47aa45fad769cdf8db4/src/cloudflare/internal/d1-api.ts#L210
155+
const queryText = statements.map((statement) => (statement as unknown as { statement?: string }).statement ?? '').join('\n');
156+
157+
return startSpan(
158+
{
159+
op: 'db.query',
160+
name: 'D1 batch',
161+
attributes: {
162+
'db.system.name': 'cloudflare-d1',
163+
'db.operation.name': 'batch',
164+
'db.query.text': queryText || undefined,
165+
'db.operation.batch.size': statements.length,
166+
[SEMANTIC_ATTRIBUTE_SENTRY_ORIGIN]: 'auto.db.cloudflare.d1',
167+
},
168+
},
169+
async () => {
170+
const res = await Reflect.apply(target, thisArg, args);
171+
createD1Breadcrumb('D1 batch', 'batch');
172+
return res;
173+
},
174+
);
175+
},
176+
});
177+
}
178+
179+
function instrumentD1Session(session: D1DatabaseSession): D1DatabaseSession {
180+
session.prepare = instrumentPrepare(session.prepare);
181+
session.batch = instrumentBatch(session.batch);
182+
return session;
183+
}
184+
185+
function _instrumentD1(db: D1Database): D1Database {
186+
db.prepare = instrumentPrepare(db.prepare);
187+
db.batch = instrumentBatch(db.batch);
188+
189+
db.exec = new Proxy(db.exec, {
190+
apply(target, thisArg, args: Parameters<typeof db.exec>) {
191+
const [query] = args;
192+
return startSpan(createStartSpanOptions(query, 'exec'), async () => {
193+
const res = await Reflect.apply(target, thisArg, args);
194+
createD1Breadcrumb(query, 'exec');
195+
return res;
196+
});
197+
},
198+
});
199+
200+
if ('withSession' in db && typeof db.withSession === 'function') {
201+
db.withSession = new Proxy(db.withSession, {
202+
apply(target, thisArg, args: [unknown]) {
203+
return instrumentD1Session(Reflect.apply(target, thisArg, args) as D1DatabaseSession);
204+
},
205+
});
206+
}
139207

140208
return db;
141209
}

0 commit comments

Comments
 (0)