Add requireParentSpan, requestHook and origin

onurtemizkan · onurtemizkan · commit edcc0164f7e7 · 2025-06-25T13:15:50.000+01:00
diff --git a/dev-packages/node-integration-tests/suites/tracing/postgresjs/test.ts b/dev-packages/node-integration-tests/suites/tracing/postgresjs/test.ts
@@ -11,15 +11,15 @@ describe('postgresjs auto instrumentation', () => {
             'db.namespace': 'test_db',
             'db.system.name': 'postgres',
             'sentry.op': 'db',
-            'sentry.origin': 'manual',
+            'sentry.origin': 'auto.db.otel.postgres',
             'server.address': 'localhost',
             'server.port': 5444,
           }),
           description:
             'CREATE TABLE "User" ("id" SERIAL NOT NULL,"createdAt" TIMESTAMP(?) NOT NULL DEFAULT CURRENT_TIMESTAMP,"email" TEXT NOT NULL,"name" TEXT,CONSTRAINT "User_pkey" PRIMARY KEY ("id"))',
           op: 'db',
           status: 'ok',
-          origin: 'manual',
+          origin: 'auto.db.otel.postgres',
           parent_span_id: expect.any(String),
           span_id: expect.any(String),
           start_timestamp: expect.any(Number),
@@ -31,15 +31,15 @@ describe('postgresjs auto instrumentation', () => {
             'db.namespace': 'test_db',
             'db.system.name': 'postgres',
             'sentry.op': 'db',
-            'sentry.origin': 'manual',
+            'sentry.origin': 'auto.db.otel.postgres',
             'server.address': 'localhost',
             'server.port': 5444,
           }),
           description:
             "select b.oid, b.typarray from pg_catalog.pg_type a left join pg_catalog.pg_type b on b.oid = a.typelem where a.typcategory = 'A' group by b.oid, b.typarray order by b.oid",
           op: 'db',
           status: 'ok',
-          origin: 'manual',
+          origin: 'auto.db.otel.postgres',
           parent_span_id: expect.any(String),
           span_id: expect.any(String),
           start_timestamp: expect.any(Number),
@@ -50,15 +50,15 @@ describe('postgresjs auto instrumentation', () => {
           data: expect.objectContaining({
             'db.namespace': 'test_db',
             'db.system.name': 'postgres',
-            'sentry.origin': 'manual',
+            'sentry.origin': 'auto.db.otel.postgres',
             'sentry.op': 'db',
             'server.address': 'localhost',
             'server.port': 5444,
           }),
           description: 'INSERT INTO "User" ("email", "name") VALUES (\'Foo\', \'bar@baz.com\')',
           op: 'db',
           status: 'ok',
-          origin: 'manual',
+          origin: 'auto.db.otel.postgres',
           parent_span_id: expect.any(String),
           span_id: expect.any(String),
           start_timestamp: expect.any(Number),
@@ -70,14 +70,14 @@ describe('postgresjs auto instrumentation', () => {
             'db.namespace': 'test_db',
             'db.system.name': 'postgres',
             'sentry.op': 'db',
-            'sentry.origin': 'manual',
+            'sentry.origin': 'auto.db.otel.postgres',
             'server.address': 'localhost',
             'server.port': 5444,
           }),
           description: 'UPDATE "User" SET "name" = \'Foo\' WHERE "email" = \'bar@baz.com\'',
           op: 'db',
           status: 'ok',
-          origin: 'manual',
+          origin: 'auto.db.otel.postgres',
           parent_span_id: expect.any(String),
           span_id: expect.any(String),
           start_timestamp: expect.any(Number),
@@ -89,14 +89,14 @@ describe('postgresjs auto instrumentation', () => {
             'db.namespace': 'test_db',
             'db.system.name': 'postgres',
             'sentry.op': 'db',
-            'sentry.origin': 'manual',
+            'sentry.origin': 'auto.db.otel.postgres',
             'server.address': 'localhost',
             'server.port': 5444,
           }),
           description: 'SELECT * FROM "User" WHERE "email" = \'bar@baz.com\'',
           op: 'db',
           status: 'ok',
-          origin: 'manual',
+          origin: 'auto.db.otel.postgres',
           parent_span_id: expect.any(String),
           span_id: expect.any(String),
           start_timestamp: expect.any(Number),
@@ -108,14 +108,14 @@ describe('postgresjs auto instrumentation', () => {
             'db.namespace': 'test_db',
             'db.system.name': 'postgres',
             'sentry.op': 'db',
-            'sentry.origin': 'manual',
+            'sentry.origin': 'auto.db.otel.postgres',
             'server.address': 'localhost',
             'server.port': 5444,
           }),
           description: 'SELECT * from generate_series(?,?) as x',
           op: 'db',
           status: 'ok',
-          origin: 'manual',
+          origin: 'auto.db.otel.postgres',
           parent_span_id: expect.any(String),
           span_id: expect.any(String),
           start_timestamp: expect.any(Number),
@@ -127,14 +127,14 @@ describe('postgresjs auto instrumentation', () => {
             'db.namespace': 'test_db',
             'db.system.name': 'postgres',
             'sentry.op': 'db',
-            'sentry.origin': 'manual',
+            'sentry.origin': 'auto.db.otel.postgres',
             'server.address': 'localhost',
             'server.port': 5444,
           }),
           description: 'DROP TABLE "User"',
           op: 'db',
           status: 'ok',
-          origin: 'manual',
+          origin: 'auto.db.otel.postgres',
           parent_span_id: expect.any(String),
           span_id: expect.any(String),
           start_timestamp: expect.any(Number),
@@ -146,14 +146,14 @@ describe('postgresjs auto instrumentation', () => {
             'db.namespace': 'test_db',
             'db.system.name': 'postgres',
             'sentry.op': 'db',
-            'sentry.origin': 'manual',
+            'sentry.origin': 'auto.db.otel.postgres',
             'server.address': 'localhost',
             'server.port': 5444,
           }),
           description: 'SELECT * FROM "User" WHERE "email" = \'foo@baz.com\'',
           op: 'db',
           status: 'unknown_error',
-          origin: 'manual',
+          origin: 'auto.db.otel.postgres',
           parent_span_id: expect.any(String),
           span_id: expect.any(String),
           start_timestamp: expect.any(Number),
diff --git a/packages/node/src/integrations/tracing/postgresjs.ts b/packages/node/src/integrations/tracing/postgresjs.ts
@@ -1,9 +1,11 @@
 // Instrumentation for https://github.com/porsager/postgres
+import { context, trace } from '@opentelemetry/api';
 import type { InstrumentationConfig } from '@opentelemetry/instrumentation';
 import {
   InstrumentationBase,
   InstrumentationNodeModuleDefinition,
   InstrumentationNodeModuleFile,
+  safeExecuteInTheMiddle,
 } from '@opentelemetry/instrumentation';
 import {
   ATTR_DB_NAMESPACE,
@@ -12,20 +14,56 @@ import {
   ATTR_SERVER_PORT,
 } from '@opentelemetry/semantic-conventions';
 import type { IntegrationFn, Span } from '@sentry/core';
-import { defineIntegration, getCurrentScope, SDK_VERSION, SPAN_STATUS_ERROR, startSpanManual } from '@sentry/core';
+import {
+  defineIntegration,
+  getCurrentScope,
+  logger,
+  SDK_VERSION,
+  SPAN_STATUS_ERROR,
+  startSpanManual,
+} from '@sentry/core';
 import { generateInstrumentOnce } from '../../otel/instrument';
+import { addOriginToSpan } from '../../utils/addOriginToSpan';
 
 const INTEGRATION_NAME = 'PostgresJs';
 const SUPPORTED_VERSIONS = ['>=3.0.0 <4'];
 
-export const instrumentPostgresJs = generateInstrumentOnce(INTEGRATION_NAME, () => new PostgresJsInstrumentation());
+type PostgresConnectionContext = {
+  ATTR_DB_NAMESPACE?: string; // Database name
+  ATTR_SERVER_ADDRESS?: string; // Hostname or IP address of the database server
+  ATTR_SERVER_PORT?: string; // Port number of the database server
+};
+
+type PostgresJsInstrumentationConfig = InstrumentationConfig & {
+  /**
+   * Whether to require a parent span for the instrumentation.
+   * If set to true, the instrumentation will only create spans if there is a parent span
+   * available in the current scope.
+   * @default true
+   */
+  requireParentSpan?: boolean;
+  /**
+   * Hook to modify the span before it is started.
+   * This can be used to set additional attributes or modify the span in any way.
+   */
+  requestHook?: (span: Span, sanitizedSqlQuery: string, postgresConnectionContext?: PostgresConnectionContext) => void;
+};
+
+export const instrumentPostgresJs = generateInstrumentOnce(
+  INTEGRATION_NAME,
+  (options?: PostgresJsInstrumentationConfig) =>
+    new PostgresJsInstrumentation({
+      requireParentSpan: options?.requireParentSpan ?? true,
+      requestHook: options?.requestHook,
+    }),
+);
 
 /**
  * Instrumentation for the [postgres](https://www.npmjs.com/package/postgres) library.
  * This instrumentation captures postgresjs queries and their attributes,
  */
-export class PostgresJsInstrumentation extends InstrumentationBase {
-  public constructor(config: InstrumentationConfig = {}) {
+export class PostgresJsInstrumentation extends InstrumentationBase<PostgresJsInstrumentationConfig> {
+  public constructor(config: PostgresJsInstrumentationConfig) {
     super('sentry-postgres-js', SDK_VERSION, config);
   }
 
@@ -58,6 +96,17 @@ export class PostgresJsInstrumentation extends InstrumentationBase {
     return [instrumentationModule];
   }
 
+  /**
+   * Determines whether a span should be created based on the current context.
+   * If `requireParentSpan` is set to true in the configuration, a span will
+   * only be created if there is a parent span available.
+   */
+  private _shouldCreateSpans(): boolean {
+    const config = this.getConfig();
+    const hasParentSpan = trace.getSpan(context.active()) !== undefined;
+    return hasParentSpan || !config.requireParentSpan;
+  }
+
   /**
    * Patches the reject method of the Query class to set the span status and end it
    */
@@ -114,6 +163,11 @@ export class PostgresJsInstrumentation extends InstrumentationBase {
         },
         handleArgs,
       ) => {
+        if (!this._shouldCreateSpans()) {
+          // If we don't need to create spans, just call the original method
+          return Reflect.apply(handleTarget, handleThisArg, handleArgs);
+        }
+
         const sanitizedSqlQuery = this._sanitizeSqlQuery(handleThisArg.strings?.[0]);
 
         return startSpanManual(
@@ -124,13 +178,24 @@ export class PostgresJsInstrumentation extends InstrumentationBase {
           (span: Span) => {
             const scope = getCurrentScope();
             const postgresConnectionContext = scope.getScopeData().contexts['postgresjsConnection'] as
-              | {
-                  ATTR_DB_NAMESPACE: string;
-                  ATTR_SERVER_ADDRESS: string;
-                  ATTR_SERVER_PORT: string;
-                }
+              | PostgresConnectionContext
               | undefined;
 
+            addOriginToSpan(span, 'auto.db.otel.postgres');
+
+            const { requestHook } = this.getConfig();
+
+            if (requestHook) {
+              safeExecuteInTheMiddle(
+                () => requestHook(span, sanitizedSqlQuery, postgresConnectionContext),
+                error => {
+                  if (error) {
+                    logger.error(`Error in requestHook for ${INTEGRATION_NAME} integration:`, error);
+                  }
+                },
+              );
+            }
+
             // ATTR_DB_NAMESPACE is used to indicate the database name and the schema name
             // It's only the database name as we don't have the schema information
             const databaseName = postgresConnectionContext?.ATTR_DB_NAMESPACE || '<unknown database>';
