fix: instrument all execution methods

Author: logaretm

packages/nuxt/src/runtime/plugins/database.server.ts

@@ -1,18 +1,24 @@
 import {
-  type SpanAttributes,
+  type StartSpanOptions,
+  addBreadcrumb,
   captureException,
   debug,
   flushIfServerless,
-  SEMANTIC_ATTRIBUTE_SENTRY_OP,
   SEMANTIC_ATTRIBUTE_SENTRY_ORIGIN,
   SPAN_STATUS_ERROR,
-  SPAN_STATUS_OK,
   startSpan,
 } from '@sentry/core';
-import type { Database } from 'db0';
+import type { Database, PreparedStatement } from 'db0';
 // eslint-disable-next-line import/no-extraneous-dependencies
 import { defineNitroPlugin, useDatabase } from 'nitropack/runtime';
 
+type PreparedStatementType = 'get' | 'run' | 'all' | 'raw';
+
+/**
+ * Keeps track of prepared statements that have been patched.
+ */
+const patchedStatement = new WeakSet<PreparedStatement>();
+
 /**
  * Creates a Nitro plugin that instruments the database calls.
  */
@@ -27,52 +33,151 @@ export default defineNitroPlugin(() => {
 });
 
 function instrumentDatabase(db: Database): void {
+  db.prepare = new Proxy(db.prepare, {
+    apply(target, thisArg, args: Parameters<typeof db.prepare>) {
+      const [query] = args;
+
+      return instrumentPreparedStatement(target.apply(thisArg, args), query, db.dialect);
+    },
+  });
+
+  // Sadly the `.sql` template tag doesn't call `db.prepare` internally and it calls the connector's `.prepare` directly
+  // So we have to patch it manually, and would mean we would have less info in the spans.
+  // https://github.com/unjs/db0/blob/main/src/database.ts#L64
   db.sql = new Proxy(db.sql, {
     apply(target, thisArg, args: Parameters<typeof db.sql>) {
-      const query = args[0]?.[0];
-      const attributes = getSpanAttributes(db, query);
-
-      return startSpan(
-        {
-          name: query || 'db.query',
-          attributes,
-        },
-        async span => {
-          try {
-            const result = await target.apply(thisArg, args);
-            span.setStatus({ code: SPAN_STATUS_OK });
-
-            return result;
-          } catch (error) {
-            span.setStatus({ code: SPAN_STATUS_ERROR, message: 'internal_error' });
-            captureException(error, {
-              mechanism: {
-                handled: false,
-                type: attributes[SEMANTIC_ATTRIBUTE_SENTRY_ORIGIN],
-              },
-            });
-
-            // Re-throw the error to be handled by the caller
-            throw error;
-          } finally {
-            await flushIfServerless();
-          }
-        },
-      );
+      const query = args[0]?.[0] ?? '';
+      const opts = createStartSpanOptions(query, db.dialect);
+
+      return startSpan(opts, async span => {
+        try {
+          const result = await target.apply(thisArg, args);
+
+          return result;
+        } catch (error) {
+          span.setStatus({ code: SPAN_STATUS_ERROR, message: 'internal_error' });
+          captureException(error, {
+            mechanism: {
+              handled: false,
+              type: opts.attributes?.[SEMANTIC_ATTRIBUTE_SENTRY_ORIGIN],
+            },
+          });
+
+          // Re-throw the error to be handled by the caller
+          throw error;
+        } finally {
+          await flushIfServerless();
+        }
+      });
+    },
+  });
+
+  db.exec = new Proxy(db.exec, {
+    apply(target, thisArg, args: Parameters<typeof db.exec>) {
+      return startSpan(createStartSpanOptions(args[0], db.dialect, 'run'), async () => {
+        const result = await target.apply(thisArg, args);
+
+        createBreadcrumb(args[0], 'run');
+
+        return result;
+      });
     },
   });
 }
 
-function getSpanAttributes(db: Database, query?: string): SpanAttributes {
-  const attributes: SpanAttributes = {
-    [SEMANTIC_ATTRIBUTE_SENTRY_ORIGIN]: 'auto.db.nuxt',
-    [SEMANTIC_ATTRIBUTE_SENTRY_OP]: 'db',
-    'db.system': db.dialect,
-  };
+/**
+ * Instruments a DB prepared statement with Sentry.
+ *
+ * This is meant to be used as a top-level call, under the hood it calls `instrumentPreparedStatementQueries`
+ * to patch the query methods. The reason for this abstraction is to ensure that the `bind` method is also patched.
+ */
+function instrumentPreparedStatement(statement: PreparedStatement, query: string, dialect: string): PreparedStatement {
+  // 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<typeof statement.bind>) {
+      return instrumentPreparedStatementQueries(target.apply(thisArg, args), query, dialect);
+    },
+  });
 
-  if (query) {
-    attributes['db.query'] = query;
+  return instrumentPreparedStatementQueries(statement, query, dialect);
+}
+
+/**
+ * Patches the query methods of a DB prepared statement with Sentry.
+ */
+function instrumentPreparedStatementQueries(
+  statement: PreparedStatement,
+  query: string,
+  dialect: string,
+): PreparedStatement {
+  if (patchedStatement.has(statement)) {
+    return statement;
   }
 
-  return attributes;
+  // eslint-disable-next-line @typescript-eslint/unbound-method
+  statement.get = new Proxy(statement.get, {
+    apply(target, thisArg, args: Parameters<typeof statement.get>) {
+      return startSpan(createStartSpanOptions(query, dialect, 'get'), async () => {
+        const result = await target.apply(thisArg, args);
+        createBreadcrumb(query, 'get');
+
+        return result;
+      });
+    },
+  });
+
+  // eslint-disable-next-line @typescript-eslint/unbound-method
+  statement.run = new Proxy(statement.run, {
+    apply(target, thisArg, args: Parameters<typeof statement.run>) {
+      return startSpan(createStartSpanOptions(query, dialect, 'run'), async () => {
+        const result = await target.apply(thisArg, args);
+        createBreadcrumb(query, 'run');
+
+        return result;
+      });
+    },
+  });
+
+  // eslint-disable-next-line @typescript-eslint/unbound-method
+  statement.all = new Proxy(statement.all, {
+    apply(target, thisArg, args: Parameters<typeof statement.all>) {
+      return startSpan(createStartSpanOptions(query, dialect, 'all'), async () => {
+        const result = await target.apply(thisArg, args);
+        // Since all has no regular shape, we can assume if it returns an array, it's a success.
+        createBreadcrumb(query, 'all');
+
+        return result;
+      });
+    },
+  });
+
+  patchedStatement.add(statement);
+
+  return statement;
+}
+
+function createBreadcrumb(query: string, type: PreparedStatementType): void {
+  addBreadcrumb({
+    category: 'query',
+    message: query,
+    data: {
+      'db.query_type': type,
+    },
+  });
+}
+
+/**
+ * Creates a start span options object.
+ */
+function createStartSpanOptions(query: string, dialect: string, type?: PreparedStatementType): StartSpanOptions {
+  return {
+    op: 'db.query',
+    name: query,
+    attributes: {
+      'db.system': dialect,
+      'db.query_type': type,
+      [SEMANTIC_ATTRIBUTE_SENTRY_ORIGIN]: 'auto.db.nuxt',
+    },
+  };
 }

packages/nuxt/src/vite/databaseConfig.ts

@@ -6,7 +6,7 @@ import type { Nuxt } from 'nuxt/schema';
  * Sets up the database instrumentation.
  */
 export function addDatabaseInstrumentation(nuxt: Nuxt): void {
-  if (!nuxt.options.nitro?.database) {
+  if (!nuxt.options.nitro?.experimental?.database && !nuxt.options.nitro?.database) {
     consoleSandbox(() => {
       // eslint-disable-next-line no-console
       console.log('[Sentry] No database configuration found. Skipping database instrumentation.');