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

Commit 1a20b7c

Browse files
kjinkjin
authored
fix: end child spans correctly in pg (#930)
1 parent bc98aa3 commit 1a20b7c

2 files changed

Lines changed: 154 additions & 83 deletions

File tree

‎src/plugins/plugin-pg.ts‎

Lines changed: 144 additions & 70 deletions
Original file line numberDiff line numberDiff line change
@@ -16,23 +16,18 @@
1616

1717
import {EventEmitter} from 'events';
1818
import * as shimmer from 'shimmer';
19-
import {Readable} from 'stream';
2019

21-
import {Patch, Plugin, Span} from '../plugin-types';
20+
import {Patch, Plugin, Span, Tracer} from '../plugin-types';
2221

2322
import {pg_6, pg_7} from './types';
2423

2524
// TS: Client#query also accepts a callback as a last argument, but TS cannot
2625
// detect this as it's a dependent type. So we don't specify it here.
2726
type ClientQueryArguments =
28-
[{submit?: Function} & pg_7.QueryConfig]|[string]|[string, {}];
27+
[Submittable & pg_7.QueryConfig]|[string]|[string, {}];
2928
type PG7QueryReturnValue = (pg_7.QueryConfig&({submit: Function}&EventEmitter)|
3029
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-
}
30+
type Callback<T> = (err: Error|null, res?: T) => void;
3631

3732
const noOp = () => {};
3833

@@ -66,46 +61,146 @@ function populateLabelsFromOutputs(
6661
}
6762
}
6863

64+
/**
65+
* Partial shape of objects returned by Client#query. Only contains methods that
66+
* are significant to the query lifecycle.
67+
*/
68+
type Submittable = {
69+
// Called when the query is completed.
70+
handleReadyForQuery: () => void;
71+
// Called when an error occurs.
72+
handleError: () => void;
73+
// A field that is populated when the Submittable is a Query object.
74+
_result?: pg_7.QueryResult;
75+
};
76+
77+
/**
78+
* Utility class to help organize patching logic.
79+
*/
80+
class PostgresPatchUtility {
81+
readonly maybePopulateLabelsFromInputs: typeof populateLabelsFromInputs;
82+
readonly maybePopulateLabelsFromOutputs: typeof populateLabelsFromOutputs;
83+
84+
constructor(private readonly tracer: Tracer) {
85+
this.maybePopulateLabelsFromInputs =
86+
tracer.enhancedDatabaseReportingEnabled() ? populateLabelsFromInputs :
87+
noOp;
88+
this.maybePopulateLabelsFromOutputs =
89+
tracer.enhancedDatabaseReportingEnabled() ? populateLabelsFromOutputs :
90+
noOp;
91+
}
92+
93+
patchSubmittable(pgQuery: Submittable, span: Span): Submittable {
94+
let spanEnded = false;
95+
const {maybePopulateLabelsFromOutputs} = this;
96+
if (pgQuery.handleError) {
97+
shimmer.wrap(pgQuery, 'handleError', (origCallback) => {
98+
// Elements of args are not individually accessed.
99+
// tslint:disable:no-any
100+
return this.tracer.wrap(function(
101+
this: Submittable, ...args: any[]): void {
102+
// tslint:enable:no-any
103+
if (!spanEnded) {
104+
const err: Error = args[0];
105+
maybePopulateLabelsFromOutputs(span, err);
106+
span.endSpan();
107+
spanEnded = true;
108+
}
109+
if (origCallback) {
110+
origCallback.apply(this, args);
111+
}
112+
});
113+
});
114+
}
115+
if (pgQuery.handleReadyForQuery) {
116+
shimmer.wrap(pgQuery, 'handleReadyForQuery', (origCallback) => {
117+
// Elements of args are not individually accessed.
118+
// tslint:disable:no-any
119+
return this.tracer.wrap(function(
120+
this: Submittable, ...args: any[]): void {
121+
// tslint:enable:no-any
122+
if (!spanEnded) {
123+
maybePopulateLabelsFromOutputs(span, null, this._result);
124+
span.endSpan();
125+
spanEnded = true;
126+
}
127+
if (origCallback) {
128+
origCallback.apply(this, args);
129+
}
130+
});
131+
});
132+
}
133+
return pgQuery;
134+
}
135+
136+
patchCallback(callback: Callback<pg_7.QueryResult>, span: Span):
137+
Callback<pg_7.QueryResult> {
138+
return this.tracer.wrap((err: Error|null, res?: pg_7.QueryResult) => {
139+
this.maybePopulateLabelsFromOutputs(span, err, res);
140+
span.endSpan();
141+
callback(err, res);
142+
});
143+
}
144+
145+
patchPromise(promise: Promise<pg_7.QueryResult>, span: Span):
146+
Promise<pg_7.QueryResult> {
147+
return promise = promise.then(
148+
(res) => {
149+
this.maybePopulateLabelsFromOutputs(span, null, res);
150+
span.endSpan();
151+
return res;
152+
},
153+
(err) => {
154+
this.maybePopulateLabelsFromOutputs(span, err);
155+
span.endSpan();
156+
throw err;
157+
});
158+
}
159+
}
160+
69161
const plugin: Plugin = [
70162
{
71163
file: 'lib/client.js',
72164
versions: '^6.x',
73165
// TS: Client is a class name.
74166
// tslint:disable-next-line:variable-name
75167
patch: (Client, api) => {
76-
const maybePopulateLabelsFromInputs =
77-
api.enhancedDatabaseReportingEnabled() ? populateLabelsFromInputs :
78-
noOp;
79-
const maybePopulateLabelsFromOutputs =
80-
api.enhancedDatabaseReportingEnabled() ? populateLabelsFromOutputs :
81-
noOp;
168+
const pgPatch = new PostgresPatchUtility(api);
169+
82170
shimmer.wrap(Client.prototype, 'query', (query) => {
83-
return function query_trace(this: pg_6.Client) {
84-
const span = api.createChildSpan({name: 'pg-query'});
85-
if (!api.isRealSpan(span)) {
86-
return query.apply(this, arguments);
87-
}
88-
const argLength = arguments.length;
89-
if (argLength >= 1) {
90-
const args: ClientQueryArguments =
91-
Array.prototype.slice.call(arguments, 0);
171+
// Every call to Client#query will have a Submittable object associated
172+
// with it. We need to patch two handlers (handleReadyForQuery and
173+
// handleError) to end a span.
174+
// There are a few things to note here:
175+
// * query accepts a Submittable or a string. A Query is a Submittable.
176+
// So if we can get a Submittable from the input we patch it
177+
// proactively, otherwise (in the case of a string) we patch the
178+
// output Query instead.
179+
// * If query is passed a callback, the callback will be invoked from
180+
// either handleReadyForQuery or handleError. So we don't need to
181+
// separately patch the callback.
182+
return function query_trace(
183+
this: pg_6.Client, ...args: ClientQueryArguments) {
184+
if (args.length >= 1) {
185+
const span = api.createChildSpan({name: 'pg-query'});
186+
if (!api.isRealSpan(span)) {
187+
return query.apply(this, args);
188+
}
92189
// Extract query text and values, if needed.
93-
maybePopulateLabelsFromInputs(span, args);
190+
pgPatch.maybePopulateLabelsFromInputs(span, args);
191+
if (typeof args[0] === 'object') {
192+
pgPatch.patchSubmittable(args[0], span);
193+
return query.apply(this, args);
194+
} else {
195+
return pgPatch.patchSubmittable(
196+
query.apply(this, args) as Submittable, span);
197+
}
198+
} else {
199+
// query was called with no arguments.
200+
// This doesn't make sense, but don't do anything that might cause
201+
// an error to get thrown here, or a span to be started.
202+
return query.apply(this, args);
94203
}
95-
const pgQuery: pg_6.QueryReturnValue = query.apply(this, arguments);
96-
api.wrapEmitter(pgQuery);
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;
109204
};
110205
});
111206
},
@@ -121,12 +216,7 @@ const plugin: Plugin = [
121216
// TS: Client is a class name.
122217
// tslint:disable-next-line:variable-name
123218
patch: (Client, api) => {
124-
const maybePopulateLabelsFromInputs =
125-
api.enhancedDatabaseReportingEnabled() ? populateLabelsFromInputs :
126-
noOp;
127-
const maybePopulateLabelsFromOutputs =
128-
api.enhancedDatabaseReportingEnabled() ? populateLabelsFromOutputs :
129-
noOp;
219+
const pgPatch = new PostgresPatchUtility(api);
130220
shimmer.wrap(Client.prototype, 'query', (query) => {
131221
return function query_trace(this: pg_7.Client) {
132222
const span = api.createChildSpan({name: 'pg-query'});
@@ -151,25 +241,18 @@ const plugin: Plugin = [
151241
Array.prototype.slice.call(arguments, 0);
152242

153243
// Extract query text and values, if needed.
154-
maybePopulateLabelsFromInputs(span, args);
244+
pgPatch.maybePopulateLabelsFromInputs(span, args);
155245

156246
// If we received a callback, bind it to the current context,
157247
// optionally adding labels as well.
158248
const callback = args[args.length - 1];
159249
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);
250+
args[args.length - 1] = pgPatch.patchCallback(
251+
callback as Callback<pg_7.QueryResult>, span);
252+
} else if (typeof args[0] === 'object') {
253+
pgPatch.patchSubmittable(args[0] as Submittable, span);
172254
}
255+
pgQuery = query.apply(this, args);
173256
} else {
174257
pgQuery = query.apply(this, arguments);
175258
}
@@ -178,19 +261,10 @@ const plugin: Plugin = [
178261
if (pgQuery instanceof EventEmitter) {
179262
api.wrapEmitter(pgQuery);
180263
} 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-
});
264+
// Unlike in pg 6, the returned value can't be both a Promise and
265+
// a Submittable. So we don't run the risk of double-patching
266+
// here.
267+
pgPatch.patchPromise(pgQuery, span);
194268
}
195269
}
196270
return pgQuery;

‎test/plugins/test-trace-pg.ts‎

Lines changed: 10 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -46,22 +46,20 @@ pgVersions.forEach(pgVersion => {
4646
client = c;
4747
releaseClient = release;
4848
assert(!err);
49-
client.query('CREATE TABLE t (name text NOT NULL, id text NOT NULL)', [],
49+
client.query('DROP TABLE t', [], function(err, res) {
50+
assert(!err || err.code == '42P01'); // table "t" does not exist
51+
client.query('CREATE TABLE t (name text NOT NULL, id text NOT NULL)', [],
5052
function(err, res) {
5153
assert(!err);
52-
common.cleanTraces();
5354
done();
5455
});
56+
});
5557
});
5658
});
5759

58-
afterEach(function(done) {
59-
client.query('DROP TABLE t', [], function(err, res) {
60-
assert(!err);
61-
releaseClient();
62-
common.cleanTraces();
63-
done();
64-
});
60+
afterEach(() => {
61+
releaseClient();
62+
common.cleanTraces();
6563
});
6664

6765
it('should perform basic operations', function(done) {
@@ -167,19 +165,18 @@ pgVersions.forEach(pgVersion => {
167165

168166
it('should work with generic Submittables', function(done) {
169167
common.runInTransaction(function(endRootSpan) {
170-
let submitCalled = false;
171168
client.query({
172169
submit: (connection) => {
173170
// Indicate that the next item may be processed.
174171
connection.emit('readyForQuery');
175-
submitCalled = true;
172+
},
173+
handleReadyForQuery: () => {
176174
endRootSpan();
177175
common.getMatchingSpan(function (span) {
178176
return span.name === 'pg-query';
179177
});
180178
done();
181-
},
182-
handleReadyForQuery: () => {}
179+
}
183180
});
184181
});
185182
});

0 commit comments

Comments
 (0)