d1-sentry-patch.ts

import { startSpan, SEMANTIC_ATTRIBUTE_SENTRY_ORIGIN, SPAN_STATUS_ERROR } from '@sentry/core';

const patchedStatement = new WeakSet<D1PreparedStatement>();

function instrumentD1PreparedStatementQueries(statement: D1PreparedStatement, query: string): D1PreparedStatement {
	if (patchedStatement.has(statement)) {
		return statement;
	}

	statement.first = new Proxy(statement.first, {
		apply(target, thisArg, args: Parameters<typeof statement.first>) {
			return startSpan(
				{
					op: 'db.query',
					name: query,
					attributes: { 'cloudflare.d1.query_type': 'first', [SEMANTIC_ATTRIBUTE_SENTRY_ORIGIN]: 'auto.db.cloudflare.d1' },
				},
				() => {
					return target.apply(thisArg, args);
				}
			);
		},
	});

	statement.run = new Proxy(statement.run, {
		apply(target, thisArg, args: Parameters<typeof statement.run>) {
			return startSpan(
				{
					op: 'db.query',
					name: query,
					attributes: { 'cloudflare.d1.query_type': 'run', [SEMANTIC_ATTRIBUTE_SENTRY_ORIGIN]: 'auto.db.cloudflare.d1' },
				},
				async (span) => {
					const d1Result = await target.apply(thisArg, args);
					if (!d1Result.success) {
						span.setStatus({ code: SPAN_STATUS_ERROR });
					}

					const { duration, rows_read, rows_written } = d1Result.meta;
					span.setAttributes({
						'cloudflare.d1.query.duration': duration,
						'cloudflare.d1.query.rows_read': rows_read,
						'cloudflare.d1.query.rows_written': rows_written,
					});

					return d1Result;
				}
			);
		},
	});

	statement.all = new Proxy(statement.all, {
		apply(target, thisArg: D1PreparedStatement, args: Parameters<typeof statement.all>) {
			return startSpan(
				{
					op: 'db.query',
					name: query,
					attributes: { 'cloudflare.d1.query_type': 'all', [SEMANTIC_ATTRIBUTE_SENTRY_ORIGIN]: 'auto.db.cloudflare.d1' },
				},
				async (span) => {
					const d1Result = await target.apply(thisArg, args);
					if (!d1Result.success) {
						span.setStatus({ code: SPAN_STATUS_ERROR });
					}

					const { duration, rows_read, rows_written } = d1Result.meta;
					span.setAttributes({
						'cloudflare.d1.query.duration': duration,
						'cloudflare.d1.query.rows_read': rows_read,
						'cloudflare.d1.query.rows_written': rows_written,
					});

					return d1Result;
				}
			);
		},
	});

	statement.raw = new Proxy(statement.raw, {
		apply(target, thisArg, args: Parameters<typeof statement.raw>) {
			return startSpan(
				{
					op: 'db.query',
					name: query,
					attributes: { 'cloudflare.d1.query_type': 'raw', [SEMANTIC_ATTRIBUTE_SENTRY_ORIGIN]: 'auto.db.cloudflare.d1' },
				},
				() => {
					return target.apply(thisArg, args);
				}
			);
		},
	});

	patchedStatement.add(statement);

	return statement;
}

// https://github.com/cloudflare/workerd/blob/cd5279e7b305003f1d9c851e73efa9d67e4b68b2/src/cloudflare/internal/d1-api.ts#L131
function instrumentD1PreparedStatement(statement: D1PreparedStatement, query: string): D1PreparedStatement {
	statement.bind = new Proxy(statement.bind, {
		apply(target, thisArg, args: Parameters<typeof statement.bind>) {
			return instrumentD1PreparedStatementQueries(target.apply(thisArg, args), query);
		},
	});

	return instrumentD1PreparedStatementQueries(statement, query);
}

function instrumentDb(db: D1Database): D1Database {
	db.prepare = new Proxy(db.prepare, {
		apply(target, thisArg, args: Parameters<typeof db.prepare>) {
			const [query] = args;
			const res = target.apply(thisArg, args);
			return instrumentD1PreparedStatement(res, query);
		},
	});

	// TODO: instrument db.exec
	// db.exec = new Proxy(db.exec, {
	// 	apply(target, thisArg, args: Parameters<typeof db.exec>) {
	// 		// TODO: this is a raw query. What do I name it?
	// 		const [query] = args;
	// 		return startSpan({ op: 'db.query.exec', name: query }, async (span) => {
	// 			const d1Result = await target.apply(thisArg, args);

	// 			const { duration } = d1Result;
	// 			span.setAttributes({
	// 				'cloudflare.d1.query.duration': duration,
	// 			});

	// 			return d1Result;
	// 		});
	// 	},
	// });

	// TODO: instrument db.batch

	return db;
}