8000 Add `pg-native` support to Postgres integration. (#3894) · GinMu/sentry-javascript@69abb97 · GitHub
[go: up one dir, main page]

Skip to content

Commit 69abb97

Browse files
Add pg-native support to Postgres integration. (getsentry#3894)
Co-authored-by: Abhijeet Prasad <aprasad@sentry.io>
1 parent 486f94a commit 69abb97

File tree

2 files changed

+116
-2
lines changed

2 files changed

+116
-2
lines changed

packages/tracing/src/integrations/postgres.ts

Lines changed: 19 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,10 @@ interface PgClient {
88
};
99
}
1010

11+
interface PgOptions {
12+
usePgNative?: boolean;
13+
}
14+
1115
/** Tracing integration for node-postgres package */
1216
export class Postgres implements Integration {
1317
/**
@@ -20,25 +24,38 @@ export class Postgres implements Integration {
2024
*/
2125
public name: string = Postgres.id;
2226

27+
private _usePgNative: boolean;
28+
29+
public constructor(options: PgOptions = {}) {
30+
this._usePgNative = !!options.usePgNative;
31+
}
32+
2333
/**
2434
* @inheritDoc
2535
*/
2636
public setupOnce(_: (callback: EventProcessor) => void, getCurrentHub: () => Hub): void {
27-
const pkg = loadModule<{ Client: PgClient }>('pg');
37+
const pkg = loadModule<{ Client: PgClient; native: { Client: PgClient } }>('pg');
2838

2939
if (!pkg) {
3040
logger.error('Postgres Integration was unable to require `pg` package.');
3141
return;
3242
}
3343

44+
if (this._usePgNative && !pkg.native?.Client) {
45+
logger.error(`Postgres Integration was unable to access 'pg-native' bindings.`);
46+
return;
47+
}
48+
49+
const { Client } = this._usePgNative ? pkg.native : pkg;
50+
3451
/**
3552
* function (query, callback) => void
3653
* function (query, params, callback) => void
3754
* function (query) => Promise
3855
* function (query, params) => Promise
3956
* function (pg.Cursor) => pg.Cursor
4057
*/
41-
fill(pkg.Client.prototype, 'query', function(orig: () => void | Promise<unknown>) {
58+
fill(Client.prototype, 'query', function(orig: () => void | Promise<unknown>) {
4259
return function(this: unknown, config: unknown, values: unknown, callback: unknown) {
4360
const scope = getCurrentHub().getScope();
4461
const parentSpan = scope?.getSpan();
Lines changed: 97 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,97 @@
1+
/* eslint-disable @typescript-eslint/unbound-method */
2+
import { Hub, Scope } from '@sentry/hub';
3+
4+
import { Postgres } from '../../src/integrations/postgres';
5+
import { Span } from '../../src/span';
6+
7+
class PgClient {
8+
// https://node-postgres.com/api/client#clientquery
9+
public query(_text: unknown, values: unknown, callback?: () => void) {
10+
if (typeof callback === 'function') {
11+
callback();
12+
return;
13+
}
14+
15+
if (typeof values === 'function') {
16+
values();
17+
return;
18+
}
19+
20+
return Promise.resolve();
21+
}
22+
}
23+
24+
// mock for 'pg' / 'pg-native' package
25+
jest.mock('@sentry/utils', () => {
26+
const actual = jest.requireActual('@sentry/utils');
27+
return {
28+
...actual,
29+
loadModule() {
30+
return {
31+
Client: PgClient,
32+
native: {
33+
Client: PgClient,
34+
},
35+
};
36+
},
37+
};
38+
});
39+
40+
describe('setupOnce', () => {
41+
['pg', 'pg-native'].forEach(pgApi => {
42+
const Client: PgClient = new PgClient();
43+
let scope = new Scope();
44+
let parentSpan: Span;
45+
let childSpan: Span;
46+
47+
beforeAll(() => {
48+
(pgApi === 'pg' ? new Postgres() : new Postgres({ usePgNative: true })).setupOnce(
49+
() => undefined,
50+
() => new Hub(undefined, scope),
51+
);
52+
});
53+
54+
beforeEach(() => {
55+
scope = new Scope();
56+
parentSpan = new Span();
57+
childSpan = parentSpan.startChild();
58+
jest.spyOn(scope, 'getSpan').mockReturnValueOnce(parentSpan);
59+
jest.spyOn(parentSpan, 'startChild').mockReturnValueOnce(childSpan);
60+
jest.spyOn(childSpan, 'finish');
61+
});
62+
63+
it(`should wrap ${pgApi}'s query method accepting callback as the last argument`, done => {
64+
Client.query('SELECT NOW()', {}, function() {
65+
expect(scope.getSpan).toBeCalled();
66+
expect(parentSpan.startChild).toBeCalledWith({
67+
description: 'SELECT NOW()',
68+
op: 'db',
69+
});
70+
expect(childSpan.finish).toBeCalled();
71+
done();
72+
}) as void;
73+
});
74+
75+
it(`should wrap ${pgApi}'s query method accepting callback as the second argument`, done => {
76+
Client.query('SELECT NOW()', function() {
77+
expect(scope.getSpan).toBeCalled();
78+
expect(parentSpan.startChild).toBeCalledWith({
79+
description: 'SELECT NOW()',
80+
op: 'db',
81+
});
82+
expect(childSpan.finish).toBeCalled();
83+
done();
84+
}) as void;
85+
});
86+
87+
it(`should wrap ${pgApi}'s query method accepting no callback as the last argument but returning promise`, async () => {
88+
await Client.query('SELECT NOW()', null);
89+
expect(scope.getSpan).toBeCalled();
90+
expect(parentSpan.startChild).toBeCalledWith({
91+
description: 'SELECT NOW()',
92+
op: 'db',
93+
});
94+
expect(childSpan.finish).toBeCalled();
95+
});
96+
});
97+
});

0 commit comments

Comments
 (0)
0