ref(node): Streamline dataloader instrumentation

.oxlintrc.base.json

@@ -142,7 +142,6 @@
     },
     {
       "files": [
-        "**/integrations/tracing/dataloader/vendored/**/*.ts",
         "**/integrations/fs/vendored/**/*.ts",
         "**/integrations/tracing/knex/vendored/**/*.ts",
         "**/integrations/tracing/mongo/vendored/**/*.ts",

dev-packages/node-integration-tests/suites/tracing/dataloader/scenario.mjs

@@ -4,15 +4,35 @@ import express from 'express';
 
 const PORT = 8008;
 
+const batchLoadFn = async keys => keys.map((_, idx) => idx);
+
 const run = async () => {
   const app = express();
-  const dataloader = new Dataloader(async keys => keys.map((_, idx) => idx), {
-    cache: false,
+
+  app.get('/load', async (_req, res) => {
+    const dataloader = new Dataloader(batchLoadFn, { cache: false });
+    const user = await dataloader.load('user-1');
+    res.send({ user });
+  });
+
+  app.get('/load-many', async (_req, res) => {
+    const dataloader = new Dataloader(batchLoadFn, { cache: false });
+    const users = await dataloader.loadMany(['user-1', 'user-2']);
+    res.send({ users });
+  });
+
+  app.get('/cache-ops', async (_req, res) => {
+    const dataloader = new Dataloader(batchLoadFn);
+    dataloader.prime('user-1', 1);
+    dataloader.clear('user-1');
+    dataloader.clearAll();
+    res.send({});
   });
 
-  app.get('/', (req, res) => {
-    const user = dataloader.load('user-1');
-    res.send(user);
+  app.get('/named', async (_req, res) => {
+    const dataloader = new Dataloader(batchLoadFn, { cache: false, name: 'usersLoader' });
+    const user = await dataloader.load('user-1');
+    res.send({ user });
   });
 
   startExpressServerAndSendPortToRunner(app, PORT);

dev-packages/node-integration-tests/suites/tracing/dataloader/test.ts

@@ -1,42 +1,96 @@
 import { afterAll, describe, expect } from 'vitest';
 import { cleanupChildProcesses, createEsmAndCjsTests } from '../../../utils/runner';
 
+const ORIGIN = 'auto.db.otel.dataloader';
+const CACHE_GET_OP = 'cache.get';
+
 describe('dataloader auto-instrumentation', () => {
-  afterAll(async () => {
+  afterAll(() => {
     cleanupChildProcesses();
   });
 
-  const EXPECTED_TRANSACTION = {
-    transaction: 'GET /',
-    spans: expect.arrayContaining([
-      expect.objectContaining({
-        data: expect.objectContaining({
-          'sentry.origin': 'auto.db.otel.dataloader',
-          'sentry.op': 'cache.get',
-        }),
-        description: 'dataloader.load',
-        origin: 'auto.db.otel.dataloader',
-        op: 'cache.get',
-        status: 'ok',
-      }),
-      expect.objectContaining({
-        data: expect.objectContaining({
-          'sentry.origin': 'auto.db.otel.dataloader',
-          'sentry.op': 'cache.get',
-        }),
-        description: 'dataloader.batch',
-        origin: 'auto.db.otel.dataloader',
-        op: 'cache.get',
-        status: 'ok',
-      }),
-    ]),
-  };
-
   createEsmAndCjsTests(__dirname, 'scenario.mjs', 'instrument.mjs', (createRunner, test) => {
-    test('should auto-instrument `dataloader` package.', async () => {
-      const runner = createRunner().expect({ transaction: EXPECTED_TRANSACTION }).start();
-      runner.makeRequest('get', '/');
+    test('instruments load, loadMany, batch, prime, clear and clearAll', async () => {
+      const runner = createRunner()
+        .expect({
+          transaction: event => {
+            expect(event.transaction).toBe('GET /load');
+
+            const spans = event.spans || [];
+
+            const loadSpan = spans.find(span => span.description === 'dataloader.load');
+            expect(loadSpan).toBeDefined();
+            expect(loadSpan?.op).toBe(CACHE_GET_OP);
+            expect(loadSpan?.origin).toBe(ORIGIN);
+            expect(loadSpan?.status).toBe('ok');
+            expect(loadSpan?.data?.['sentry.origin']).toBe(ORIGIN);
+            expect(loadSpan?.data?.['sentry.op']).toBe(CACHE_GET_OP);
+
+            const batchSpan = spans.find(span => span.description === 'dataloader.batch');
+            expect(batchSpan).toBeDefined();
+            expect(batchSpan?.op).toBe(CACHE_GET_OP);
+            expect(batchSpan?.origin).toBe(ORIGIN);
+            expect(batchSpan?.status).toBe('ok');
+
+            // The batch span links back to the load span that triggered it
+            expect(batchSpan?.links).toEqual([
+              expect.objectContaining({
+                trace_id: loadSpan?.trace_id,
+                span_id: loadSpan?.span_id,
+              }),
+            ]);
+          },
+        })
+        .expect({
+          transaction: event => {
+            expect(event.transaction).toBe('GET /load-many');
+
+            const loadManySpan = (event.spans || []).find(span => span.description === 'dataloader.loadMany');
+            expect(loadManySpan).toBeDefined();
+            expect(loadManySpan?.op).toBe(CACHE_GET_OP);
+            expect(loadManySpan?.origin).toBe(ORIGIN);
+            expect(loadManySpan?.status).toBe('ok');
+            expect(loadManySpan?.data?.['sentry.origin']).toBe(ORIGIN);
+            expect(loadManySpan?.data?.['sentry.op']).toBe(CACHE_GET_OP);
+          },
+        })
+        .expect({
+          transaction: event => {
+            expect(event.transaction).toBe('GET /cache-ops');
+
+            const spans = event.spans || [];
+
+            // prime/clear/clearAll are not cache reads, so they get an origin but no `op`
+            for (const operation of ['prime', 'clear', 'clearAll']) {
+              const span = spans.find(s => s.description === `dataloader.${operation}`);
+              expect(span, `expected a dataloader.${operation} span`).toBeDefined();
+              expect(span?.origin).toBe(ORIGIN);
+              expect(span?.status).toBe('ok');
+              expect(span?.op).toBeUndefined();
+              expect(span?.data?.['sentry.origin']).toBe(ORIGIN);
+              expect(span?.data?.['sentry.op']).toBeUndefined();
+            }
+          },
+        })
+        .expect({
+          transaction: event => {
+            expect(event.transaction).toBe('GET /named');
+
+            // A named dataloader includes its name in the span description
+            const namedLoadSpan = (event.spans || []).find(span => span.description === 'dataloader.load usersLoader');
+            expect(namedLoadSpan).toBeDefined();
+            expect(namedLoadSpan?.op).toBe(CACHE_GET_OP);
+            expect(namedLoadSpan?.origin).toBe(ORIGIN);
+            expect(namedLoadSpan?.status).toBe('ok');
+          },
+        })
+        .start();
+
+      await runner.makeRequest('get', '/load');
+      await runner.makeRequest('get', '/load-many');
+      await runner.makeRequest('get', '/cache-ops');
+      await runner.makeRequest('get', '/named');
       await runner.completed();
-    });
+    }, 30_000);
   });
 });

packages/node/src/integrations/tracing/dataloader/index.ts

@@ -1,56 +1,17 @@
 import { DataloaderInstrumentation } from './vendored/instrumentation';
 import type { IntegrationFn } from '@sentry/core';
-import {
-  defineIntegration,
-  SEMANTIC_ATTRIBUTE_SENTRY_OP,
-  SEMANTIC_ATTRIBUTE_SENTRY_ORIGIN,
-  spanToJSON,
-} from '@sentry/core';
-import { generateInstrumentOnce, instrumentWhenWrapped } from '@sentry/node-core';
+import { defineIntegration } from '@sentry/core';
+import { generateInstrumentOnce } from '@sentry/node-core';
 
 const INTEGRATION_NAME = 'Dataloader';
 
-export const instrumentDataloader = generateInstrumentOnce(
-  INTEGRATION_NAME,
-  () =>
-    new DataloaderInstrumentation({
-      requireParentSpan: true,
-    }),
-);
+export const instrumentDataloader = generateInstrumentOnce(INTEGRATION_NAME, () => new DataloaderInstrumentation());
 
 const _dataloaderIntegration = (() => {
-  let instrumentationWrappedCallback: undefined | ((callback: () => void) => void);
-
   return {
     name: INTEGRATION_NAME,
     setupOnce() {
-      const instrumentation = instrumentDataloader();
-      instrumentationWrappedCallback = instrumentWhenWrapped(instrumentation);
-    },
-
-    setup(client) {
-      // This is called either immediately or when the instrumentation is wrapped
-      instrumentationWrappedCallback?.(() => {
-        client.on('spanStart', span => {
-          const spanJSON = spanToJSON(span);
-          if (spanJSON.description?.startsWith('dataloader')) {
-            span.setAttribute(SEMANTIC_ATTRIBUTE_SENTRY_ORIGIN, 'auto.db.otel.dataloader');
-          }
-
-          // These are all possible dataloader span descriptions
-          // Still checking for the future versions
-          // in case they add support for `clear` and `prime`
-          if (
-            spanJSON.description === 'dataloader.load' ||
-            spanJSON.description === 'dataloader.loadMany' ||
-            spanJSON.description === 'dataloader.batch'
-          ) {
-            span.setAttribute(SEMANTIC_ATTRIBUTE_SENTRY_OP, 'cache.get');
-            // TODO: We can try adding `key` to the `data` attribute upstream.
-            // Or alternatively, we can add `requestHook` to the dataloader instrumentation.
-          }
-        });
-      });
+      instrumentDataloader();
     },
   };
 }) satisfies IntegrationFn;