From 0300e4a7df551f0afec758803767242e6c27c1a7 Mon Sep 17 00:00:00 2001 From: Andrei Borza Date: Mon, 6 Jul 2026 14:39:01 +0200 Subject: [PATCH 1/5] feat(server-utils): Migrate postgres.js instrumentation to orchestrion Co-Authored-By: Claude Opus 4.8 --- .../postgresjs/instrument-requestHook.mjs | 39 +- .../suites/tracing/postgresjs/test.ts | 86 ++-- packages/core/src/server-exports.ts | 6 +- .../tracing-channel/postgres-js.ts | 372 ++++++++++++++++++ .../server-utils/src/orchestrion/channels.ts | 2 + .../src/orchestrion/config/index.ts | 2 + .../src/orchestrion/config/postgres.ts | 43 ++ .../server-utils/src/orchestrion/index.ts | 4 + .../tracing-channel/postgres-js.test.ts | 327 +++++++++++++++ 9 files changed, 824 insertions(+), 57 deletions(-) create mode 100644 packages/server-utils/src/integrations/tracing-channel/postgres-js.ts create mode 100644 packages/server-utils/src/orchestrion/config/postgres.ts create mode 100644 packages/server-utils/test/integrations/tracing-channel/postgres-js.test.ts diff --git a/dev-packages/node-integration-tests/suites/tracing/postgresjs/instrument-requestHook.mjs b/dev-packages/node-integration-tests/suites/tracing/postgresjs/instrument-requestHook.mjs index 885c6198100b..afe717c95548 100644 --- a/dev-packages/node-integration-tests/suites/tracing/postgresjs/instrument-requestHook.mjs +++ b/dev-packages/node-integration-tests/suites/tracing/postgresjs/instrument-requestHook.mjs @@ -1,25 +1,32 @@ import * as Sentry from '@sentry/node'; import { loggingTransport } from '@sentry-internal/node-integration-tests'; +const requestHook = (span, sanitizedSqlQuery, connectionContext) => { + // Add custom attributes to demonstrate requestHook functionality + span.setAttribute('custom.requestHook', 'called'); + + // Set context information as extras for test validation + Sentry.setExtra('requestHookCalled', { + sanitizedQuery: sanitizedSqlQuery, + database: connectionContext?.ATTR_DB_NAMESPACE, + host: connectionContext?.ATTR_SERVER_ADDRESS, + port: connectionContext?.ATTR_SERVER_PORT, + }); +}; + +// Under orchestrion (INJECT_ORCHESTRION), `experimentalUseDiagnosticsChannelInjection()` has already run +// and the default `PostgresJs` OTel integration is swapped for the channel one. Passing the OTel +// `postgresJsIntegration()` explicitly here would override that swap and silently re-test the old path, +// so use the channel integration instead — configured with the same requestHook. +const postgresJsIntegration = + process.env.INJECT_ORCHESTRION === 'true' + ? Sentry.diagnosticsChannelInjectionIntegrations().postgresJsIntegration({ requestHook }) + : Sentry.postgresJsIntegration({ requestHook }); + Sentry.init({ dsn: 'https://public@dsn.ingest.sentry.io/1337', release: '1.0', tracesSampleRate: 1.0, transport: loggingTransport, - integrations: [ - Sentry.postgresJsIntegration({ - requestHook: (span, sanitizedSqlQuery, connectionContext) => { - // Add custom attributes to demonstrate requestHook functionality - span.setAttribute('custom.requestHook', 'called'); - - // Set context information as extras for test validation - Sentry.setExtra('requestHookCalled', { - sanitizedQuery: sanitizedSqlQuery, - database: connectionContext?.ATTR_DB_NAMESPACE, - host: connectionContext?.ATTR_SERVER_ADDRESS, - port: connectionContext?.ATTR_SERVER_PORT, - }); - }, - }), - ], + integrations: [postgresJsIntegration], }); diff --git a/dev-packages/node-integration-tests/suites/tracing/postgresjs/test.ts b/dev-packages/node-integration-tests/suites/tracing/postgresjs/test.ts index dd36478857fe..535135b2ddc4 100644 --- a/dev-packages/node-integration-tests/suites/tracing/postgresjs/test.ts +++ b/dev-packages/node-integration-tests/suites/tracing/postgresjs/test.ts @@ -1,4 +1,5 @@ import { afterAll, describe, expect } from 'vitest'; +import { isOrchestrionEnabled } from '../../../utils'; import { cleanupChildProcesses, createEsmAndCjsTests } from '../../../utils/runner'; describe('postgresjs auto instrumentation', () => { @@ -6,6 +7,11 @@ describe('postgresjs auto instrumentation', () => { cleanupChildProcesses(); }); + // Under orchestrion (INJECT_ORCHESTRION), the OTel `PostgresJs` integration is + // swapped for the diagnostics-channel one, so query spans carry a different + // origin. Every other attribute is identical. + const ORIGIN = isOrchestrionEnabled() ? 'auto.db.orchestrion.postgresjs' : 'auto.db.postgresjs'; + describe('basic', () => { const EXPECTED_TRANSACTION = { transaction: 'Test Transaction', @@ -18,7 +24,7 @@ describe('postgresjs auto instrumentation', () => { 'db.query.text': '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"))', 'sentry.op': 'db', - 'sentry.origin': 'auto.db.postgresjs', + 'sentry.origin': ORIGIN, 'server.address': 'localhost', 'server.port': 5444, }), @@ -26,7 +32,7 @@ describe('postgresjs auto instrumentation', () => { '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: 'auto.db.postgresjs', + origin: ORIGIN, parent_span_id: expect.any(String), span_id: expect.any(String), start_timestamp: expect.any(Number), @@ -39,7 +45,7 @@ describe('postgresjs auto instrumentation', () => { 'db.system.name': 'postgres', 'db.operation.name': 'INSERT', 'db.query.text': 'INSERT INTO "User" ("email", "name") VALUES (?, ?)', - 'sentry.origin': 'auto.db.postgresjs', + 'sentry.origin': ORIGIN, 'sentry.op': 'db', 'server.address': 'localhost', 'server.port': 5444, @@ -47,7 +53,7 @@ describe('postgresjs auto instrumentation', () => { description: 'INSERT INTO "User" ("email", "name") VALUES (?, ?)', op: 'db', status: 'ok', - origin: 'auto.db.postgresjs', + origin: ORIGIN, parent_span_id: expect.any(String), span_id: expect.any(String), start_timestamp: expect.any(Number), @@ -61,14 +67,14 @@ describe('postgresjs auto instrumentation', () => { 'db.operation.name': 'UPDATE', 'db.query.text': 'UPDATE "User" SET "name" = ? WHERE "email" = ?', 'sentry.op': 'db', - 'sentry.origin': 'auto.db.postgresjs', + 'sentry.origin': ORIGIN, 'server.address': 'localhost', 'server.port': 5444, }), description: 'UPDATE "User" SET "name" = ? WHERE "email" = ?', op: 'db', status: 'ok', - origin: 'auto.db.postgresjs', + origin: ORIGIN, parent_span_id: expect.any(String), span_id: expect.any(String), start_timestamp: expect.any(Number), @@ -82,14 +88,14 @@ describe('postgresjs auto instrumentation', () => { 'db.operation.name': 'SELECT', 'db.query.text': 'SELECT * FROM "User" WHERE "email" = ?', 'sentry.op': 'db', - 'sentry.origin': 'auto.db.postgresjs', + 'sentry.origin': ORIGIN, 'server.address': 'localhost', 'server.port': 5444, }), description: 'SELECT * FROM "User" WHERE "email" = ?', op: 'db', status: 'ok', - origin: 'auto.db.postgresjs', + origin: ORIGIN, parent_span_id: expect.any(String), span_id: expect.any(String), start_timestamp: expect.any(Number), @@ -106,14 +112,14 @@ describe('postgresjs auto instrumentation', () => { 'db.operation.name': 'SELECT', 'db.query.text': 'SELECT * FROM "User" WHERE "email" = $1 AND "name" = $2', 'sentry.op': 'db', - 'sentry.origin': 'auto.db.postgresjs', + 'sentry.origin': ORIGIN, 'server.address': 'localhost', 'server.port': 5444, }), description: 'SELECT * FROM "User" WHERE "email" = $1 AND "name" = $2', op: 'db', status: 'ok', - origin: 'auto.db.postgresjs', + origin: ORIGIN, parent_span_id: expect.any(String), span_id: expect.any(String), start_timestamp: expect.any(Number), @@ -127,14 +133,14 @@ describe('postgresjs auto instrumentation', () => { 'db.operation.name': 'SELECT', 'db.query.text': 'SELECT * from generate_series(?,?) as x', 'sentry.op': 'db', - 'sentry.origin': 'auto.db.postgresjs', + 'sentry.origin': ORIGIN, 'server.address': 'localhost', 'server.port': 5444, }), description: 'SELECT * from generate_series(?,?) as x', op: 'db', status: 'ok', - origin: 'auto.db.postgresjs', + origin: ORIGIN, parent_span_id: expect.any(String), span_id: expect.any(String), start_timestamp: expect.any(Number), @@ -148,14 +154,14 @@ describe('postgresjs auto instrumentation', () => { 'db.operation.name': 'DROP TABLE', 'db.query.text': 'DROP TABLE "User"', 'sentry.op': 'db', - 'sentry.origin': 'auto.db.postgresjs', + 'sentry.origin': ORIGIN, 'server.address': 'localhost', 'server.port': 5444, }), description: 'DROP TABLE "User"', op: 'db', status: 'ok', - origin: 'auto.db.postgresjs', + origin: ORIGIN, parent_span_id: expect.any(String), span_id: expect.any(String), start_timestamp: expect.any(Number), @@ -171,14 +177,14 @@ describe('postgresjs auto instrumentation', () => { 'error.type': 'PostgresError', 'db.query.text': 'SELECT * FROM "User" WHERE "email" = ?', 'sentry.op': 'db', - 'sentry.origin': 'auto.db.postgresjs', + 'sentry.origin': ORIGIN, 'server.address': 'localhost', 'server.port': 5444, }), description: 'SELECT * FROM "User" WHERE "email" = ?', op: 'db', status: 'internal_error', - origin: 'auto.db.postgresjs', + origin: ORIGIN, parent_span_id: expect.any(String), span_id: expect.any(String), start_timestamp: expect.any(Number), @@ -252,7 +258,7 @@ describe('postgresjs auto instrumentation', () => { '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"))', 'custom.requestHook': 'called', 'sentry.op': 'db', - 'sentry.origin': 'auto.db.postgresjs', + 'sentry.origin': ORIGIN, 'server.address': 'localhost', 'server.port': 5444, }), @@ -260,7 +266,7 @@ describe('postgresjs auto instrumentation', () => { '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: 'auto.db.postgresjs', + origin: ORIGIN, }), expect.objectContaining({ data: expect.objectContaining({ @@ -270,14 +276,14 @@ describe('postgresjs auto instrumentation', () => { 'db.query.text': 'INSERT INTO "User" ("email", "name") VALUES (?, ?)', 'custom.requestHook': 'called', 'sentry.op': 'db', - 'sentry.origin': 'auto.db.postgresjs', + 'sentry.origin': ORIGIN, 'server.address': 'localhost', 'server.port': 5444, }), description: 'INSERT INTO "User" ("email", "name") VALUES (?, ?)', op: 'db', status: 'ok', - origin: 'auto.db.postgresjs', + origin: ORIGIN, }), expect.objectContaining({ data: expect.objectContaining({ @@ -287,14 +293,14 @@ describe('postgresjs auto instrumentation', () => { 'db.query.text': 'SELECT * FROM "User" WHERE "email" = ?', 'custom.requestHook': 'called', 'sentry.op': 'db', - 'sentry.origin': 'auto.db.postgresjs', + 'sentry.origin': ORIGIN, 'server.address': 'localhost', 'server.port': 5444, }), description: 'SELECT * FROM "User" WHERE "email" = ?', op: 'db', status: 'ok', - origin: 'auto.db.postgresjs', + origin: ORIGIN, }), expect.objectContaining({ data: expect.objectContaining({ @@ -304,14 +310,14 @@ describe('postgresjs auto instrumentation', () => { 'db.query.text': 'DROP TABLE "User"', 'custom.requestHook': 'called', 'sentry.op': 'db', - 'sentry.origin': 'auto.db.postgresjs', + 'sentry.origin': ORIGIN, 'server.address': 'localhost', 'server.port': 5444, }), description: 'DROP TABLE "User"', op: 'db', status: 'ok', - origin: 'auto.db.postgresjs', + origin: ORIGIN, }), ]), extra: expect.objectContaining({ @@ -353,7 +359,7 @@ describe('postgresjs auto instrumentation', () => { 'db.query.text': '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"))', 'sentry.op': 'db', - 'sentry.origin': 'auto.db.postgresjs', + 'sentry.origin': ORIGIN, 'server.address': 'localhost', 'server.port': 5444, }), @@ -361,7 +367,7 @@ describe('postgresjs auto instrumentation', () => { '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: 'auto.db.postgresjs', + origin: ORIGIN, }), expect.objectContaining({ data: expect.objectContaining({ @@ -370,14 +376,14 @@ describe('postgresjs auto instrumentation', () => { 'db.operation.name': 'INSERT', 'db.query.text': 'INSERT INTO "User" ("email", "name") VALUES (?, ?)', 'sentry.op': 'db', - 'sentry.origin': 'auto.db.postgresjs', + 'sentry.origin': ORIGIN, 'server.address': 'localhost', 'server.port': 5444, }), description: 'INSERT INTO "User" ("email", "name") VALUES (?, ?)', op: 'db', status: 'ok', - origin: 'auto.db.postgresjs', + origin: ORIGIN, }), expect.objectContaining({ data: expect.objectContaining({ @@ -386,14 +392,14 @@ describe('postgresjs auto instrumentation', () => { 'db.operation.name': 'SELECT', 'db.query.text': 'SELECT * FROM "User" WHERE "email" = ?', 'sentry.op': 'db', - 'sentry.origin': 'auto.db.postgresjs', + 'sentry.origin': ORIGIN, 'server.address': 'localhost', 'server.port': 5444, }), description: 'SELECT * FROM "User" WHERE "email" = ?', op: 'db', status: 'ok', - origin: 'auto.db.postgresjs', + origin: ORIGIN, }), expect.objectContaining({ data: expect.objectContaining({ @@ -402,14 +408,14 @@ describe('postgresjs auto instrumentation', () => { 'db.operation.name': 'DELETE', 'db.query.text': 'DELETE FROM "User" WHERE "email" = ?', 'sentry.op': 'db', - 'sentry.origin': 'auto.db.postgresjs', + 'sentry.origin': ORIGIN, 'server.address': 'localhost', 'server.port': 5444, }), description: 'DELETE FROM "User" WHERE "email" = ?', op: 'db', status: 'ok', - origin: 'auto.db.postgresjs', + origin: ORIGIN, }), ]), }; @@ -442,14 +448,14 @@ describe('postgresjs auto instrumentation', () => { 'db.operation.name': 'CREATE TABLE', 'db.query.text': 'CREATE TABLE "User" ("id" SERIAL NOT NULL, "email" TEXT NOT NULL, PRIMARY KEY ("id"))', 'sentry.op': 'db', - 'sentry.origin': 'auto.db.postgresjs', + 'sentry.origin': ORIGIN, 'server.address': 'localhost', 'server.port': 5444, }), description: 'CREATE TABLE "User" ("id" SERIAL NOT NULL, "email" TEXT NOT NULL, PRIMARY KEY ("id"))', op: 'db', status: 'ok', - origin: 'auto.db.postgresjs', + origin: ORIGIN, }), // sql.unsafe() with $1 placeholders - preserved per OTEL spec expect.objectContaining({ @@ -459,14 +465,14 @@ describe('postgresjs auto instrumentation', () => { 'db.operation.name': 'INSERT', 'db.query.text': 'INSERT INTO "User" ("email") VALUES ($1)', 'sentry.op': 'db', - 'sentry.origin': 'auto.db.postgresjs', + 'sentry.origin': ORIGIN, 'server.address': 'localhost', 'server.port': 5444, }), description: 'INSERT INTO "User" ("email") VALUES ($1)', op: 'db', status: 'ok', - origin: 'auto.db.postgresjs', + origin: ORIGIN, }), expect.objectContaining({ data: expect.objectContaining({ @@ -475,14 +481,14 @@ describe('postgresjs auto instrumentation', () => { 'db.operation.name': 'SELECT', 'db.query.text': 'SELECT * FROM "User" WHERE "email" = $1', 'sentry.op': 'db', - 'sentry.origin': 'auto.db.postgresjs', + 'sentry.origin': ORIGIN, 'server.address': 'localhost', 'server.port': 5444, }), description: 'SELECT * FROM "User" WHERE "email" = $1', op: 'db', status: 'ok', - origin: 'auto.db.postgresjs', + origin: ORIGIN, }), expect.objectContaining({ data: expect.objectContaining({ @@ -491,14 +497,14 @@ describe('postgresjs auto instrumentation', () => { 'db.operation.name': 'DROP TABLE', 'db.query.text': 'DROP TABLE "User"', 'sentry.op': 'db', - 'sentry.origin': 'auto.db.postgresjs', + 'sentry.origin': ORIGIN, 'server.address': 'localhost', 'server.port': 5444, }), description: 'DROP TABLE "User"', op: 'db', status: 'ok', - origin: 'auto.db.postgresjs', + origin: ORIGIN, }), ]), }; diff --git a/packages/core/src/server-exports.ts b/packages/core/src/server-exports.ts index f97e5e677c5e..66f08352917c 100644 --- a/packages/core/src/server-exports.ts +++ b/packages/core/src/server-exports.ts @@ -23,7 +23,11 @@ export type { ExpressMiddleware, ExpressErrorMiddleware, } from './integrations/express/types'; -export { instrumentPostgresJsSql, _sanitizeSqlQuery as _INTERNAL_sanitizeSqlQuery } from './integrations/postgresjs'; +export { + instrumentPostgresJsSql, + _sanitizeSqlQuery as _INTERNAL_sanitizeSqlQuery, + _reconstructQuery as _INTERNAL_reconstructPostgresQuery, +} from './integrations/postgresjs'; export { getSqlQuerySummary as _INTERNAL_getSqlQuerySummary } from './utils/sql'; export { patchHttpModuleClient } from './integrations/http/client-patch'; diff --git a/packages/server-utils/src/integrations/tracing-channel/postgres-js.ts b/packages/server-utils/src/integrations/tracing-channel/postgres-js.ts new file mode 100644 index 000000000000..e01c942cc890 --- /dev/null +++ b/packages/server-utils/src/integrations/tracing-channel/postgres-js.ts @@ -0,0 +1,372 @@ +import * as diagnosticsChannel from 'node:diagnostics_channel'; +import { + DB_NAMESPACE, + DB_OPERATION_NAME, + DB_QUERY_TEXT, + DB_SYSTEM_NAME, + ERROR_TYPE, + SERVER_ADDRESS, + SERVER_PORT, +} from '@sentry/conventions/attributes'; +import type { IntegrationFn, Span } from '@sentry/core'; +import { + _INTERNAL_reconstructPostgresQuery, + _INTERNAL_sanitizeSqlQuery, + debug, + defineIntegration, + getActiveSpan, + SEMANTIC_ATTRIBUTE_SENTRY_ORIGIN, + SPAN_KIND, + SPAN_STATUS_ERROR, + startInactiveSpan, + waitForTracingChannelBinding, +} from '@sentry/core'; +import { DEBUG_BUILD } from '../../debug-build'; +import { CHANNELS } from '../../orchestrion/channels'; +import { bindTracingChannelToSpan } from '../../tracing-channel'; + +// Same name as the OTel `PostgresJs` integration by design: when this is +// enabled, the OTel integration of the same name is dropped from the default +// set (see `experimentalUseDiagnosticsChannelInjection`). +const INTEGRATION_NAME = 'PostgresJs' as const; + +const ORIGIN = 'auto.db.orchestrion.postgresjs'; + +// Not part of `@sentry/conventions`, so we keep it inline (matches the OTel +// `PostgresJsInstrumentation`). +const DB_RESPONSE_STATUS_CODE = 'db.response.status_code'; + +const SQL_OPERATION_REGEX = /^(SELECT|INSERT|UPDATE|DELETE|CREATE|DROP|ALTER)/i; + +// Same `Symbol.for()` marker the core `instrumentPostgresJsSql` wrapper sets on +// queries it manually instruments, so we skip them here and never double-span. +const QUERY_FROM_INSTRUMENTED_SQL = Symbol.for('sentry.query.from.instrumented.sql'); +// The query span, stashed on the `Query` so the `execute` channel can attach +// connection attributes to it. +const QUERY_SPAN = Symbol('sentryPostgresJsSpan'); +// Set once connection attributes are on the span, so the fallback and the +// `execute` channel don't both write them. +const CONNECTION_ATTRS_SET = Symbol('sentryPostgresJsConnectionAttrsSet'); +// Set on the channel context once the resolve/reject wrappers have ended the +// span, so `deferSpanEnd` knows the wrappers own the lifecycle. +const SPAN_ENDED = Symbol('sentryPostgresJsSpanEnded'); + +/** + * Connection attributes resolved from postgres.js' parsed options. Port is kept + * as a string to match the `requestHook` contract of the OTel integration + * (and the portable `instrumentPostgresJsSql`); it's coerced to a number when + * set on the span (semantic conventions expect a number for `server.port`). + */ +interface PostgresConnectionContext { + ATTR_DB_NAMESPACE?: string; + ATTR_SERVER_ADDRESS?: string; + ATTR_SERVER_PORT?: string; +} + +export interface PostgresJsChannelIntegrationOptions { + /** + * Only create spans when there's already an active parent span. Defaults to + * `true`, matching the OTel `postgresJsIntegration`. + */ + requireParentSpan?: boolean; + /** + * Hook to modify the query span before the query runs. Receives the span, the + * sanitized SQL, and (when resolvable) the connection context. + */ + requestHook?: (span: Span, sanitizedSqlQuery: string, postgresConnectionContext?: PostgresConnectionContext) => void; +} + +/** The `Query` instance postgres.js passes as `self` to `Query.prototype.handle`. */ +interface PostgresQuery { + strings?: string[]; + executed?: boolean; + resolve?: (...args: unknown[]) => unknown; + reject?: (...args: unknown[]) => unknown; +} + +interface PostgresJsQueryContext { + arguments?: unknown[]; + self?: PostgresQuery; + result?: unknown; + error?: unknown; +} + +// postgres.js parses `host`/`port` into arrays (it can connect to multiple +// hosts); the `Connection` factory receives this parsed options object. +interface PostgresParsedOptions { + host?: string[]; + port?: number[]; + database?: string; +} + +const NOOP = (): void => {}; + +// A connection object -> its resolved context, populated on the `connection` +// channel and read on the `execute` channel (keyed by the same object). +const connectionContexts = new WeakMap(); +// Distinct endpoints seen so far (value-compared, so N connections to one DB +// count once). When exactly one endpoint exists — the common case, and the only +// one the tests exercise — every query resolves to it at handle-start. +const endpointRegistry: PostgresConnectionContext[] = []; + +function registerEndpoint(context: PostgresConnectionContext): void { + const alreadyKnown = endpointRegistry.some( + e => + e.ATTR_SERVER_ADDRESS === context.ATTR_SERVER_ADDRESS && + e.ATTR_SERVER_PORT === context.ATTR_SERVER_PORT && + e.ATTR_DB_NAMESPACE === context.ATTR_DB_NAMESPACE, + ); + if (!alreadyKnown) { + endpointRegistry.push(context); + } +} + +/** The single known endpoint, or `undefined` when zero or multiple are known. */ +function resolveSingleEndpoint(): PostgresConnectionContext | undefined { + return endpointRegistry.length === 1 ? endpointRegistry[0] : undefined; +} + +function buildConnectionContext(options: PostgresParsedOptions): PostgresConnectionContext { + // postgres.js defaults to 'localhost'/5432, but be defensive. + const host = options.host?.[0] || 'localhost'; + const port = options.port?.[0] || 5432; + return { + ATTR_DB_NAMESPACE: typeof options.database === 'string' && options.database !== '' ? options.database : undefined, + ATTR_SERVER_ADDRESS: host, + ATTR_SERVER_PORT: String(port), + }; +} + +function setConnectionAttributes(span: Span, query: PostgresQuery, context: PostgresConnectionContext): void { + const queryRecord = query as Record; + if (queryRecord[CONNECTION_ATTRS_SET]) { + return; + } + queryRecord[CONNECTION_ATTRS_SET] = true; + + if (context.ATTR_SERVER_ADDRESS) { + span.setAttribute(SERVER_ADDRESS, context.ATTR_SERVER_ADDRESS); + } + if (context.ATTR_SERVER_PORT !== undefined) { + const port = parseInt(context.ATTR_SERVER_PORT, 10); + if (!Number.isNaN(port)) { + span.setAttribute(SERVER_PORT, port); + } + } + if (context.ATTR_DB_NAMESPACE) { + span.setAttribute(DB_NAMESPACE, context.ATTR_DB_NAMESPACE); + } +} + +function setOperationName(span: Span, sanitizedQuery: string | undefined, command?: string): void { + if (command) { + span.setAttribute(DB_OPERATION_NAME, command); + return; + } + const operationMatch = sanitizedQuery?.match(SQL_OPERATION_REGEX); + if (operationMatch?.[1]) { + span.setAttribute(DB_OPERATION_NAME, operationMatch[1].toUpperCase()); + } +} + +/** + * Wrap `query.resolve`/`query.reject` so the span ends when the query settles. + * + * `Query extends Promise` and `async handle()` only dispatches — its promise + * resolves immediately, long before the query completes. postgres.js signals + * completion by calling `this.resolve`/`this.reject`, so we own the span end + * there. Wrapping happens at handle-start because `reject` can fire + * synchronously during dispatch and `cursor()` reassigns both before executing. + */ +function wrapQuerySettlement(data: PostgresJsQueryContext, span: Span, sanitizedSqlQuery: string): void { + const query = data.self; + if (!query) { + return; + } + + // Claim ownership of ending the span up front, so `deferSpanEnd` defers to the + // wrapper even if `span.end()` below throws. + const markEnded = (): void => { + (data as Record)[SPAN_ENDED] = true; + }; + + const originalResolve = query.resolve; + if (typeof originalResolve === 'function') { + query.resolve = function (this: unknown, ...resolveArgs: unknown[]): unknown { + markEnded(); + try { + const command = (resolveArgs[0] as { command?: string } | undefined)?.command; + setOperationName(span, sanitizedSqlQuery, command); + span.end(); + } catch (e) { + DEBUG_BUILD && debug.error('[orchestrion:postgresjs] error ending span in resolve:', e); + } + return originalResolve.apply(this, resolveArgs); + }; + } + + const originalReject = query.reject; + if (typeof originalReject === 'function') { + query.reject = function (this: unknown, ...rejectArgs: unknown[]): unknown { + markEnded(); + try { + const err = rejectArgs[0] as { message?: string; code?: string; name?: string } | undefined; + span.setStatus({ code: SPAN_STATUS_ERROR, message: err?.message || 'unknown_error' }); + span.setAttribute(DB_RESPONSE_STATUS_CODE, err?.code || 'unknown'); + span.setAttribute(ERROR_TYPE, err?.name || 'unknown'); + setOperationName(span, sanitizedSqlQuery); + span.end(); + } catch (e) { + DEBUG_BUILD && debug.error('[orchestrion:postgresjs] error ending span in reject:', e); + } + return originalReject.apply(this, rejectArgs); + }; + } +} + +const _postgresJsChannelIntegration = ((options: PostgresJsChannelIntegrationOptions = {}) => { + const { requireParentSpan, requestHook } = options; + + return { + name: INTEGRATION_NAME, + setupOnce() { + // `tracingChannel` is unavailable before Node 18.19. + if (!diagnosticsChannel.tracingChannel) { + return; + } + + DEBUG_BUILD && debug.log(`[orchestrion:postgresjs] subscribing to "${CHANNELS.POSTGRESJS_HANDLE}"`); + + // Connection + execute are pure observers (no span, no async binding), so + // subscribe immediately — factory-time `Connection()` calls happen before + // `waitForTracingChannelBinding` resolves and must still be recorded. + diagnosticsChannel.tracingChannel(CHANNELS.POSTGRESJS_CONNECTION).subscribe({ + start: NOOP, + asyncStart: NOOP, + asyncEnd: NOOP, + error: NOOP, + end(message) { + const connection = message.result; + const connectionOptions = message.arguments?.[0] as PostgresParsedOptions | undefined; + if (!connection || typeof connection !== 'object' || !connectionOptions) { + return; + } + const context = buildConnectionContext(connectionOptions); + connectionContexts.set(connection, context); + registerEndpoint(context); + }, + }); + + diagnosticsChannel.tracingChannel(CHANNELS.POSTGRESJS_EXECUTE).subscribe({ + end: NOOP, + asyncStart: NOOP, + asyncEnd: NOOP, + error: NOOP, + start(message) { + // `execute` is also called bare (`self === undefined`) for the first + // query on each connection, `fetchState` and `retry`; those miss here + // and rely on the handle-start fallback instead. + const connection = message.self as object | undefined; + const query = message.arguments?.[0] as PostgresQuery | undefined; + if (!connection || !query) { + return; + } + const span = (query as Record)[QUERY_SPAN] as Span | undefined; + const context = connectionContexts.get(connection); + if (span && context) { + setConnectionAttributes(span, query, context); + } + }, + }); + + // The span-creating `handle` subscription needs the async-context binding + // that `initOpenTelemetry()` registers after integration setup. + waitForTracingChannelBinding(() => { + bindTracingChannelToSpan( + diagnosticsChannel.tracingChannel(CHANNELS.POSTGRESJS_HANDLE), + data => { + const query = data.self; + if (!query) { + return undefined; + } + + // Opt out of: re-entrant `handle()` calls (then/catch/finally re-invoke + // it, guarded by `executed`), queries already wrapped by the portable + // `instrumentPostgresJsSql`, and (by default) queries with no parent span. + if (query.executed === true || (query as Record)[QUERY_FROM_INSTRUMENTED_SQL]) { + return undefined; + } + if (requireParentSpan !== false && !getActiveSpan()) { + return undefined; + } + + const fullQuery = _INTERNAL_reconstructPostgresQuery(query.strings); + const sanitizedSqlQuery = _INTERNAL_sanitizeSqlQuery(fullQuery); + + // `kind: CLIENT` matches the mysql/pg channel subscribers. + const span = startInactiveSpan({ + name: sanitizedSqlQuery || 'postgresjs.query', + op: 'db', + kind: SPAN_KIND.CLIENT, + attributes: { + [SEMANTIC_ATTRIBUTE_SENTRY_ORIGIN]: ORIGIN, + [DB_SYSTEM_NAME]: 'postgres', + [DB_QUERY_TEXT]: sanitizedSqlQuery, + }, + }); + + // Stash for the `execute` channel to attach per-connection attributes. + (query as Record)[QUERY_SPAN] = span; + + // Single-endpoint fallback: resolve context now so `requestHook` has it + // and the first-query-per-connection (bare `execute`) path still gets attrs. + const context = resolveSingleEndpoint(); + if (context) { + setConnectionAttributes(span, query, context); + } + + if (requestHook) { + try { + requestHook(span, sanitizedSqlQuery, context); + } catch (e) { + span.setAttribute('sentry.hook.error', 'requestHook failed'); + DEBUG_BUILD && debug.error('[orchestrion:postgresjs] error in requestHook:', e); + } + } + + wrapQuerySettlement(data, span, sanitizedSqlQuery); + + return span; + }, + { + deferSpanEnd({ data }) { + // `handle` is async: its promise settles on dispatch (asyncEnd), long + // before the query does. The resolve/reject wrappers own the ending. + if ((data as Record)[SPAN_ENDED]) { + return true; // wrappers already ended it + } + if ('error' in data) { + return false; // `handle()` itself threw; the error subscriber annotated the span, let the helper end it + } + // NOTE: for a cursor consumed as an async iterator, only the first batch + // reaches `handle` (the `executed` guard blocks the rest), so the span + // ends on the first batch — a pre-existing flaw kept for parity. + return true; // query in flight; the wrappers will end the span when it settles + }, + }, + ); + }); + }, + }; +}) satisfies IntegrationFn; + +/** + * EXPERIMENTAL — orchestrion-driven postgres.js (`postgres` v3.x) integration. + * + * Subscribes to the `orchestrion:postgres:handle` / `:connection` / `:execute` + * diagnostics channels injected into postgres.js' `Query.prototype.handle` and + * `Connection`/`execute` (in `src/*` and `cjs/src/*`) and creates db spans + * matching the OTel `postgresJsIntegration`. Requires the orchestrion runtime + * hook or bundler plugin. + */ +export const postgresJsChannelIntegration = defineIntegration(_postgresJsChannelIntegration); diff --git a/packages/server-utils/src/orchestrion/channels.ts b/packages/server-utils/src/orchestrion/channels.ts index b008aad618a2..4117be46bbc4 100644 --- a/packages/server-utils/src/orchestrion/channels.ts +++ b/packages/server-utils/src/orchestrion/channels.ts @@ -2,6 +2,7 @@ import { mysqlChannels } from './config/mysql'; import { lruMemoizerChannels } from './config/lru-memoizer'; import { ioredisChannels } from './config/ioredis'; import { pgChannels } from './config/pg'; +import { postgresJsChannels } from './config/postgres'; import { openaiChannels } from './config/openai'; import { anthropicAiChannels } from './config/anthropic-ai'; import { vercelAiChannels } from './config/vercel-ai'; @@ -24,6 +25,7 @@ export const CHANNELS = { ...lruMemoizerChannels, ...ioredisChannels, ...pgChannels, + ...postgresJsChannels, ...openaiChannels, ...anthropicAiChannels, ...vercelAiChannels, diff --git a/packages/server-utils/src/orchestrion/config/index.ts b/packages/server-utils/src/orchestrion/config/index.ts index 218cbf7e8875..819a8360a72c 100644 --- a/packages/server-utils/src/orchestrion/config/index.ts +++ b/packages/server-utils/src/orchestrion/config/index.ts @@ -4,6 +4,7 @@ import { lruMemoizerConfig } from './lru-memoizer'; import { ioredisConfig } from './ioredis'; import { openaiConfig } from './openai'; import { pgConfig } from './pg'; +import { postgresJsConfig } from './postgres'; import { anthropicAiConfig } from './anthropic-ai'; import { vercelAiConfig } from './vercel-ai'; @@ -13,6 +14,7 @@ export const SENTRY_INSTRUMENTATIONS: InstrumentationConfig[] = [ ...ioredisConfig, ...openaiConfig, ...pgConfig, + ...postgresJsConfig, ...anthropicAiConfig, ...vercelAiConfig, ]; diff --git a/packages/server-utils/src/orchestrion/config/postgres.ts b/packages/server-utils/src/orchestrion/config/postgres.ts new file mode 100644 index 000000000000..6f0cb5639035 --- /dev/null +++ b/packages/server-utils/src/orchestrion/config/postgres.ts @@ -0,0 +1,43 @@ +import type { InstrumentationConfig } from '@apm-js-collab/code-transformer'; + +// postgres.js (`postgres` npm package, v3.x). Named after the npm package; +// `postgres` doesn't collide with `pg.ts` (that file instruments `pg`/`pg-pool`). +// +// The ESM build lives under `src/*`, the CJS build under `cjs/src/*` (the +// `cf/*` workerd build has no channel subscribers, see the integration). +// Both builds share the same class/function shapes, so a single `flatMap` +// over the two dirs emits one entry per (dir, target). +const postgresJsInstrumentationConfig = (dir: string): InstrumentationConfig[] => [ + // `Query.prototype.handle` (`class Query extends Promise`) is the single + // funnel every query passes through (`then`/`catch`/`finally`/`.execute()`/ + // `.forEach()`/cursor all call it), guarded by `this.executed`. `Async` + // because `handle` is `async`. + { + channelName: 'handle', + module: { name: 'postgres', versionRange: '>=3.0.0 <4', filePath: `${dir}/query.js` }, + functionQuery: { className: 'Query', methodName: 'handle', kind: 'Async' }, + }, + // `function Connection(options, ...)` (default export of `connection.js`) + // returns the connection object; used to build the endpoint registry that + // resolves `server.address`/`server.port`/`db.namespace`. + { + channelName: 'connection', + module: { name: 'postgres', versionRange: '>=3.0.0 <4', filePath: `${dir}/connection.js` }, + functionQuery: { functionName: 'Connection', kind: 'Sync' }, + }, + // The nested `function execute(q)` inside `Connection`; the per-connection + // hook that attaches connection attributes to the query's span. + { + channelName: 'execute', + module: { name: 'postgres', versionRange: '>=3.0.0 <4', filePath: `${dir}/connection.js` }, + functionQuery: { functionName: 'execute', kind: 'Sync' }, + }, +]; + +export const postgresJsConfig = ['src', 'cjs/src'].flatMap(postgresJsInstrumentationConfig); + +export const postgresJsChannels = { + POSTGRESJS_HANDLE: 'orchestrion:postgres:handle', + POSTGRESJS_CONNECTION: 'orchestrion:postgres:connection', + POSTGRESJS_EXECUTE: 'orchestrion:postgres:execute', +} as const; diff --git a/packages/server-utils/src/orchestrion/index.ts b/packages/server-utils/src/orchestrion/index.ts index 3682ca752232..499876223123 100644 --- a/packages/server-utils/src/orchestrion/index.ts +++ b/packages/server-utils/src/orchestrion/index.ts @@ -4,6 +4,7 @@ import { lruMemoizerChannelIntegration } from '../integrations/tracing-channel/l import { mysqlChannelIntegration } from '../integrations/tracing-channel/mysql'; import { openaiChannelIntegration } from '../integrations/tracing-channel/openai'; import { postgresChannelIntegration } from '../integrations/tracing-channel/postgres'; +import { postgresJsChannelIntegration } from '../integrations/tracing-channel/postgres-js'; import { vercelAiChannelIntegration } from '../integrations/tracing-channel/vercel-ai'; export { detectOrchestrionSetup, isOrchestrionInjected } from './detect'; @@ -14,9 +15,11 @@ export { mysqlChannelIntegration, openaiChannelIntegration, postgresChannelIntegration, + postgresJsChannelIntegration, vercelAiChannelIntegration, }; export type { IORedisChannelIntegrationOptions, IORedisResponseHook } from '../integrations/tracing-channel/ioredis'; +export type { PostgresJsChannelIntegrationOptions } from '../integrations/tracing-channel/postgres-js'; /** * The canonical set of orchestrion diagnostics-channel integrations, keyed by their public @@ -33,6 +36,7 @@ export type { IORedisChannelIntegrationOptions, IORedisResponseHook } from '../i */ export const channelIntegrations = { postgresIntegration: postgresChannelIntegration, + postgresJsIntegration: postgresJsChannelIntegration, mysqlIntegration: mysqlChannelIntegration, lruMemoizerIntegration: lruMemoizerChannelIntegration, openaiIntegration: openaiChannelIntegration, diff --git a/packages/server-utils/test/integrations/tracing-channel/postgres-js.test.ts b/packages/server-utils/test/integrations/tracing-channel/postgres-js.test.ts new file mode 100644 index 000000000000..a427dd31cc20 --- /dev/null +++ b/packages/server-utils/test/integrations/tracing-channel/postgres-js.test.ts @@ -0,0 +1,327 @@ +import { AsyncLocalStorage } from 'node:async_hooks'; +import { tracingChannel } from 'node:diagnostics_channel'; +import type { Scope, Span } from '@sentry/core'; +import { + _INTERNAL_setSpanForScope, + Client, + createTransport, + getClient, + getCurrentScope, + getDefaultCurrentScope, + getDefaultIsolationScope, + getGlobalScope, + initAndBind, + resolvedSyncPromise, + setAsyncContextStrategy, + spanToJSON, + startSpan, +} from '@sentry/core'; +import { afterAll, beforeAll, beforeEach, describe, expect, it, vi } from 'vitest'; +import { postgresJsChannelIntegration } from '../../../src/integrations/tracing-channel/postgres-js'; +import { CHANNELS } from '../../../src/orchestrion/channels'; + +interface TestStore { + scope: Scope; + isolationScope: Scope; +} + +class TestClient extends Client { + public eventFromException(): PromiseLike { + return resolvedSyncPromise({}); + } + public eventFromMessage(): PromiseLike { + return resolvedSyncPromise({}); + } +} + +function initTestClient(): void { + initAndBind(TestClient, { + dsn: 'https://username@domain/123', + integrations: [], + sendClientReports: false, + stackParser: () => [], + tracesSampleRate: 1, + transport: () => createTransport({ recordDroppedEvent: () => undefined }, () => resolvedSyncPromise({})), + }); +} + +function installTestAsyncContextStrategy(): void { + const asyncStorage = new AsyncLocalStorage(); + + function getScopes(): TestStore { + return ( + asyncStorage.getStore() || { + scope: getDefaultCurrentScope(), + isolationScope: getDefaultIsolationScope(), + } + ); + } + + setAsyncContextStrategy({ + withScope: callback => { + const scope = getScopes().scope.clone(); + const isolationScope = getScopes().isolationScope; + return asyncStorage.run({ scope, isolationScope }, () => callback(scope)); + }, + withSetScope: (scope, callback) => { + const isolationScope = getScopes().isolationScope; + return asyncStorage.run({ scope, isolationScope }, () => callback(scope)); + }, + withIsolationScope: callback => { + const scope = getScopes().scope; + const isolationScope = getScopes().isolationScope.clone(); + return asyncStorage.run({ scope, isolationScope }, () => callback(isolationScope)); + }, + withSetIsolationScope: (isolationScope, callback) => { + const scope = getScopes().scope; + return asyncStorage.run({ scope, isolationScope }, () => callback(isolationScope)); + }, + getCurrentScope: () => getScopes().scope, + getIsolationScope: () => getScopes().isolationScope, + getTracingChannelBinding: () => ({ + asyncLocalStorage: asyncStorage, + getStoreWithActiveSpan: span => { + const scope = getScopes().scope.clone(); + const isolationScope = getScopes().isolationScope; + _INTERNAL_setSpanForScope(scope, span); + return { scope, isolationScope }; + }, + }), + }); +} + +/** A minimal stand-in for a postgres.js `Query`: what `Query.prototype.handle` sees as `self`. */ +interface FakeQuery { + strings?: string[]; + executed?: boolean; + resolve: (...args: unknown[]) => unknown; + reject: (...args: unknown[]) => unknown; + [key: symbol]: unknown; +} + +function makeQuery(strings: string[] | undefined, extra: Partial = {}): FakeQuery { + return { + strings, + executed: false, + resolve: vi.fn(), + reject: vi.fn(), + ...extra, + }; +} + +const endedSpans: Span[] = []; +let requestHookSpy: ReturnType | undefined; + +/** + * Drive `Query.prototype.handle` through its tracing channel, mirroring how + * postgres.js dispatches: `handle()` (async) returns immediately, so the span is + * still open when the promise settles. The caller ends it by invoking + * `query.resolve`/`query.reject` afterwards (as postgres.js does on completion). + */ +async function driveHandle(query: FakeQuery, { withParent = true }: { withParent?: boolean } = {}): Promise { + const drive = async (): Promise => { + await tracingChannel(CHANNELS.POSTGRESJS_HANDLE).tracePromise(async () => {}, { arguments: [], self: query }); + }; + + if (withParent) { + await startSpan({ name: 'parent' }, drive); + } else { + await drive(); + } +} + +function publishConnection(connection: object, options: Record): void { + tracingChannel(CHANNELS.POSTGRESJS_CONNECTION).traceSync(() => connection, { arguments: [options] }); +} + +function publishExecute(connection: object, query: FakeQuery): void { + tracingChannel(CHANNELS.POSTGRESJS_EXECUTE).traceSync(() => undefined, { self: connection, arguments: [query] }); +} + +function lastPgSpan(): Span | undefined { + return endedSpans.filter(s => spanToJSON(s).data['sentry.origin'] === 'auto.db.orchestrion.postgresjs').at(-1); +} + +describe('postgresJsChannelIntegration', () => { + beforeAll(() => { + installTestAsyncContextStrategy(); + initTestClient(); + const integration = postgresJsChannelIntegration({ requestHook: (...args) => requestHookSpy?.(...args) }); + integration.setupOnce?.(); + getClient()?.on('spanEnd', span => endedSpans.push(span)); + }); + + afterAll(() => { + setAsyncContextStrategy(undefined); + getCurrentScope().clear(); + getCurrentScope().setClient(undefined); + getGlobalScope().clear(); + vi.clearAllMocks(); + }); + + beforeEach(() => { + endedSpans.length = 0; + requestHookSpy = vi.fn(); + }); + + // These run with an empty endpoint registry (no `connection` channel driven + // yet), so query spans carry no connection attributes here — those are + // covered in the "connection context" block below. + describe('query span (handle channel)', () => { + it('creates a db span with the NEW semconv shape and orchestrion origin', async () => { + const query = makeQuery(['SELECT * FROM "User"']); + await driveHandle(query); + query.resolve({ command: 'SELECT' }); + + const span = lastPgSpan(); + expect(span).toBeDefined(); + const json = spanToJSON(span!); + expect(json.description).toBe('SELECT * FROM "User"'); + expect(json.op).toBe('db'); + // A successful core span leaves the status unset (the OTel pipeline maps it to 'ok'). + expect(json.status).toBeUndefined(); + expect(json.data['sentry.origin']).toBe('auto.db.orchestrion.postgresjs'); + expect(json.data['db.system.name']).toBe('postgres'); + expect(json.data['db.query.text']).toBe('SELECT * FROM "User"'); + expect(json.data['db.operation.name']).toBe('SELECT'); + }); + + it('reconstructs $n placeholders from tagged-template strings', async () => { + const query = makeQuery(['SELECT * FROM users WHERE id = ', ' AND name = ', '']); + await driveHandle(query); + query.resolve({ command: 'SELECT' }); + + const json = spanToJSON(lastPgSpan()!); + expect(json.description).toBe('SELECT * FROM users WHERE id = $1 AND name = $2'); + expect(json.data['db.query.text']).toBe('SELECT * FROM users WHERE id = $1 AND name = $2'); + }); + + it('sets error status and error attributes when the query rejects', async () => { + const query = makeQuery(['SELECT * FROM "Missing"']); + await driveHandle(query); + query.reject({ message: 'relation "Missing" does not exist', code: '42P01', name: 'PostgresError' }); + + const json = spanToJSON(lastPgSpan()!); + expect(json.status).toBe('relation "Missing" does not exist'); + expect(json.data['db.response.status_code']).toBe('42P01'); + expect(json.data['error.type']).toBe('PostgresError'); + // Falls back to the leading SQL keyword when the rejection carries no command. + expect(json.data['db.operation.name']).toBe('SELECT'); + }); + + it('does not create a span for re-entrant handle() calls (then/catch/finally)', async () => { + const query = makeQuery(['SELECT 1'], { executed: true }); + await driveHandle(query); + query.resolve({ command: 'SELECT' }); + + expect(lastPgSpan()).toBeUndefined(); + }); + + it('does not create a span for queries already wrapped by instrumentPostgresJsSql', async () => { + const query = makeQuery(['SELECT 1'], { + [Symbol.for('sentry.query.from.instrumented.sql')]: true, + }); + await driveHandle(query); + query.resolve({ command: 'SELECT' }); + + expect(lastPgSpan()).toBeUndefined(); + }); + + it('does not create a span without an active parent span by default', async () => { + const query = makeQuery(['SELECT 1']); + await driveHandle(query, { withParent: false }); + query.resolve({ command: 'SELECT' }); + + expect(lastPgSpan()).toBeUndefined(); + }); + + it('runs requestHook with the span and sanitized query', async () => { + const query = makeQuery(['SELECT ', '']); + await driveHandle(query); + query.resolve({ command: 'SELECT' }); + + const span = lastPgSpan(); + expect(requestHookSpy).toHaveBeenCalledTimes(1); + // No connection context yet (empty registry), so the third arg is undefined here. + expect(requestHookSpy).toHaveBeenCalledWith(span, 'SELECT $1', undefined); + }); + }); + + // NOTE: these tests are order-dependent. The endpoint registry only grows + // (it models connections discovered over the process lifetime), so the + // single-endpoint case must run before the two-endpoint case adds a second. + describe('connection context', () => { + const connectionX = { id: 'x' }; + const optionsX = { host: ['localhost'], port: [5444], database: 'test_db' }; + const connectionY = { id: 'y' }; + const optionsY = { host: ['localhost'], port: [5455], database: 'other_db' }; + + it('with one known endpoint, attaches connection attrs at handle-start and passes context to requestHook', async () => { + publishConnection(connectionX, optionsX); + + const query = makeQuery(['SELECT 1']); + await driveHandle(query); + query.resolve({ command: 'SELECT' }); + + const json = spanToJSON(lastPgSpan()!); + expect(json.data['server.address']).toBe('localhost'); + expect(json.data['server.port']).toBe(5444); + expect(json.data['db.namespace']).toBe('test_db'); + + expect(requestHookSpy).toHaveBeenCalledWith(expect.anything(), 'SELECT ?', { + ATTR_DB_NAMESPACE: 'test_db', + ATTR_SERVER_ADDRESS: 'localhost', + ATTR_SERVER_PORT: '5444', + }); + }); + + it('with multiple known endpoints, does not attach connection attrs at handle-start', async () => { + publishConnection(connectionY, optionsY); + + const query = makeQuery(['SELECT 1']); + await driveHandle(query); + query.resolve({ command: 'SELECT' }); + + const json = spanToJSON(lastPgSpan()!); + expect(json.data['server.address']).toBeUndefined(); + expect(json.data['server.port']).toBeUndefined(); + expect(json.data['db.namespace']).toBeUndefined(); + expect(requestHookSpy).toHaveBeenCalledWith(expect.anything(), 'SELECT ?', undefined); + }); + + it('attaches per-connection attrs via the execute channel when the fallback cannot resolve them', async () => { + // Registry has two endpoints, so handle-start leaves the span without attrs; + // the execute channel resolves them from the connection object. + const query = makeQuery(['SELECT 1']); + await driveHandle(query); + publishExecute(connectionX, query); + query.resolve({ command: 'SELECT' }); + + const json = spanToJSON(lastPgSpan()!); + expect(json.data['server.address']).toBe('localhost'); + expect(json.data['server.port']).toBe(5444); + expect(json.data['db.namespace']).toBe('test_db'); + }); + }); + + // Own integration instance so `requireParentSpan: false` is exercised in + // isolation. Its handle producer replaces the default one on the shared + // channel, so this block runs last and only drives unparented queries. + describe('requireParentSpan: false', () => { + beforeAll(() => { + const integration = postgresJsChannelIntegration({ requireParentSpan: false }); + integration.setupOnce?.(); + }); + + it('creates a span even without an active parent span', async () => { + const query = makeQuery(['SELECT 1']); + await driveHandle(query, { withParent: false }); + query.resolve({ command: 'SELECT' }); + + const span = lastPgSpan(); + expect(span).toBeDefined(); + expect(spanToJSON(span!).description).toBe('SELECT ?'); + expect(spanToJSON(span!).data['db.operation.name']).toBe('SELECT'); + }); + }); +}); From 2ab13492047b0d82c21b2ebc463be493830932da Mon Sep 17 00:00:00 2001 From: Andrei Borza Date: Mon, 6 Jul 2026 21:54:08 +0200 Subject: [PATCH 2/5] Backfill connection attributes for the first query per connection via a connect channel --- .../tracing-channel/postgres-js.ts | 60 +++++++++++++------ .../src/orchestrion/config/postgres.ts | 12 ++++ .../tracing-channel/postgres-js.test.ts | 19 ++++++ 3 files changed, 72 insertions(+), 19 deletions(-) diff --git a/packages/server-utils/src/integrations/tracing-channel/postgres-js.ts b/packages/server-utils/src/integrations/tracing-channel/postgres-js.ts index e01c942cc890..8ea391be41fd 100644 --- a/packages/server-utils/src/integrations/tracing-channel/postgres-js.ts +++ b/packages/server-utils/src/integrations/tracing-channel/postgres-js.ts @@ -158,6 +158,25 @@ function setConnectionAttributes(span: Span, query: PostgresQuery, context: Post } } +/** + * Backfill connection attributes onto a query's span from a channel whose `self` + * is the connection object and `arguments[0]` the query. Shared by the `execute` + * and `connect` channels; both carry that shape and both resolve the context via + * the `connectionContexts` WeakMap. Idempotent (guarded inside `setConnectionAttributes`). + */ +function attachConnectionAttributesFromChannel(message: PostgresJsQueryContext): void { + const connection = message.self as object | undefined; + const query = message.arguments?.[0] as PostgresQuery | undefined; + if (!connection || !query) { + return; + } + const span = (query as Record)[QUERY_SPAN] as Span | undefined; + const context = connectionContexts.get(connection); + if (span && context) { + setConnectionAttributes(span, query, context); + } +} + function setOperationName(span: Span, sanitizedQuery: string | undefined, command?: string): void { if (command) { span.setAttribute(DB_OPERATION_NAME, command); @@ -257,26 +276,29 @@ const _postgresJsChannelIntegration = ((options: PostgresJsChannelIntegrationOpt }, }); + // Per-connection attributes for queries reusing an already-open connection + // (`c.execute(q)`, `self === c`). `execute` is also called bare + // (`self === undefined`) for the first query on each connection, `fetchState` + // and `retry`; those miss here (the `connect` channel below covers the first + // user query, and the single-endpoint fallback covers the common case). diagnosticsChannel.tracingChannel(CHANNELS.POSTGRESJS_EXECUTE).subscribe({ end: NOOP, asyncStart: NOOP, asyncEnd: NOOP, error: NOOP, - start(message) { - // `execute` is also called bare (`self === undefined`) for the first - // query on each connection, `fetchState` and `retry`; those miss here - // and rely on the handle-start fallback instead. - const connection = message.self as object | undefined; - const query = message.arguments?.[0] as PostgresQuery | undefined; - if (!connection || !query) { - return; - } - const span = (query as Record)[QUERY_SPAN] as Span | undefined; - const context = connectionContexts.get(connection); - if (span && context) { - setConnectionAttributes(span, query, context); - } - }, + start: attachConnectionAttributesFromChannel, + }); + + // The connection's `connect(query)` method (`self === c`, `arguments[0]` the + // query) fires when a fresh connection is opened for a query. That first query + // is later dispatched via a bare `execute` (no `self`), so this is where it + // gets its connection attributes in multi-endpoint apps. + diagnosticsChannel.tracingChannel(CHANNELS.POSTGRESJS_CONNECT).subscribe({ + end: NOOP, + asyncStart: NOOP, + asyncEnd: NOOP, + error: NOOP, + start: attachConnectionAttributesFromChannel, }); // The span-creating `handle` subscription needs the async-context binding @@ -363,10 +385,10 @@ const _postgresJsChannelIntegration = ((options: PostgresJsChannelIntegrationOpt /** * EXPERIMENTAL — orchestrion-driven postgres.js (`postgres` v3.x) integration. * - * Subscribes to the `orchestrion:postgres:handle` / `:connection` / `:execute` - * diagnostics channels injected into postgres.js' `Query.prototype.handle` and - * `Connection`/`execute` (in `src/*` and `cjs/src/*`) and creates db spans - * matching the OTel `postgresJsIntegration`. Requires the orchestrion runtime + * Subscribes to the `orchestrion:postgres:handle` / `:connection` / `:execute` / + * `:connect` diagnostics channels injected into postgres.js' `Query.prototype.handle` + * and `Connection`/`execute`/`connect` (in `src/*` and `cjs/src/*`) and creates db + * spans matching the OTel `postgresJsIntegration`. Requires the orchestrion runtime * hook or bundler plugin. */ export const postgresJsChannelIntegration = defineIntegration(_postgresJsChannelIntegration); diff --git a/packages/server-utils/src/orchestrion/config/postgres.ts b/packages/server-utils/src/orchestrion/config/postgres.ts index 6f0cb5639035..7710a7132c6a 100644 --- a/packages/server-utils/src/orchestrion/config/postgres.ts +++ b/packages/server-utils/src/orchestrion/config/postgres.ts @@ -32,6 +32,17 @@ const postgresJsInstrumentationConfig = (dir: string): InstrumentationConfig[] = module: { name: 'postgres', versionRange: '>=3.0.0 <4', filePath: `${dir}/connection.js` }, functionQuery: { functionName: 'execute', kind: 'Sync' }, }, + // The connection object's `connect(query)` method. Matched by `methodName` + // (an object-literal method): `functionName` would hit the unrelated + // socket-level `async function connect()` in the same file. `self` is the + // connection object and `arguments[0]` the query, so the first query that + // opens a connection (dispatched via a bare `execute` with no `self`) still + // gets connection attributes in multi-endpoint apps. + { + channelName: 'connect', + module: { name: 'postgres', versionRange: '>=3.0.0 <4', filePath: `${dir}/connection.js` }, + functionQuery: { methodName: 'connect', kind: 'Sync' }, + }, ]; export const postgresJsConfig = ['src', 'cjs/src'].flatMap(postgresJsInstrumentationConfig); @@ -40,4 +51,5 @@ export const postgresJsChannels = { POSTGRESJS_HANDLE: 'orchestrion:postgres:handle', POSTGRESJS_CONNECTION: 'orchestrion:postgres:connection', POSTGRESJS_EXECUTE: 'orchestrion:postgres:execute', + POSTGRESJS_CONNECT: 'orchestrion:postgres:connect', } as const; diff --git a/packages/server-utils/test/integrations/tracing-channel/postgres-js.test.ts b/packages/server-utils/test/integrations/tracing-channel/postgres-js.test.ts index a427dd31cc20..775d92582e98 100644 --- a/packages/server-utils/test/integrations/tracing-channel/postgres-js.test.ts +++ b/packages/server-utils/test/integrations/tracing-channel/postgres-js.test.ts @@ -138,6 +138,10 @@ function publishExecute(connection: object, query: FakeQuery): void { tracingChannel(CHANNELS.POSTGRESJS_EXECUTE).traceSync(() => undefined, { self: connection, arguments: [query] }); } +function publishConnect(connection: object, query: FakeQuery): void { + tracingChannel(CHANNELS.POSTGRESJS_CONNECT).traceSync(() => undefined, { self: connection, arguments: [query] }); +} + function lastPgSpan(): Span | undefined { return endedSpans.filter(s => spanToJSON(s).data['sentry.origin'] === 'auto.db.orchestrion.postgresjs').at(-1); } @@ -302,6 +306,21 @@ describe('postgresJsChannelIntegration', () => { expect(json.data['server.port']).toBe(5444); expect(json.data['db.namespace']).toBe('test_db'); }); + + it('attaches per-connection attrs via the connect channel (first query opening a connection)', async () => { + // Registry has two endpoints (no handle-start fallback). The first query on a + // fresh connection reaches `execute` bare (no self), but `c.connect(query)` + // carries the connection, so the span still gets its attributes. + const query = makeQuery(['SELECT 1']); + await driveHandle(query); + publishConnect(connectionX, query); + query.resolve({ command: 'SELECT' }); + + const json = spanToJSON(lastPgSpan()!); + expect(json.data['server.address']).toBe('localhost'); + expect(json.data['server.port']).toBe(5444); + expect(json.data['db.namespace']).toBe('test_db'); + }); }); // Own integration instance so `requireParentSpan: false` is exercised in From c65871a85dce50427ec7b739f74ab1bae63d9aa5 Mon Sep 17 00:00:00 2001 From: Andrei Borza Date: Tue, 7 Jul 2026 23:50:54 +0200 Subject: [PATCH 3/5] Extract postgres.js connection and query-span helpers into postgres-js-utils --- .../tracing-channel/postgres-js-utils.ts | 206 ++++++++++++++++ .../tracing-channel/postgres-js.ts | 224 ++---------------- 2 files changed, 222 insertions(+), 208 deletions(-) create mode 100644 packages/server-utils/src/integrations/tracing-channel/postgres-js-utils.ts diff --git a/packages/server-utils/src/integrations/tracing-channel/postgres-js-utils.ts b/packages/server-utils/src/integrations/tracing-channel/postgres-js-utils.ts new file mode 100644 index 000000000000..777413702bdb --- /dev/null +++ b/packages/server-utils/src/integrations/tracing-channel/postgres-js-utils.ts @@ -0,0 +1,206 @@ +import { + DB_NAMESPACE, + DB_OPERATION_NAME, + ERROR_TYPE, + SERVER_ADDRESS, + SERVER_PORT, +} from '@sentry/conventions/attributes'; +import type { Span } from '@sentry/core'; +import { debug, SPAN_STATUS_ERROR } from '@sentry/core'; +import { DEBUG_BUILD } from '../../debug-build'; + +// Not part of `@sentry/conventions`, so we keep it inline (matches the OTel +// `PostgresJsInstrumentation`). +const DB_RESPONSE_STATUS_CODE = 'db.response.status_code'; + +const SQL_OPERATION_REGEX = /^(SELECT|INSERT|UPDATE|DELETE|CREATE|DROP|ALTER)/i; + +// Same `Symbol.for()` marker the core `instrumentPostgresJsSql` wrapper sets on +// queries it manually instruments, so we skip them there and never double-span. +export const QUERY_FROM_INSTRUMENTED_SQL = Symbol.for('sentry.query.from.instrumented.sql'); +// The query span, stashed on the `Query` so the `execute`/`connect` channels can +// attach connection attributes to it. +export const QUERY_SPAN = Symbol('sentryPostgresJsSpan'); +// Set once connection attributes are on the span, so the fallback and the +// `execute`/`connect` channels don't both write them. +const CONNECTION_ATTRS_SET = Symbol('sentryPostgresJsConnectionAttrsSet'); +// Set on the channel context once the resolve/reject wrappers have ended the +// span, so `deferSpanEnd` knows the wrappers own the lifecycle. +export const SPAN_ENDED = Symbol('sentryPostgresJsSpanEnded'); + +/** + * Connection attributes resolved from postgres.js' parsed options. Port is kept + * as a string to match the `requestHook` contract of the OTel integration + * (and the portable `instrumentPostgresJsSql`); it's coerced to a number when + * set on the span (semantic conventions expect a number for `server.port`). + */ +export interface PostgresConnectionContext { + ATTR_DB_NAMESPACE?: string; + ATTR_SERVER_ADDRESS?: string; + ATTR_SERVER_PORT?: string; +} + +/** The `Query` instance postgres.js passes as `self` to `Query.prototype.handle`. */ +export interface PostgresQuery { + strings?: string[]; + executed?: boolean; + resolve?: (...args: unknown[]) => unknown; + reject?: (...args: unknown[]) => unknown; +} + +export interface PostgresJsQueryContext { + arguments?: unknown[]; + self?: PostgresQuery; + result?: unknown; + error?: unknown; +} + +// postgres.js parses `host`/`port` into arrays (it can connect to multiple +// hosts); the `Connection` factory receives this parsed options object. +export interface PostgresParsedOptions { + host?: string[]; + port?: number[]; + database?: string; +} + +// A connection object -> its resolved context, populated on the `connection` +// channel and read on the `execute`/`connect` channels (keyed by the same object). +export const connectionContexts = new WeakMap(); +// Distinct endpoints seen so far (value-compared, so N connections to one DB +// count once). When exactly one endpoint exists — the common case, and the only +// one the tests exercise — every query resolves to it at handle-start. +const endpointRegistry: PostgresConnectionContext[] = []; + +export function registerEndpoint(context: PostgresConnectionContext): void { + const alreadyKnown = endpointRegistry.some( + e => + e.ATTR_SERVER_ADDRESS === context.ATTR_SERVER_ADDRESS && + e.ATTR_SERVER_PORT === context.ATTR_SERVER_PORT && + e.ATTR_DB_NAMESPACE === context.ATTR_DB_NAMESPACE, + ); + if (!alreadyKnown) { + endpointRegistry.push(context); + } +} + +/** The single known endpoint, or `undefined` when zero or multiple are known. */ +export function resolveSingleEndpoint(): PostgresConnectionContext | undefined { + return endpointRegistry.length === 1 ? endpointRegistry[0] : undefined; +} + +export function buildConnectionContext(options: PostgresParsedOptions): PostgresConnectionContext { + // postgres.js defaults to 'localhost'/5432, but be defensive. + const host = options.host?.[0] || 'localhost'; + const port = options.port?.[0] || 5432; + return { + ATTR_DB_NAMESPACE: typeof options.database === 'string' && options.database !== '' ? options.database : undefined, + ATTR_SERVER_ADDRESS: host, + ATTR_SERVER_PORT: String(port), + }; +} + +export function setConnectionAttributes(span: Span, query: PostgresQuery, context: PostgresConnectionContext): void { + const queryRecord = query as Record; + if (queryRecord[CONNECTION_ATTRS_SET]) { + return; + } + queryRecord[CONNECTION_ATTRS_SET] = true; + + if (context.ATTR_SERVER_ADDRESS) { + span.setAttribute(SERVER_ADDRESS, context.ATTR_SERVER_ADDRESS); + } + if (context.ATTR_SERVER_PORT !== undefined) { + const port = parseInt(context.ATTR_SERVER_PORT, 10); + if (!Number.isNaN(port)) { + span.setAttribute(SERVER_PORT, port); + } + } + if (context.ATTR_DB_NAMESPACE) { + span.setAttribute(DB_NAMESPACE, context.ATTR_DB_NAMESPACE); + } +} + +/** + * Backfill connection attributes onto a query's span from a channel whose `self` + * is the connection object and `arguments[0]` the query. Shared by the `execute` + * and `connect` channels; both carry that shape and both resolve the context via + * the `connectionContexts` WeakMap. Idempotent (guarded inside `setConnectionAttributes`). + */ +export function attachConnectionAttributesFromChannel(message: PostgresJsQueryContext): void { + const connection = message.self as object | undefined; + const query = message.arguments?.[0] as PostgresQuery | undefined; + if (!connection || !query) { + return; + } + const span = (query as Record)[QUERY_SPAN] as Span | undefined; + const context = connectionContexts.get(connection); + if (span && context) { + setConnectionAttributes(span, query, context); + } +} + +function setOperationName(span: Span, sanitizedQuery: string | undefined, command?: string): void { + if (command) { + span.setAttribute(DB_OPERATION_NAME, command); + return; + } + const operationMatch = sanitizedQuery?.match(SQL_OPERATION_REGEX); + if (operationMatch?.[1]) { + span.setAttribute(DB_OPERATION_NAME, operationMatch[1].toUpperCase()); + } +} + +/** + * Wrap `query.resolve`/`query.reject` so the span ends when the query settles. + * + * `Query extends Promise` and `async handle()` only dispatches — its promise + * resolves immediately, long before the query completes. postgres.js signals + * completion by calling `this.resolve`/`this.reject`, so we own the span end + * there. Wrapping happens at handle-start because `reject` can fire + * synchronously during dispatch and `cursor()` reassigns both before executing. + */ +export function wrapQuerySettlement(data: PostgresJsQueryContext, span: Span, sanitizedSqlQuery: string): void { + const query = data.self; + if (!query) { + return; + } + + // Claim ownership of ending the span up front, so `deferSpanEnd` defers to the + // wrapper even if `span.end()` below throws. + const markEnded = (): void => { + (data as Record)[SPAN_ENDED] = true; + }; + + const originalResolve = query.resolve; + if (typeof originalResolve === 'function') { + query.resolve = function (this: unknown, ...resolveArgs: unknown[]): unknown { + markEnded(); + try { + const command = (resolveArgs[0] as { command?: string } | undefined)?.command; + setOperationName(span, sanitizedSqlQuery, command); + span.end(); + } catch (e) { + DEBUG_BUILD && debug.error('[orchestrion:postgresjs] error ending span in resolve:', e); + } + return originalResolve.apply(this, resolveArgs); + }; + } + + const originalReject = query.reject; + if (typeof originalReject === 'function') { + query.reject = function (this: unknown, ...rejectArgs: unknown[]): unknown { + markEnded(); + try { + const err = rejectArgs[0] as { message?: string; code?: string; name?: string } | undefined; + span.setStatus({ code: SPAN_STATUS_ERROR, message: err?.message || 'unknown_error' }); + span.setAttribute(DB_RESPONSE_STATUS_CODE, err?.code || 'unknown'); + span.setAttribute(ERROR_TYPE, err?.name || 'unknown'); + setOperationName(span, sanitizedSqlQuery); + span.end(); + } catch (e) { + DEBUG_BUILD && debug.error('[orchestrion:postgresjs] error ending span in reject:', e); + } + return originalReject.apply(this, rejectArgs); + }; + } +} diff --git a/packages/server-utils/src/integrations/tracing-channel/postgres-js.ts b/packages/server-utils/src/integrations/tracing-channel/postgres-js.ts index 8ea391be41fd..24917de6fe17 100644 --- a/packages/server-utils/src/integrations/tracing-channel/postgres-js.ts +++ b/packages/server-utils/src/integrations/tracing-channel/postgres-js.ts @@ -1,13 +1,5 @@ import * as diagnosticsChannel from 'node:diagnostics_channel'; -import { - DB_NAMESPACE, - DB_OPERATION_NAME, - DB_QUERY_TEXT, - DB_SYSTEM_NAME, - ERROR_TYPE, - SERVER_ADDRESS, - SERVER_PORT, -} from '@sentry/conventions/attributes'; +import { DB_QUERY_TEXT, DB_SYSTEM_NAME } from '@sentry/conventions/attributes'; import type { IntegrationFn, Span } from '@sentry/core'; import { _INTERNAL_reconstructPostgresQuery, @@ -17,13 +9,25 @@ import { getActiveSpan, SEMANTIC_ATTRIBUTE_SENTRY_ORIGIN, SPAN_KIND, - SPAN_STATUS_ERROR, startInactiveSpan, waitForTracingChannelBinding, } from '@sentry/core'; import { DEBUG_BUILD } from '../../debug-build'; import { CHANNELS } from '../../orchestrion/channels'; import { bindTracingChannelToSpan } from '../../tracing-channel'; +import type { PostgresConnectionContext, PostgresJsQueryContext, PostgresParsedOptions } from './postgres-js-utils'; +import { + attachConnectionAttributesFromChannel, + buildConnectionContext, + connectionContexts, + QUERY_FROM_INSTRUMENTED_SQL, + QUERY_SPAN, + registerEndpoint, + resolveSingleEndpoint, + setConnectionAttributes, + SPAN_ENDED, + wrapQuerySettlement, +} from './postgres-js-utils'; // Same name as the OTel `PostgresJs` integration by design: when this is // enabled, the OTel integration of the same name is dropped from the default @@ -32,36 +36,7 @@ const INTEGRATION_NAME = 'PostgresJs' as const; const ORIGIN = 'auto.db.orchestrion.postgresjs'; -// Not part of `@sentry/conventions`, so we keep it inline (matches the OTel -// `PostgresJsInstrumentation`). -const DB_RESPONSE_STATUS_CODE = 'db.response.status_code'; - -const SQL_OPERATION_REGEX = /^(SELECT|INSERT|UPDATE|DELETE|CREATE|DROP|ALTER)/i; - -// Same `Symbol.for()` marker the core `instrumentPostgresJsSql` wrapper sets on -// queries it manually instruments, so we skip them here and never double-span. -const QUERY_FROM_INSTRUMENTED_SQL = Symbol.for('sentry.query.from.instrumented.sql'); -// The query span, stashed on the `Query` so the `execute` channel can attach -// connection attributes to it. -const QUERY_SPAN = Symbol('sentryPostgresJsSpan'); -// Set once connection attributes are on the span, so the fallback and the -// `execute` channel don't both write them. -const CONNECTION_ATTRS_SET = Symbol('sentryPostgresJsConnectionAttrsSet'); -// Set on the channel context once the resolve/reject wrappers have ended the -// span, so `deferSpanEnd` knows the wrappers own the lifecycle. -const SPAN_ENDED = Symbol('sentryPostgresJsSpanEnded'); - -/** - * Connection attributes resolved from postgres.js' parsed options. Port is kept - * as a string to match the `requestHook` contract of the OTel integration - * (and the portable `instrumentPostgresJsSql`); it's coerced to a number when - * set on the span (semantic conventions expect a number for `server.port`). - */ -interface PostgresConnectionContext { - ATTR_DB_NAMESPACE?: string; - ATTR_SERVER_ADDRESS?: string; - ATTR_SERVER_PORT?: string; -} +const NOOP = (): void => {}; export interface PostgresJsChannelIntegrationOptions { /** @@ -76,173 +51,6 @@ export interface PostgresJsChannelIntegrationOptions { requestHook?: (span: Span, sanitizedSqlQuery: string, postgresConnectionContext?: PostgresConnectionContext) => void; } -/** The `Query` instance postgres.js passes as `self` to `Query.prototype.handle`. */ -interface PostgresQuery { - strings?: string[]; - executed?: boolean; - resolve?: (...args: unknown[]) => unknown; - reject?: (...args: unknown[]) => unknown; -} - -interface PostgresJsQueryContext { - arguments?: unknown[]; - self?: PostgresQuery; - result?: unknown; - error?: unknown; -} - -// postgres.js parses `host`/`port` into arrays (it can connect to multiple -// hosts); the `Connection` factory receives this parsed options object. -interface PostgresParsedOptions { - host?: string[]; - port?: number[]; - database?: string; -} - -const NOOP = (): void => {}; - -// A connection object -> its resolved context, populated on the `connection` -// channel and read on the `execute` channel (keyed by the same object). -const connectionContexts = new WeakMap(); -// Distinct endpoints seen so far (value-compared, so N connections to one DB -// count once). When exactly one endpoint exists — the common case, and the only -// one the tests exercise — every query resolves to it at handle-start. -const endpointRegistry: PostgresConnectionContext[] = []; - -function registerEndpoint(context: PostgresConnectionContext): void { - const alreadyKnown = endpointRegistry.some( - e => - e.ATTR_SERVER_ADDRESS === context.ATTR_SERVER_ADDRESS && - e.ATTR_SERVER_PORT === context.ATTR_SERVER_PORT && - e.ATTR_DB_NAMESPACE === context.ATTR_DB_NAMESPACE, - ); - if (!alreadyKnown) { - endpointRegistry.push(context); - } -} - -/** The single known endpoint, or `undefined` when zero or multiple are known. */ -function resolveSingleEndpoint(): PostgresConnectionContext | undefined { - return endpointRegistry.length === 1 ? endpointRegistry[0] : undefined; -} - -function buildConnectionContext(options: PostgresParsedOptions): PostgresConnectionContext { - // postgres.js defaults to 'localhost'/5432, but be defensive. - const host = options.host?.[0] || 'localhost'; - const port = options.port?.[0] || 5432; - return { - ATTR_DB_NAMESPACE: typeof options.database === 'string' && options.database !== '' ? options.database : undefined, - ATTR_SERVER_ADDRESS: host, - ATTR_SERVER_PORT: String(port), - }; -} - -function setConnectionAttributes(span: Span, query: PostgresQuery, context: PostgresConnectionContext): void { - const queryRecord = query as Record; - if (queryRecord[CONNECTION_ATTRS_SET]) { - return; - } - queryRecord[CONNECTION_ATTRS_SET] = true; - - if (context.ATTR_SERVER_ADDRESS) { - span.setAttribute(SERVER_ADDRESS, context.ATTR_SERVER_ADDRESS); - } - if (context.ATTR_SERVER_PORT !== undefined) { - const port = parseInt(context.ATTR_SERVER_PORT, 10); - if (!Number.isNaN(port)) { - span.setAttribute(SERVER_PORT, port); - } - } - if (context.ATTR_DB_NAMESPACE) { - span.setAttribute(DB_NAMESPACE, context.ATTR_DB_NAMESPACE); - } -} - -/** - * Backfill connection attributes onto a query's span from a channel whose `self` - * is the connection object and `arguments[0]` the query. Shared by the `execute` - * and `connect` channels; both carry that shape and both resolve the context via - * the `connectionContexts` WeakMap. Idempotent (guarded inside `setConnectionAttributes`). - */ -function attachConnectionAttributesFromChannel(message: PostgresJsQueryContext): void { - const connection = message.self as object | undefined; - const query = message.arguments?.[0] as PostgresQuery | undefined; - if (!connection || !query) { - return; - } - const span = (query as Record)[QUERY_SPAN] as Span | undefined; - const context = connectionContexts.get(connection); - if (span && context) { - setConnectionAttributes(span, query, context); - } -} - -function setOperationName(span: Span, sanitizedQuery: string | undefined, command?: string): void { - if (command) { - span.setAttribute(DB_OPERATION_NAME, command); - return; - } - const operationMatch = sanitizedQuery?.match(SQL_OPERATION_REGEX); - if (operationMatch?.[1]) { - span.setAttribute(DB_OPERATION_NAME, operationMatch[1].toUpperCase()); - } -} - -/** - * Wrap `query.resolve`/`query.reject` so the span ends when the query settles. - * - * `Query extends Promise` and `async handle()` only dispatches — its promise - * resolves immediately, long before the query completes. postgres.js signals - * completion by calling `this.resolve`/`this.reject`, so we own the span end - * there. Wrapping happens at handle-start because `reject` can fire - * synchronously during dispatch and `cursor()` reassigns both before executing. - */ -function wrapQuerySettlement(data: PostgresJsQueryContext, span: Span, sanitizedSqlQuery: string): void { - const query = data.self; - if (!query) { - return; - } - - // Claim ownership of ending the span up front, so `deferSpanEnd` defers to the - // wrapper even if `span.end()` below throws. - const markEnded = (): void => { - (data as Record)[SPAN_ENDED] = true; - }; - - const originalResolve = query.resolve; - if (typeof originalResolve === 'function') { - query.resolve = function (this: unknown, ...resolveArgs: unknown[]): unknown { - markEnded(); - try { - const command = (resolveArgs[0] as { command?: string } | undefined)?.command; - setOperationName(span, sanitizedSqlQuery, command); - span.end(); - } catch (e) { - DEBUG_BUILD && debug.error('[orchestrion:postgresjs] error ending span in resolve:', e); - } - return originalResolve.apply(this, resolveArgs); - }; - } - - const originalReject = query.reject; - if (typeof originalReject === 'function') { - query.reject = function (this: unknown, ...rejectArgs: unknown[]): unknown { - markEnded(); - try { - const err = rejectArgs[0] as { message?: string; code?: string; name?: string } | undefined; - span.setStatus({ code: SPAN_STATUS_ERROR, message: err?.message || 'unknown_error' }); - span.setAttribute(DB_RESPONSE_STATUS_CODE, err?.code || 'unknown'); - span.setAttribute(ERROR_TYPE, err?.name || 'unknown'); - setOperationName(span, sanitizedSqlQuery); - span.end(); - } catch (e) { - DEBUG_BUILD && debug.error('[orchestrion:postgresjs] error ending span in reject:', e); - } - return originalReject.apply(this, rejectArgs); - }; - } -} - const _postgresJsChannelIntegration = ((options: PostgresJsChannelIntegrationOptions = {}) => { const { requireParentSpan, requestHook } = options; @@ -337,7 +145,7 @@ const _postgresJsChannelIntegration = ((options: PostgresJsChannelIntegrationOpt }, }); - // Stash for the `execute` channel to attach per-connection attributes. + // Stash for the `execute`/`connect` channels to attach per-connection attributes. (query as Record)[QUERY_SPAN] = span; // Single-endpoint fallback: resolve context now so `requestHook` has it From 2e903c6bde1f6f8865b3f7ecc7484c8113c40050 Mon Sep 17 00:00:00 2001 From: Andrei Borza Date: Wed, 8 Jul 2026 00:10:34 +0200 Subject: [PATCH 4/5] Move postgres.js integration into a folder and reuse core query/connection helpers --- packages/core/src/integrations/postgresjs.ts | 44 ++++--- packages/core/src/server-exports.ts | 4 + .../{postgres-js.ts => postgres-js/index.ts} | 27 ++-- .../tracing-channel/postgres-js/types.ts | 22 ++++ .../utils.ts} | 117 +++++------------- 5 files changed, 93 insertions(+), 121 deletions(-) rename packages/server-utils/src/integrations/tracing-channel/{postgres-js.ts => postgres-js/index.ts} (89%) create mode 100644 packages/server-utils/src/integrations/tracing-channel/postgres-js/types.ts rename packages/server-utils/src/integrations/tracing-channel/{postgres-js-utils.ts => postgres-js/utils.ts} (63%) diff --git a/packages/core/src/integrations/postgresjs.ts b/packages/core/src/integrations/postgresjs.ts index 93602d71a6e5..bf3a3fb8337d 100644 --- a/packages/core/src/integrations/postgresjs.ts +++ b/packages/core/src/integrations/postgresjs.ts @@ -11,7 +11,7 @@ import { getActiveSpan } from '../utils/spanUtils'; const SQL_OPERATION_REGEX = /^(SELECT|INSERT|UPDATE|DELETE|CREATE|DROP|ALTER)/i; -type PostgresConnectionContext = { +export type PostgresConnectionContext = { ATTR_DB_NAMESPACE?: string; ATTR_SERVER_ADDRESS?: string; ATTR_SERVER_PORT?: string; @@ -398,8 +398,10 @@ export function _sanitizeSqlQuery(sqlQuery: string | undefined): string { /** * Sets connection context attributes on a span. + * + * @internal Exported for the orchestrion (diagnostics-channel) integration. */ -function _setConnectionAttributes(span: Span, connectionContext: PostgresConnectionContext | undefined): void { +export function _setConnectionAttributes(span: Span, connectionContext: PostgresConnectionContext | undefined): void { if (!connectionContext) { return; } @@ -421,8 +423,10 @@ function _setConnectionAttributes(span: Span, connectionContext: PostgresConnect /** * Extracts DB operation name from SQL query and sets it on the span. + * + * @internal Exported for the orchestrion (diagnostics-channel) integration. */ -function _setOperationName(span: Span, sanitizedQuery: string | undefined, command?: string): void { +export function _setOperationName(span: Span, sanitizedQuery: string | undefined, command?: string): void { if (command) { span.setAttribute('db.operation.name', command); return; @@ -434,6 +438,26 @@ function _setOperationName(span: Span, sanitizedQuery: string | undefined, comma } } +/** + * Builds a {@link PostgresConnectionContext} from postgres.js' parsed options + * (which store `host`/`port` as arrays). Defaults to 'localhost'/5432. + * + * @internal Exported for the orchestrion (diagnostics-channel) integration. + */ +export function _buildConnectionContext(options: { + host?: string[]; + port?: number[]; + database?: string; +}): PostgresConnectionContext { + const host = options.host?.[0] || 'localhost'; + const port = options.port?.[0] || 5432; + return { + ATTR_DB_NAMESPACE: typeof options.database === 'string' && options.database !== '' ? options.database : undefined, + ATTR_SERVER_ADDRESS: host, + ATTR_SERVER_PORT: String(port), + }; +} + /** * Extracts and stores connection context from sql.options. */ @@ -443,17 +467,5 @@ function _attachConnectionContext(sql: unknown, proxiedSql: Record unknown; + reject?: (...args: unknown[]) => unknown; +} + +export interface PostgresJsQueryContext { + arguments?: unknown[]; + self?: PostgresQuery; + result?: unknown; + error?: unknown; +} + +// postgres.js parses `host`/`port` into arrays (it can connect to multiple +// hosts); the `Connection` factory receives this parsed options object. +export interface PostgresParsedOptions { + host?: string[]; + port?: number[]; + database?: string; +} diff --git a/packages/server-utils/src/integrations/tracing-channel/postgres-js-utils.ts b/packages/server-utils/src/integrations/tracing-channel/postgres-js/utils.ts similarity index 63% rename from packages/server-utils/src/integrations/tracing-channel/postgres-js-utils.ts rename to packages/server-utils/src/integrations/tracing-channel/postgres-js/utils.ts index 777413702bdb..ed22034344fe 100644 --- a/packages/server-utils/src/integrations/tracing-channel/postgres-js-utils.ts +++ b/packages/server-utils/src/integrations/tracing-channel/postgres-js/utils.ts @@ -1,20 +1,18 @@ +import { ERROR_TYPE } from '@sentry/conventions/attributes'; +import type { PostgresConnectionContext, Span } from '@sentry/core'; import { - DB_NAMESPACE, - DB_OPERATION_NAME, - ERROR_TYPE, - SERVER_ADDRESS, - SERVER_PORT, -} from '@sentry/conventions/attributes'; -import type { Span } from '@sentry/core'; -import { debug, SPAN_STATUS_ERROR } from '@sentry/core'; -import { DEBUG_BUILD } from '../../debug-build'; - -// Not part of `@sentry/conventions`, so we keep it inline (matches the OTel -// `PostgresJsInstrumentation`). + _INTERNAL_buildPostgresConnectionContext, + _INTERNAL_setPostgresConnectionAttributes, + _INTERNAL_setPostgresOperationName, + debug, + SPAN_STATUS_ERROR, +} from '@sentry/core'; +import { DEBUG_BUILD } from '../../../debug-build'; +import type { PostgresJsQueryContext, PostgresParsedOptions, PostgresQuery } from './types'; + +// Not part of `@sentry/conventions`, so we keep it inline (matches the OTel `PostgresJsInstrumentation`). const DB_RESPONSE_STATUS_CODE = 'db.response.status_code'; -const SQL_OPERATION_REGEX = /^(SELECT|INSERT|UPDATE|DELETE|CREATE|DROP|ALTER)/i; - // Same `Symbol.for()` marker the core `instrumentPostgresJsSql` wrapper sets on // queries it manually instruments, so we skip them there and never double-span. export const QUERY_FROM_INSTRUMENTED_SQL = Symbol.for('sentry.query.from.instrumented.sql'); @@ -28,50 +26,15 @@ const CONNECTION_ATTRS_SET = Symbol('sentryPostgresJsConnectionAttrsSet'); // span, so `deferSpanEnd` knows the wrappers own the lifecycle. export const SPAN_ENDED = Symbol('sentryPostgresJsSpanEnded'); -/** - * Connection attributes resolved from postgres.js' parsed options. Port is kept - * as a string to match the `requestHook` contract of the OTel integration - * (and the portable `instrumentPostgresJsSql`); it's coerced to a number when - * set on the span (semantic conventions expect a number for `server.port`). - */ -export interface PostgresConnectionContext { - ATTR_DB_NAMESPACE?: string; - ATTR_SERVER_ADDRESS?: string; - ATTR_SERVER_PORT?: string; -} - -/** The `Query` instance postgres.js passes as `self` to `Query.prototype.handle`. */ -export interface PostgresQuery { - strings?: string[]; - executed?: boolean; - resolve?: (...args: unknown[]) => unknown; - reject?: (...args: unknown[]) => unknown; -} - -export interface PostgresJsQueryContext { - arguments?: unknown[]; - self?: PostgresQuery; - result?: unknown; - error?: unknown; -} - -// postgres.js parses `host`/`port` into arrays (it can connect to multiple -// hosts); the `Connection` factory receives this parsed options object. -export interface PostgresParsedOptions { - host?: string[]; - port?: number[]; - database?: string; -} - // A connection object -> its resolved context, populated on the `connection` // channel and read on the `execute`/`connect` channels (keyed by the same object). -export const connectionContexts = new WeakMap(); +const connectionContexts = new WeakMap(); // Distinct endpoints seen so far (value-compared, so N connections to one DB // count once). When exactly one endpoint exists — the common case, and the only // one the tests exercise — every query resolves to it at handle-start. const endpointRegistry: PostgresConnectionContext[] = []; -export function registerEndpoint(context: PostgresConnectionContext): void { +function registerEndpoint(context: PostgresConnectionContext): void { const alreadyKnown = endpointRegistry.some( e => e.ATTR_SERVER_ADDRESS === context.ATTR_SERVER_ADDRESS && @@ -88,15 +51,20 @@ export function resolveSingleEndpoint(): PostgresConnectionContext | undefined { return endpointRegistry.length === 1 ? endpointRegistry[0] : undefined; } -export function buildConnectionContext(options: PostgresParsedOptions): PostgresConnectionContext { - // postgres.js defaults to 'localhost'/5432, but be defensive. - const host = options.host?.[0] || 'localhost'; - const port = options.port?.[0] || 5432; - return { - ATTR_DB_NAMESPACE: typeof options.database === 'string' && options.database !== '' ? options.database : undefined, - ATTR_SERVER_ADDRESS: host, - ATTR_SERVER_PORT: String(port), - }; +/** + * Record a connection from the `connection` channel `end` (`result` is the + * connection object, `arguments[0]` the parsed options), keying its resolved + * context by the connection object and tracking its endpoint. + */ +export function recordConnectionFromChannel(message: PostgresJsQueryContext): void { + const connection = message.result; + const options = message.arguments?.[0] as PostgresParsedOptions | undefined; + if (!connection || typeof connection !== 'object' || !options) { + return; + } + const context = _INTERNAL_buildPostgresConnectionContext(options); + connectionContexts.set(connection, context); + registerEndpoint(context); } export function setConnectionAttributes(span: Span, query: PostgresQuery, context: PostgresConnectionContext): void { @@ -105,19 +73,7 @@ export function setConnectionAttributes(span: Span, query: PostgresQuery, contex return; } queryRecord[CONNECTION_ATTRS_SET] = true; - - if (context.ATTR_SERVER_ADDRESS) { - span.setAttribute(SERVER_ADDRESS, context.ATTR_SERVER_ADDRESS); - } - if (context.ATTR_SERVER_PORT !== undefined) { - const port = parseInt(context.ATTR_SERVER_PORT, 10); - if (!Number.isNaN(port)) { - span.setAttribute(SERVER_PORT, port); - } - } - if (context.ATTR_DB_NAMESPACE) { - span.setAttribute(DB_NAMESPACE, context.ATTR_DB_NAMESPACE); - } + _INTERNAL_setPostgresConnectionAttributes(span, context); } /** @@ -139,17 +95,6 @@ export function attachConnectionAttributesFromChannel(message: PostgresJsQueryCo } } -function setOperationName(span: Span, sanitizedQuery: string | undefined, command?: string): void { - if (command) { - span.setAttribute(DB_OPERATION_NAME, command); - return; - } - const operationMatch = sanitizedQuery?.match(SQL_OPERATION_REGEX); - if (operationMatch?.[1]) { - span.setAttribute(DB_OPERATION_NAME, operationMatch[1].toUpperCase()); - } -} - /** * Wrap `query.resolve`/`query.reject` so the span ends when the query settles. * @@ -177,7 +122,7 @@ export function wrapQuerySettlement(data: PostgresJsQueryContext, span: Span, sa markEnded(); try { const command = (resolveArgs[0] as { command?: string } | undefined)?.command; - setOperationName(span, sanitizedSqlQuery, command); + _INTERNAL_setPostgresOperationName(span, sanitizedSqlQuery, command); span.end(); } catch (e) { DEBUG_BUILD && debug.error('[orchestrion:postgresjs] error ending span in resolve:', e); @@ -195,7 +140,7 @@ export function wrapQuerySettlement(data: PostgresJsQueryContext, span: Span, sa span.setStatus({ code: SPAN_STATUS_ERROR, message: err?.message || 'unknown_error' }); span.setAttribute(DB_RESPONSE_STATUS_CODE, err?.code || 'unknown'); span.setAttribute(ERROR_TYPE, err?.name || 'unknown'); - setOperationName(span, sanitizedSqlQuery); + _INTERNAL_setPostgresOperationName(span, sanitizedSqlQuery); span.end(); } catch (e) { DEBUG_BUILD && debug.error('[orchestrion:postgresjs] error ending span in reject:', e); From bc2b480e4bac1ebe2a8ad38bea6e963e92080a34 Mon Sep 17 00:00:00 2001 From: Andrei Borza Date: Wed, 8 Jul 2026 00:15:39 +0200 Subject: [PATCH 5/5] Inline postgres.js integration back into a single file --- .../{postgres-js/index.ts => postgres-js.ts} | 178 ++++++++++++++++-- .../tracing-channel/postgres-js/types.ts | 22 --- .../tracing-channel/postgres-js/utils.ts | 151 --------------- 3 files changed, 163 insertions(+), 188 deletions(-) rename packages/server-utils/src/integrations/tracing-channel/{postgres-js/index.ts => postgres-js.ts} (51%) delete mode 100644 packages/server-utils/src/integrations/tracing-channel/postgres-js/types.ts delete mode 100644 packages/server-utils/src/integrations/tracing-channel/postgres-js/utils.ts diff --git a/packages/server-utils/src/integrations/tracing-channel/postgres-js/index.ts b/packages/server-utils/src/integrations/tracing-channel/postgres-js.ts similarity index 51% rename from packages/server-utils/src/integrations/tracing-channel/postgres-js/index.ts rename to packages/server-utils/src/integrations/tracing-channel/postgres-js.ts index 251125815627..0370dba92454 100644 --- a/packages/server-utils/src/integrations/tracing-channel/postgres-js/index.ts +++ b/packages/server-utils/src/integrations/tracing-channel/postgres-js.ts @@ -1,31 +1,24 @@ import * as diagnosticsChannel from 'node:diagnostics_channel'; -import { DB_QUERY_TEXT, DB_SYSTEM_NAME } from '@sentry/conventions/attributes'; +import { DB_QUERY_TEXT, DB_SYSTEM_NAME, ERROR_TYPE } from '@sentry/conventions/attributes'; import type { IntegrationFn, PostgresConnectionContext, Span } from '@sentry/core'; import { + _INTERNAL_buildPostgresConnectionContext, _INTERNAL_reconstructPostgresQuery, _INTERNAL_sanitizeSqlQuery, + _INTERNAL_setPostgresConnectionAttributes, + _INTERNAL_setPostgresOperationName, debug, defineIntegration, getActiveSpan, SEMANTIC_ATTRIBUTE_SENTRY_ORIGIN, SPAN_KIND, + SPAN_STATUS_ERROR, startInactiveSpan, waitForTracingChannelBinding, } from '@sentry/core'; -import { DEBUG_BUILD } from '../../../debug-build'; -import { CHANNELS } from '../../../orchestrion/channels'; -import { bindTracingChannelToSpan } from '../../../tracing-channel'; -import type { PostgresJsQueryContext } from './types'; -import { - attachConnectionAttributesFromChannel, - QUERY_FROM_INSTRUMENTED_SQL, - QUERY_SPAN, - recordConnectionFromChannel, - resolveSingleEndpoint, - setConnectionAttributes, - SPAN_ENDED, - wrapQuerySettlement, -} from './utils'; +import { DEBUG_BUILD } from '../../debug-build'; +import { CHANNELS } from '../../orchestrion/channels'; +import { bindTracingChannelToSpan } from '../../tracing-channel'; // Same name as the OTel `PostgresJs` integration by design: when this is // enabled, the OTel integration of the same name is dropped from the default @@ -34,8 +27,24 @@ const INTEGRATION_NAME = 'PostgresJs' as const; const ORIGIN = 'auto.db.orchestrion.postgresjs'; +// Not part of `@sentry/conventions`, so we keep it inline (matches the OTel `PostgresJsInstrumentation`). +const DB_RESPONSE_STATUS_CODE = 'db.response.status_code'; + const NOOP = (): void => {}; +// Same `Symbol.for()` marker the core `instrumentPostgresJsSql` wrapper sets on +// queries it manually instruments, so we skip them there and never double-span. +const QUERY_FROM_INSTRUMENTED_SQL = Symbol.for('sentry.query.from.instrumented.sql'); +// The query span, stashed on the `Query` so the `execute`/`connect` channels can +// attach connection attributes to it. +const QUERY_SPAN = Symbol('sentryPostgresJsSpan'); +// Set once connection attributes are on the span, so the fallback and the +// `execute`/`connect` channels don't both write them. +const CONNECTION_ATTRS_SET = Symbol('sentryPostgresJsConnectionAttrsSet'); +// Set on the channel context once the resolve/reject wrappers have ended the +// span, so `deferSpanEnd` knows the wrappers own the lifecycle. +const SPAN_ENDED = Symbol('sentryPostgresJsSpanEnded'); + export interface PostgresJsChannelIntegrationOptions { /** * Only create spans when there's already an active parent span. Defaults to @@ -49,6 +58,145 @@ export interface PostgresJsChannelIntegrationOptions { requestHook?: (span: Span, sanitizedSqlQuery: string, postgresConnectionContext?: PostgresConnectionContext) => void; } +/** The `Query` instance postgres.js passes as `self` to `Query.prototype.handle`. */ +interface PostgresQuery { + strings?: string[]; + executed?: boolean; + resolve?: (...args: unknown[]) => unknown; + reject?: (...args: unknown[]) => unknown; +} + +interface PostgresJsQueryContext { + arguments?: unknown[]; + self?: PostgresQuery; + result?: unknown; + error?: unknown; +} + +// A connection object -> its resolved context, populated on the `connection` +// channel and read on the `execute`/`connect` channels (keyed by the same object). +const connectionContexts = new WeakMap(); +// Distinct endpoints seen so far (value-compared, so N connections to one DB +// count once). When exactly one endpoint exists — the common case, and the only +// one the tests exercise — every query resolves to it at handle-start. +const endpointRegistry: PostgresConnectionContext[] = []; + +function registerEndpoint(context: PostgresConnectionContext): void { + const alreadyKnown = endpointRegistry.some( + e => + e.ATTR_SERVER_ADDRESS === context.ATTR_SERVER_ADDRESS && + e.ATTR_SERVER_PORT === context.ATTR_SERVER_PORT && + e.ATTR_DB_NAMESPACE === context.ATTR_DB_NAMESPACE, + ); + if (!alreadyKnown) { + endpointRegistry.push(context); + } +} + +/** The single known endpoint, or `undefined` when zero or multiple are known. */ +function resolveSingleEndpoint(): PostgresConnectionContext | undefined { + return endpointRegistry.length === 1 ? endpointRegistry[0] : undefined; +} + +/** + * Record a connection from the `connection` channel `end` (`result` is the + * connection object, `arguments[0]` the parsed options), keying its resolved + * context by the connection object and tracking its endpoint. + */ +function recordConnectionFromChannel(message: PostgresJsQueryContext): void { + const connection = message.result; + const options = message.arguments?.[0] as { host?: string[]; port?: number[]; database?: string } | undefined; + if (!connection || typeof connection !== 'object' || !options) { + return; + } + const context = _INTERNAL_buildPostgresConnectionContext(options); + connectionContexts.set(connection, context); + registerEndpoint(context); +} + +function setConnectionAttributes(span: Span, query: PostgresQuery, context: PostgresConnectionContext): void { + const queryRecord = query as Record; + if (queryRecord[CONNECTION_ATTRS_SET]) { + return; + } + queryRecord[CONNECTION_ATTRS_SET] = true; + _INTERNAL_setPostgresConnectionAttributes(span, context); +} + +/** + * Backfill connection attributes onto a query's span from a channel whose `self` + * is the connection object and `arguments[0]` the query. Shared by the `execute` + * and `connect` channels; both carry that shape and both resolve the context via + * the `connectionContexts` WeakMap. Idempotent (guarded inside `setConnectionAttributes`). + */ +function attachConnectionAttributesFromChannel(message: PostgresJsQueryContext): void { + const connection = message.self as object | undefined; + const query = message.arguments?.[0] as PostgresQuery | undefined; + if (!connection || !query) { + return; + } + const span = (query as Record)[QUERY_SPAN] as Span | undefined; + const context = connectionContexts.get(connection); + if (span && context) { + setConnectionAttributes(span, query, context); + } +} + +/** + * Wrap `query.resolve`/`query.reject` so the span ends when the query settles. + * + * `Query extends Promise` and `async handle()` only dispatches — its promise + * resolves immediately, long before the query completes. postgres.js signals + * completion by calling `this.resolve`/`this.reject`, so we own the span end + * there. Wrapping happens at handle-start because `reject` can fire + * synchronously during dispatch and `cursor()` reassigns both before executing. + */ +function wrapQuerySettlement(data: PostgresJsQueryContext, span: Span, sanitizedSqlQuery: string): void { + const query = data.self; + if (!query) { + return; + } + + // Claim ownership of ending the span up front, so `deferSpanEnd` defers to the + // wrapper even if `span.end()` below throws. + const markEnded = (): void => { + (data as Record)[SPAN_ENDED] = true; + }; + + const originalResolve = query.resolve; + if (typeof originalResolve === 'function') { + query.resolve = function (this: unknown, ...resolveArgs: unknown[]): unknown { + markEnded(); + try { + const command = (resolveArgs[0] as { command?: string } | undefined)?.command; + _INTERNAL_setPostgresOperationName(span, sanitizedSqlQuery, command); + span.end(); + } catch (e) { + DEBUG_BUILD && debug.error('[orchestrion:postgresjs] error ending span in resolve:', e); + } + return originalResolve.apply(this, resolveArgs); + }; + } + + const originalReject = query.reject; + if (typeof originalReject === 'function') { + query.reject = function (this: unknown, ...rejectArgs: unknown[]): unknown { + markEnded(); + try { + const err = rejectArgs[0] as { message?: string; code?: string; name?: string } | undefined; + span.setStatus({ code: SPAN_STATUS_ERROR, message: err?.message || 'unknown_error' }); + span.setAttribute(DB_RESPONSE_STATUS_CODE, err?.code || 'unknown'); + span.setAttribute(ERROR_TYPE, err?.name || 'unknown'); + _INTERNAL_setPostgresOperationName(span, sanitizedSqlQuery); + span.end(); + } catch (e) { + DEBUG_BUILD && debug.error('[orchestrion:postgresjs] error ending span in reject:', e); + } + return originalReject.apply(this, rejectArgs); + }; + } +} + const _postgresJsChannelIntegration = ((options: PostgresJsChannelIntegrationOptions = {}) => { const { requireParentSpan, requestHook } = options; diff --git a/packages/server-utils/src/integrations/tracing-channel/postgres-js/types.ts b/packages/server-utils/src/integrations/tracing-channel/postgres-js/types.ts deleted file mode 100644 index 8fbd20f528ec..000000000000 --- a/packages/server-utils/src/integrations/tracing-channel/postgres-js/types.ts +++ /dev/null @@ -1,22 +0,0 @@ -/** The `Query` instance postgres.js passes as `self` to `Query.prototype.handle`. */ -export interface PostgresQuery { - strings?: string[]; - executed?: boolean; - resolve?: (...args: unknown[]) => unknown; - reject?: (...args: unknown[]) => unknown; -} - -export interface PostgresJsQueryContext { - arguments?: unknown[]; - self?: PostgresQuery; - result?: unknown; - error?: unknown; -} - -// postgres.js parses `host`/`port` into arrays (it can connect to multiple -// hosts); the `Connection` factory receives this parsed options object. -export interface PostgresParsedOptions { - host?: string[]; - port?: number[]; - database?: string; -} diff --git a/packages/server-utils/src/integrations/tracing-channel/postgres-js/utils.ts b/packages/server-utils/src/integrations/tracing-channel/postgres-js/utils.ts deleted file mode 100644 index ed22034344fe..000000000000 --- a/packages/server-utils/src/integrations/tracing-channel/postgres-js/utils.ts +++ /dev/null @@ -1,151 +0,0 @@ -import { ERROR_TYPE } from '@sentry/conventions/attributes'; -import type { PostgresConnectionContext, Span } from '@sentry/core'; -import { - _INTERNAL_buildPostgresConnectionContext, - _INTERNAL_setPostgresConnectionAttributes, - _INTERNAL_setPostgresOperationName, - debug, - SPAN_STATUS_ERROR, -} from '@sentry/core'; -import { DEBUG_BUILD } from '../../../debug-build'; -import type { PostgresJsQueryContext, PostgresParsedOptions, PostgresQuery } from './types'; - -// Not part of `@sentry/conventions`, so we keep it inline (matches the OTel `PostgresJsInstrumentation`). -const DB_RESPONSE_STATUS_CODE = 'db.response.status_code'; - -// Same `Symbol.for()` marker the core `instrumentPostgresJsSql` wrapper sets on -// queries it manually instruments, so we skip them there and never double-span. -export const QUERY_FROM_INSTRUMENTED_SQL = Symbol.for('sentry.query.from.instrumented.sql'); -// The query span, stashed on the `Query` so the `execute`/`connect` channels can -// attach connection attributes to it. -export const QUERY_SPAN = Symbol('sentryPostgresJsSpan'); -// Set once connection attributes are on the span, so the fallback and the -// `execute`/`connect` channels don't both write them. -const CONNECTION_ATTRS_SET = Symbol('sentryPostgresJsConnectionAttrsSet'); -// Set on the channel context once the resolve/reject wrappers have ended the -// span, so `deferSpanEnd` knows the wrappers own the lifecycle. -export const SPAN_ENDED = Symbol('sentryPostgresJsSpanEnded'); - -// A connection object -> its resolved context, populated on the `connection` -// channel and read on the `execute`/`connect` channels (keyed by the same object). -const connectionContexts = new WeakMap(); -// Distinct endpoints seen so far (value-compared, so N connections to one DB -// count once). When exactly one endpoint exists — the common case, and the only -// one the tests exercise — every query resolves to it at handle-start. -const endpointRegistry: PostgresConnectionContext[] = []; - -function registerEndpoint(context: PostgresConnectionContext): void { - const alreadyKnown = endpointRegistry.some( - e => - e.ATTR_SERVER_ADDRESS === context.ATTR_SERVER_ADDRESS && - e.ATTR_SERVER_PORT === context.ATTR_SERVER_PORT && - e.ATTR_DB_NAMESPACE === context.ATTR_DB_NAMESPACE, - ); - if (!alreadyKnown) { - endpointRegistry.push(context); - } -} - -/** The single known endpoint, or `undefined` when zero or multiple are known. */ -export function resolveSingleEndpoint(): PostgresConnectionContext | undefined { - return endpointRegistry.length === 1 ? endpointRegistry[0] : undefined; -} - -/** - * Record a connection from the `connection` channel `end` (`result` is the - * connection object, `arguments[0]` the parsed options), keying its resolved - * context by the connection object and tracking its endpoint. - */ -export function recordConnectionFromChannel(message: PostgresJsQueryContext): void { - const connection = message.result; - const options = message.arguments?.[0] as PostgresParsedOptions | undefined; - if (!connection || typeof connection !== 'object' || !options) { - return; - } - const context = _INTERNAL_buildPostgresConnectionContext(options); - connectionContexts.set(connection, context); - registerEndpoint(context); -} - -export function setConnectionAttributes(span: Span, query: PostgresQuery, context: PostgresConnectionContext): void { - const queryRecord = query as Record; - if (queryRecord[CONNECTION_ATTRS_SET]) { - return; - } - queryRecord[CONNECTION_ATTRS_SET] = true; - _INTERNAL_setPostgresConnectionAttributes(span, context); -} - -/** - * Backfill connection attributes onto a query's span from a channel whose `self` - * is the connection object and `arguments[0]` the query. Shared by the `execute` - * and `connect` channels; both carry that shape and both resolve the context via - * the `connectionContexts` WeakMap. Idempotent (guarded inside `setConnectionAttributes`). - */ -export function attachConnectionAttributesFromChannel(message: PostgresJsQueryContext): void { - const connection = message.self as object | undefined; - const query = message.arguments?.[0] as PostgresQuery | undefined; - if (!connection || !query) { - return; - } - const span = (query as Record)[QUERY_SPAN] as Span | undefined; - const context = connectionContexts.get(connection); - if (span && context) { - setConnectionAttributes(span, query, context); - } -} - -/** - * Wrap `query.resolve`/`query.reject` so the span ends when the query settles. - * - * `Query extends Promise` and `async handle()` only dispatches — its promise - * resolves immediately, long before the query completes. postgres.js signals - * completion by calling `this.resolve`/`this.reject`, so we own the span end - * there. Wrapping happens at handle-start because `reject` can fire - * synchronously during dispatch and `cursor()` reassigns both before executing. - */ -export function wrapQuerySettlement(data: PostgresJsQueryContext, span: Span, sanitizedSqlQuery: string): void { - const query = data.self; - if (!query) { - return; - } - - // Claim ownership of ending the span up front, so `deferSpanEnd` defers to the - // wrapper even if `span.end()` below throws. - const markEnded = (): void => { - (data as Record)[SPAN_ENDED] = true; - }; - - const originalResolve = query.resolve; - if (typeof originalResolve === 'function') { - query.resolve = function (this: unknown, ...resolveArgs: unknown[]): unknown { - markEnded(); - try { - const command = (resolveArgs[0] as { command?: string } | undefined)?.command; - _INTERNAL_setPostgresOperationName(span, sanitizedSqlQuery, command); - span.end(); - } catch (e) { - DEBUG_BUILD && debug.error('[orchestrion:postgresjs] error ending span in resolve:', e); - } - return originalResolve.apply(this, resolveArgs); - }; - } - - const originalReject = query.reject; - if (typeof originalReject === 'function') { - query.reject = function (this: unknown, ...rejectArgs: unknown[]): unknown { - markEnded(); - try { - const err = rejectArgs[0] as { message?: string; code?: string; name?: string } | undefined; - span.setStatus({ code: SPAN_STATUS_ERROR, message: err?.message || 'unknown_error' }); - span.setAttribute(DB_RESPONSE_STATUS_CODE, err?.code || 'unknown'); - span.setAttribute(ERROR_TYPE, err?.name || 'unknown'); - _INTERNAL_setPostgresOperationName(span, sanitizedSqlQuery); - span.end(); - } catch (e) { - DEBUG_BUILD && debug.error('[orchestrion:postgresjs] error ending span in reject:', e); - } - return originalReject.apply(this, rejectArgs); - }; - } -}