|
| 1 | +import { Hub } from '@sentry/hub'; |
| 2 | +import { EventProcessor, Integration } from '@sentry/types'; |
| 3 | +import { fill, logger } from '@sentry/utils'; |
| 4 | + |
| 5 | +interface PgClient { |
| 6 | + prototype: { |
| 7 | + query: () => void | Promise<unknown>; |
| 8 | + }; |
| 9 | +} |
| 10 | + |
| 11 | +/** Tracing integration for node-postgres package */ |
| 12 | +export class Postgres implements Integration { |
| 13 | + /** |
| 14 | + * @inheritDoc |
| 15 | + */ |
| 16 | + public static id: string = 'Postgres'; |
| 17 | + |
| 18 | + /** |
| 19 | + * @inheritDoc |
| 20 | + */ |
| 21 | + public name: string = Postgres.id; |
| 22 | + |
| 23 | + /** |
| 24 | + * @inheritDoc |
| 25 | + */ |
| 26 | + public setupOnce(_: (callback: EventProcessor) => void, getCurrentHub: () => Hub): void { |
| 27 | + let client: PgClient; |
| 28 | + |
| 29 | + try { |
| 30 | + /* eslint-disable @typescript-eslint/no-var-requires, @typescript-eslint/no-unsafe-member-access */ |
| 31 | + const pgModule = require('pg'); |
| 32 | + client = pgModule.Client; |
| 33 | + /* eslint-enable @typescript-eslint/no-var-requires, @typescript-eslint/no-unsafe-member-access */ |
| 34 | + } catch (e) { |
| 35 | + logger.error('Postgres Integration was unable to require `pg` package.'); |
| 36 | + return; |
| 37 | + } |
| 38 | + |
| 39 | + /** |
| 40 | + * function (query, callback) => void |
| 41 | + * function (query, params, callback) => void |
| 42 | + * function (query) => Promise |
| 43 | + * function (query, params) => Promise |
| 44 | + */ |
| 45 | + fill(client.prototype, 'query', function(orig: () => void | Promise<unknown>) { |
| 46 | + return function(this: unknown, config: unknown, values: unknown, callback: unknown) { |
| 47 | + const scope = getCurrentHub().getScope(); |
| 48 | + const transaction = scope?.getTransaction(); |
| 49 | + const span = transaction?.startChild({ |
| 50 | + description: typeof config === 'string' ? config : (config as { text: string }).text, |
| 51 | + op: `query`, |
| 52 | + }); |
| 53 | + |
| 54 | + if (typeof callback === 'function') { |
| 55 | + return orig.call(this, config, values, function(err: Error, result: unknown) { |
| 56 | + if (span) span.finish(); |
| 57 | + callback(err, result); |
| 58 | + }); |
| 59 | + } |
| 60 | + |
| 61 | + if (typeof values === 'function') { |
| 62 | + return orig.call(this, config, function(err: Error, result: unknown) { |
| 63 | + if (span) span.finish(); |
| 64 | + values(err, result); |
| 65 | + }); |
| 66 | + } |
| 67 | + |
| 68 | + return (orig.call(this, config, values) as Promise<unknown>).then((res: unknown) => { |
| 69 | + if (span) span.finish(); |
| 70 | + return res; |
| 71 | + }); |
| 72 | + }; |
| 73 | + }); |
| 74 | + } |
| 75 | +} |
0 commit comments