Skip to content
This repository was archived by the owner on Jan 21, 2026. It is now read-only.

Commit f070636

Browse files
kjinkjin
authored
fix: add support for pg 7 changes (#702)
PR-URL: #702
1 parent c13a3bf commit f070636

7 files changed

Lines changed: 386 additions & 133 deletions

File tree

‎package-lock.json‎

Lines changed: 26 additions & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎package.json‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -64,6 +64,7 @@
6464
"@types/nock": "^9.1.2",
6565
"@types/node": "^9.4.6",
6666
"@types/once": "^1.4.0",
67+
"@types/pg": "^7.4.5",
6768
"@types/pify": "^3.0.0",
6869
"@types/proxyquire": "^1.3.28",
6970
"@types/request": "^2.0.8",

‎src/plugins/plugin-pg.ts‎

Lines changed: 175 additions & 40 deletions
Original file line numberDiff line numberDiff line change
@@ -13,61 +13,196 @@
1313
* See the License for the specific language governing permissions and
1414
* limitations under the License.
1515
*/
16-
'use strict';
1716

18-
var shimmer = require('shimmer');
17+
import {EventEmitter} from 'events';
18+
import * as shimmer from 'shimmer';
19+
import {Readable} from 'stream';
1920

20-
var SUPPORTED_VERSIONS = '^6.x || ^7.x';
21+
import {Patch, Plugin, SpanData} from '../plugin-types';
2122

22-
module.exports = [
23+
import {pg_6, pg_7} from './types';
24+
25+
// TS: Client#query also accepts a callback as a last argument, but TS cannot
26+
// detect this as it's a dependent type. So we don't specify it here.
27+
type ClientQueryArguments =
28+
[{submit?: Function} & pg_7.QueryConfig]|[string]|[string, {}];
29+
type PG7QueryReturnValue = (pg_7.QueryConfig&({submit: Function}&EventEmitter)|
30+
pg_7.Query)|Promise<pg_7.QueryResult>;
31+
32+
// tslint:disable-next-line:no-any
33+
function isSubmittable(obj: any): obj is {submit: Function} {
34+
return typeof obj.submit === 'function';
35+
}
36+
37+
const noOp = () => {};
38+
39+
function populateLabelsFromInputs(span: SpanData, args: ClientQueryArguments) {
40+
const queryObj = args[0];
41+
if (typeof queryObj === 'object') {
42+
if (queryObj.text) {
43+
span.addLabel('query', queryObj.text);
44+
}
45+
if (queryObj.values) {
46+
span.addLabel('values', queryObj.values);
47+
}
48+
} else if (typeof queryObj === 'string') {
49+
span.addLabel('query', queryObj);
50+
if (args.length >= 2 && typeof args[1] !== 'function') {
51+
span.addLabel('values', args[1]);
52+
}
53+
}
54+
}
55+
56+
function populateLabelsFromOutputs(
57+
span: SpanData, err: Error|null, res?: pg_7.QueryResult) {
58+
if (err) {
59+
span.addLabel('error', err);
60+
}
61+
if (res) {
62+
span.addLabel('row_count', res.rowCount);
63+
span.addLabel('oid', res.oid);
64+
span.addLabel('rows', res.rows);
65+
span.addLabel('fields', res.fields);
66+
}
67+
}
68+
69+
const plugin: Plugin = [
2370
{
2471
file: 'lib/client.js',
25-
versions: SUPPORTED_VERSIONS,
26-
patch: function(Client, api) {
27-
function queryWrap(query) {
28-
return function query_trace() {
29-
var span = api.createChildSpan({
30-
name: 'pg-query'
31-
});
32-
var pgQuery = query.apply(this, arguments);
72+
versions: '^6.x',
73+
// TS: Client is a class name.
74+
// tslint:disable-next-line:variable-name
75+
patch: (Client, api) => {
76+
const maybePopulateLabelsFromInputs =
77+
api.enhancedDatabaseReportingEnabled() ? populateLabelsFromInputs :
78+
noOp;
79+
const maybePopulateLabelsFromOutputs =
80+
api.enhancedDatabaseReportingEnabled() ? populateLabelsFromOutputs :
81+
noOp;
82+
shimmer.wrap(Client.prototype, 'query', (query) => {
83+
return function query_trace(this: pg_6.Client) {
84+
const span = api.createChildSpan({name: 'pg-query'});
3385
if (!api.isRealSpan(span)) {
34-
return pgQuery;
86+
return query.apply(this, arguments);
3587
}
36-
if (api.enhancedDatabaseReportingEnabled()) {
37-
span.addLabel('query', pgQuery.text);
38-
if (pgQuery.values) {
39-
span.addLabel('values', pgQuery.values);
40-
}
88+
const argLength = arguments.length;
89+
if (argLength >= 1) {
90+
const args: ClientQueryArguments =
91+
Array.prototype.slice.call(arguments, 0);
92+
// Extract query text and values, if needed.
93+
maybePopulateLabelsFromInputs(span, args);
4194
}
95+
const pgQuery: pg_6.QueryReturnValue = query.apply(this, arguments);
4296
api.wrapEmitter(pgQuery);
43-
var done = pgQuery.callback;
44-
pgQuery.callback = api.wrap(function(err, res) {
45-
if (api.enhancedDatabaseReportingEnabled()) {
46-
if (err) {
47-
span.addLabel('error', err);
48-
}
49-
if (res) {
50-
span.addLabel('row_count', res.rowCount);
51-
span.addLabel('oid', res.oid);
52-
span.addLabel('rows', res.rows);
53-
span.addLabel('fields', res.fields);
54-
}
97+
const done = pgQuery.callback;
98+
// TODO(kjin): Clean up this line a little bit by casting the function
99+
// passed to api.wrap as a NonNullable<typeof done>.
100+
pgQuery.callback =
101+
api.wrap((err: Error|null, res?: pg_7.QueryResult) => {
102+
maybePopulateLabelsFromOutputs(span, err, res);
103+
span.endSpan();
104+
if (done) {
105+
done(err, res);
106+
}
107+
});
108+
return pgQuery;
109+
};
110+
});
111+
},
112+
// TS: Client is a class name.
113+
// tslint:disable-next-line:variable-name
114+
unpatch(Client) {
115+
shimmer.unwrap(Client.prototype, 'query');
116+
}
117+
} as Patch<typeof pg_6.Client>,
118+
{
119+
file: 'lib/client.js',
120+
versions: '^7.x',
121+
// TS: Client is a class name.
122+
// tslint:disable-next-line:variable-name
123+
patch: (Client, api) => {
124+
const maybePopulateLabelsFromInputs =
125+
api.enhancedDatabaseReportingEnabled() ? populateLabelsFromInputs :
126+
noOp;
127+
const maybePopulateLabelsFromOutputs =
128+
api.enhancedDatabaseReportingEnabled() ? populateLabelsFromOutputs :
129+
noOp;
130+
shimmer.wrap(Client.prototype, 'query', (query) => {
131+
return function query_trace(this: pg_7.Client) {
132+
const span = api.createChildSpan({name: 'pg-query'});
133+
if (!api.isRealSpan(span)) {
134+
return query.apply(this, arguments);
135+
}
136+
137+
let pgQuery: PG7QueryReturnValue;
138+
// In 7.x, the value of pgQuery depends on how the query() was called.
139+
// It can be one of:
140+
// - (query: pg.Submittable) => EventEmitter
141+
// - Note: return value is the same as the argument.
142+
// - ([*], callback: (err, res: pg.Result) => void) => void
143+
// - ([*]) => Promise<pg.Result>
144+
// where [*] is one of:
145+
// - ...[query: { text: string, values?: Array<any> }]
146+
// - ...[text: string, values?: Array<any>]
147+
// See: https://node-postgres.com/guides/upgrading
148+
const argLength = arguments.length;
149+
if (argLength >= 1) {
150+
const args: ClientQueryArguments =
151+
Array.prototype.slice.call(arguments, 0);
152+
153+
// Extract query text and values, if needed.
154+
maybePopulateLabelsFromInputs(span, args);
155+
156+
// If we received a callback, bind it to the current context,
157+
// optionally adding labels as well.
158+
const callback = args[args.length - 1];
159+
if (typeof callback === 'function') {
160+
args[args.length - 1] =
161+
api.wrap((err: Error|null, res?: pg_7.QueryResult) => {
162+
maybePopulateLabelsFromOutputs(span, err, res);
163+
span.endSpan();
164+
// TS: Type cast is safe as we know that callback is a
165+
// Function.
166+
(callback as (err: Error|null, res?: pg_7.QueryResult) =>
167+
void)(err, res);
168+
});
169+
pgQuery = query.apply(this, args);
170+
} else {
171+
pgQuery = query.apply(this, arguments);
55172
}
56-
span.endSpan();
57-
if (done) {
58-
done(err, res);
173+
} else {
174+
pgQuery = query.apply(this, arguments);
175+
}
176+
177+
if (pgQuery) {
178+
if (pgQuery instanceof EventEmitter) {
179+
api.wrapEmitter(pgQuery);
180+
} else if (typeof pgQuery.then === 'function') {
181+
// Ensure that the span is ended, optionally adding labels as
182+
// well.
183+
pgQuery = pgQuery.then(
184+
(res) => {
185+
maybePopulateLabelsFromOutputs(span, null, res);
186+
span.endSpan();
187+
return res;
188+
},
189+
(err) => {
190+
maybePopulateLabelsFromOutputs(span, err);
191+
span.endSpan();
192+
throw err;
193+
});
59194
}
60-
});
195+
}
61196
return pgQuery;
62197
};
63-
}
64-
65-
shimmer.wrap(Client.prototype, 'query', queryWrap);
198+
});
66199
},
67-
unpatch: function(Client) {
200+
// TS: Client is a class name.
201+
// tslint:disable-next-line:variable-name
202+
unpatch(Client) {
68203
shimmer.unwrap(Client.prototype, 'query');
69204
}
70-
}
205+
} as Patch<typeof pg_7.Client>
71206
];
72207

73-
export default {};
208+
export = plugin;

‎src/plugins/types/index.d.ts‎

Lines changed: 27 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -2,12 +2,14 @@ import * as connect_3 from 'connect'; // connect@3
22
import * as express_4 from 'express'; // express@4
33
import * as hapi_16 from 'hapi'; // hapi@16
44
import * as koa_2 from 'koa'; // koa@2
5+
import * as pg_7 from 'pg'; // pg@7
56
import * as restify_5 from 'restify'; // restify@5
67

7-
//---koa@1---//
8-
98
import { EventEmitter } from 'events';
109
import { Server } from 'http';
10+
import { Readable } from 'stream';
11+
12+
//---koa@1---//
1113

1214
declare class koa_1 extends EventEmitter {
1315
use(middleware: koa_1.Middleware): this;
@@ -28,6 +30,27 @@ declare namespace koa_1 {
2830
interface Context extends koa_2.Context {}
2931
}
3032

33+
//---pg@6---//
34+
35+
declare namespace pg_6 {
36+
// PG 6's method signature for Client#query differs from that of PG 7 in that
37+
// the return value is either a Submittable if one was passed in, or a
38+
// pg.Query object instead. (In PG 6, pg.Query is PromiseLike and contains
39+
// values passed in as the query configuration.)
40+
//
41+
// References:
42+
// https://node-postgres.com/guides/upgrading#client-query-on
43+
// https://github.com/brianc/node-postgres/blob/v6.4.2/lib/client.js#L355
44+
type QueryReturnValue = (
45+
pg_7.QueryConfig &
46+
{ callback?: (err: Error|null, res?: pg_7.QueryResult) => void }
47+
) & (({ submit: Function } & Readable) | (pg_7.Query & PromiseLike<any>));
48+
49+
class Client {
50+
query(...args: any[]): QueryReturnValue;
51+
}
52+
}
53+
3154
//---exports---//
3255

3356
export {
@@ -36,5 +59,7 @@ export {
3659
hapi_16,
3760
koa_1,
3861
koa_2,
62+
pg_6,
63+
pg_7,
3964
restify_5
4065
};

‎test/fixtures/plugin-fixtures.json‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -152,6 +152,11 @@
152152
"pg": "^6.1.2"
153153
}
154154
},
155+
"pg7": {
156+
"dependencies": {
157+
"pg": "^7.4.1"
158+
}
159+
},
155160
"redis0.12": {
156161
"dependencies": {
157162
"redis": "^0.12.1"

0 commit comments

Comments
 (0)