|
| 1 | +import * as diagnosticsChannel from '../utils/diagnosticsChannel'; |
| 2 | +import { |
| 3 | + DB_OPERATION_BATCH_SIZE, |
| 4 | + DB_QUERY_SUMMARY, |
| 5 | + DB_QUERY_TEXT, |
| 6 | + DB_SYSTEM_NAME, |
| 7 | + SENTRY_KIND, |
| 8 | + SENTRY_OP, |
| 9 | + SENTRY_ORIGIN, |
| 10 | +} from '@sentry/conventions/attributes'; |
| 11 | +import { DB } from '@sentry/conventions/op'; |
| 12 | +import type { IntegrationFn, Span } from '@sentry/core'; |
| 13 | +import { defineIntegration, getActiveSpan, getClient, hasSpanStreamingEnabled, startInactiveSpan } from '@sentry/core'; |
| 14 | +import { sanitizeSqlQueryWithSummary } from '../utils/sql'; |
| 15 | +import { CHANNELS } from '../orchestrion/channels'; |
| 16 | +import { neonModuleNames } from '../orchestrion/config/neon'; |
| 17 | +import { invokeOrchestrionInstrumentation } from '../orchestrion/instrumentation'; |
| 18 | +import { bindTracingChannelToSpan } from '../tracing-channel'; |
| 19 | +import { bindPgQueryChannel, getConnectionAttributes } from './postgres'; |
| 20 | + |
| 21 | +const INTEGRATION_NAME = 'Neon' as const; |
| 22 | +const ORIGIN = 'auto.db.neon'; |
| 23 | +const DB_SYSTEM_POSTGRESQL = 'postgresql'; |
| 24 | +const SPAN_QUERY_FALLBACK = 'neon.query'; |
| 25 | +const DEFAULT_PORT = 5432; |
| 26 | +const NOOP = (): void => {}; |
| 27 | + |
| 28 | +// `sql.query(text, params)` builds the first shape, a tagged-template call the second. |
| 29 | +type NeonQueryData = { query: string; params?: unknown[] } | NeonSqlTemplate; |
| 30 | + |
| 31 | +interface NeonSqlTemplate { |
| 32 | + strings: readonly string[]; |
| 33 | + values: unknown[]; |
| 34 | + toParameterizedQuery?: () => { query: string }; |
| 35 | +} |
| 36 | + |
| 37 | +interface NeonHttpQueryContext { |
| 38 | + // `[queryData, opts]` for a single query, `[queryData[], opts[]]` for a `transaction()` batch. |
| 39 | + arguments: unknown[]; |
| 40 | +} |
| 41 | + |
| 42 | +interface NeonResolveConnectionContext { |
| 43 | + result?: { resolvedURL?: URL }; |
| 44 | +} |
| 45 | + |
| 46 | +// The HTTP query spans this integration started, so the connection-resolver subscriber only |
| 47 | +// enriches those and never a user's own active span. |
| 48 | +const httpQuerySpans = new WeakSet<Span>(); |
| 49 | + |
| 50 | +const _neonIntegration = (() => { |
| 51 | + return { |
| 52 | + name: INTEGRATION_NAME, |
| 53 | + setup(client) { |
| 54 | + invokeOrchestrionInstrumentation(client, neonModuleNames, instrumentNeon, []); |
| 55 | + }, |
| 56 | + }; |
| 57 | +}) satisfies IntegrationFn; |
| 58 | + |
| 59 | +function instrumentNeon(): void { |
| 60 | + // The WebSocket driver is pg's `Client`, bundled into the package, so its query channel gets |
| 61 | + // the pg span shape. |
| 62 | + bindPgQueryChannel(CHANNELS.NEON_QUERY, { origin: ORIGIN }); |
| 63 | + |
| 64 | + bindTracingChannelToSpan( |
| 65 | + diagnosticsChannel.tracingChannel<NeonHttpQueryContext>(CHANNELS.NEON_HTTP_QUERY), |
| 66 | + data => { |
| 67 | + const span = startInactiveSpan(httpQuerySpanOptions(data.arguments)); |
| 68 | + httpQuerySpans.add(span); |
| 69 | + return span; |
| 70 | + }, |
| 71 | + { requiresParentSpan: true }, |
| 72 | + ); |
| 73 | + |
| 74 | + // The connection string is resolved inside the HTTP executor, so the query span is active |
| 75 | + // when this settles. The executor's own context never sees the string. |
| 76 | + diagnosticsChannel.tracingChannel<NeonResolveConnectionContext>(CHANNELS.NEON_RESOLVE_CONNECTION).subscribe({ |
| 77 | + start: NOOP, |
| 78 | + end: NOOP, |
| 79 | + asyncStart: NOOP, |
| 80 | + error: NOOP, |
| 81 | + asyncEnd(data) { |
| 82 | + const span = getActiveSpan(); |
| 83 | + const url = data.result?.resolvedURL; |
| 84 | + if (!span || !httpQuerySpans.has(span) || !url) { |
| 85 | + return; |
| 86 | + } |
| 87 | + span.setAttributes( |
| 88 | + getConnectionAttributes({ |
| 89 | + host: url.hostname, |
| 90 | + port: Number(url.port) || DEFAULT_PORT, |
| 91 | + database: url.pathname.slice(1), |
| 92 | + user: url.username || undefined, |
| 93 | + }), |
| 94 | + ); |
| 95 | + }, |
| 96 | + }); |
| 97 | +} |
| 98 | + |
| 99 | +function httpQuerySpanOptions(args: unknown[]): Parameters<typeof startInactiveSpan>[0] { |
| 100 | + const queryData = args[0]; |
| 101 | + const batchSize = Array.isArray(queryData) ? queryData.length : undefined; |
| 102 | + // A batch gets one span whose text and summary list every statement, in order. |
| 103 | + const statements = (Array.isArray(queryData) ? queryData : [queryData]) |
| 104 | + .map(getQueryText) |
| 105 | + .filter((text): text is string => !!text) |
| 106 | + .map(text => sanitizeSqlQueryWithSummary(text)); |
| 107 | + const queryText = joinDefined(statements.map(s => s.queryText)); |
| 108 | + const querySummary = joinDefined(statements.map(s => s.querySummary)); |
| 109 | + const client = getClient(); |
| 110 | + const name = |
| 111 | + client && hasSpanStreamingEnabled(client) |
| 112 | + ? querySummary || DB_SYSTEM_POSTGRESQL |
| 113 | + : (queryText ?? SPAN_QUERY_FALLBACK); |
| 114 | + |
| 115 | + return { |
| 116 | + name, |
| 117 | + attributes: { |
| 118 | + [SENTRY_OP]: DB, |
| 119 | + [SENTRY_ORIGIN]: ORIGIN, |
| 120 | + [SENTRY_KIND]: 'client', |
| 121 | + [DB_SYSTEM_NAME]: DB_SYSTEM_POSTGRESQL, |
| 122 | + [DB_QUERY_TEXT]: queryText || undefined, |
| 123 | + [DB_QUERY_SUMMARY]: querySummary, |
| 124 | + [DB_OPERATION_BATCH_SIZE]: batchSize, |
| 125 | + }, |
| 126 | + }; |
| 127 | +} |
| 128 | + |
| 129 | +function joinDefined(parts: (string | undefined)[]): string | undefined { |
| 130 | + const defined = parts.filter((part): part is string => !!part); |
| 131 | + return defined.length ? defined.join('; ') : undefined; |
| 132 | +} |
| 133 | + |
| 134 | +function getQueryText(queryData: unknown): string | undefined { |
| 135 | + const data = queryData as NeonQueryData | undefined; |
| 136 | + if (!data || typeof data !== 'object') { |
| 137 | + return undefined; |
| 138 | + } |
| 139 | + if ('query' in data && typeof data.query === 'string') { |
| 140 | + return data.query; |
| 141 | + } |
| 142 | + if ('strings' in data && Array.isArray(data.strings)) { |
| 143 | + try { |
| 144 | + // Neon's own rendering also inlines `sql.unsafe()` fragments and nested templates. |
| 145 | + if (typeof data.toParameterizedQuery === 'function') { |
| 146 | + return data.toParameterizedQuery().query; |
| 147 | + } |
| 148 | + } catch { |
| 149 | + // fall through to the plain placeholder rendering |
| 150 | + } |
| 151 | + return data.strings.reduce((text, part, i) => `${text}$${i}${part}`); |
| 152 | + } |
| 153 | + return undefined; |
| 154 | +} |
| 155 | + |
| 156 | +/** |
| 157 | + * Diagnostics-channel-based `@neondatabase/serverless` integration. |
| 158 | + * |
| 159 | + * Subscribes to the channels Sentry's code transform injects into the package's bundled pg |
| 160 | + * `Client.prototype.query` (WebSocket driver) and into the HTTP executor behind `neon()`, so both |
| 161 | + * drivers emit `db` spans with the same shape as the `pg` integration. Requires the Sentry runtime |
| 162 | + * hook or bundler plugin to be active. |
| 163 | + */ |
| 164 | +export const neonIntegration = defineIntegration(_neonIntegration); |
0 commit comments