Skip to content

Commit

Permalink
chore: address comments
Browse files Browse the repository at this point in the history
  • Loading branch information
@kjin
kjin committed Mar 29, 2018
1 parent db6338c commit c070d4a
Show file tree
Hide file tree
Showing 2 changed files with 76 additions and 79 deletions.
153 changes: 75 additions & 78 deletions src/plugins/plugin-pg.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -18,13 +18,13 @@ import {EventEmitter} from 'events';
import * as shimmer from 'shimmer';
import {Readable} from 'stream';

import {Patch, Plugin} from '../plugin-types';
import {Patch, Plugin, SpanData} from '../plugin-types';

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

// TS: Client#query also accepts a callback as a last argument, but TS cannot
// detect this as it's a dependent type. So we don't specify it here.
type PG7QueryArguments =
type ClientQueryArguments =
[{submit?: Function} & pg_7.QueryConfig]|[string]|[string, {}];
type PG7QueryReturnValue = (pg_7.QueryConfig&({submit: Function}&EventEmitter)|
pg_7.Query)|Promise<pg_7.QueryResult>;
Expand All @@ -34,60 +34,77 @@ function isSubmittable(obj: any): obj is {submit: Function} {
return typeof obj.submit === 'function';
}

const noOp = () => {};

function populateLabelsFromInputs(span: SpanData, args: ClientQueryArguments) {
const queryObj = args[0];
if (typeof queryObj === 'object') {
if (queryObj.text) {
span.addLabel('query', queryObj.text);
}
if (queryObj.values) {
span.addLabel('values', queryObj.values);
}
} else if (typeof queryObj === 'string') {
span.addLabel('query', queryObj);
if (args.length >= 2 && typeof args[1] !== 'function') {
span.addLabel('values', args[1]);
}
}
}

function populateLabelsFromOutputs(
span: SpanData, err: Error|null, res?: pg_7.QueryResult) {
if (err) {
span.addLabel('error', err);
}
if (res) {
span.addLabel('row_count', res.rowCount);
span.addLabel('oid', res.oid);
span.addLabel('rows', res.rows);
span.addLabel('fields', res.fields);
}
}

const plugin: Plugin = [
{
file: 'lib/client.js',
versions: '^6.x',
// TS: Client is a class name.
// tslint:disable-next-line:variable-name
patch: (Client, api) => {
const maybePopulateLabelsFromInputs =
api.enhancedDatabaseReportingEnabled() ? populateLabelsFromInputs :
noOp;
const maybePopulateLabelsFromOutputs =
api.enhancedDatabaseReportingEnabled() ? populateLabelsFromOutputs :
noOp;
shimmer.wrap(Client.prototype, 'query', (query) => {
return function query_trace(this: pg_6.Client) {
const span = api.createChildSpan({name: 'pg-query'});
if (!api.isRealSpan(span)) {
return query.apply(this, arguments);
}
const args: pg_6.QueryReturnValue[] =
Array.prototype.slice.call(arguments, 0);
if (args.length >= 1) {
const argLength = arguments.length;
if (argLength >= 1) {
const args: ClientQueryArguments =
Array.prototype.slice.call(arguments, 0);
// Extract query text and values, if needed.
if (api.enhancedDatabaseReportingEnabled()) {
const queryObj = args[0];
if (typeof queryObj === 'object') {
if (queryObj.text) {
span.addLabel('query', queryObj.text);
}
if (queryObj.values) {
span.addLabel('values', queryObj.values);
}
} else if (typeof queryObj === 'string') {
span.addLabel('query', queryObj);
if (args.length >= 2 && typeof args[1] !== 'function') {
span.addLabel('values', args[1]);
}
}
}
maybePopulateLabelsFromInputs(span, args);
}
const pgQuery: pg_6.QueryReturnValue = query.apply(this, arguments);
api.wrapEmitter(pgQuery);
const done = pgQuery.callback;
pgQuery.callback = api.wrap((err, res) => {
if (api.enhancedDatabaseReportingEnabled()) {
if (err) {
span.addLabel('error', err);
}
if (res) {
span.addLabel('row_count', res.rowCount);
span.addLabel('oid', res.oid);
span.addLabel('rows', res.rows);
span.addLabel('fields', res.fields);
}
}
span.endSpan();
if (done) {
done(err, res);
}
});
// TODO(kjin): Clean up this line a little bit by casting the function
// passed to api.wrap as a NonNullable<typeof done>.
pgQuery.callback =
api.wrap((err: Error|null, res?: pg_7.QueryResult) => {
maybePopulateLabelsFromOutputs(span, err, res);
span.endSpan();
if (done) {
done(err, res);
}
});
return pgQuery;
};
});
Expand All @@ -104,15 +121,20 @@ const plugin: Plugin = [
// TS: Client is a class name.
// tslint:disable-next-line:variable-name
patch: (Client, api) => {
const maybePopulateLabelsFromInputs =
api.enhancedDatabaseReportingEnabled() ? populateLabelsFromInputs :
noOp;
const maybePopulateLabelsFromOutputs =
api.enhancedDatabaseReportingEnabled() ? populateLabelsFromOutputs :
noOp;
shimmer.wrap(Client.prototype, 'query', (query) => {
return function query_trace(this: pg_7.Client) {
const span = api.createChildSpan({name: 'pg-query'});
if (!api.isRealSpan(span)) {
return query.apply(this, arguments);
}
const args: PG7QueryArguments =
Array.prototype.slice.call(arguments, 0);

let pgQuery: PG7QueryReturnValue;
// In 7.x, the value of pgQuery depends on how the query() was called.
// It can be one of:
// - (query: pg.Submittable) => EventEmitter
Expand All @@ -123,51 +145,33 @@ const plugin: Plugin = [
// - ...[query: { text: string, values?: Array<any> }]
// - ...[text: string, values?: Array<any>]
// See: https://node-postgres.com/guides/upgrading
const argLength = arguments.length;
if (argLength >= 1) {
const args: ClientQueryArguments =
Array.prototype.slice.call(arguments, 0);

if (args.length >= 1) {
// Extract query text and values, if needed.
if (api.enhancedDatabaseReportingEnabled()) {
const queryObj = args[0];
if (typeof queryObj === 'object') {
if (queryObj.text) {
span.addLabel('query', queryObj.text);
}
if (queryObj.values) {
span.addLabel('values', queryObj.values);
}
} else if (typeof queryObj === 'string') {
span.addLabel('query', queryObj);
if (args.length >= 2 && typeof args[1] !== 'function') {
span.addLabel('values', args[1]);
}
}
}
maybePopulateLabelsFromInputs(span, args);

// If we received a callback, bind it to the current context,
// optionally adding labels as well.
const callback = args[args.length - 1];
if (typeof callback === 'function') {
args[args.length - 1] = api.wrap((err, res) => {
if (api.enhancedDatabaseReportingEnabled()) {
if (err) {
span.addLabel('error', err);
}
if (res) {
span.addLabel('row_count', res.rowCount);
span.addLabel('oid', res.oid);
span.addLabel('rows', res.rows);
span.addLabel('fields', res.fields);
}
}
maybePopulateLabelsFromOutputs(span, err, res);
span.endSpan();
// TS: Type cast is safe as we know that callback is a Function.
(callback as (err: Error, res: pg_7.QueryArrayResult) => void)(
err, res);
});
pgQuery = query.apply(this, args);
} else {
pgQuery = query.apply(this, arguments);
}
} else {
pgQuery = query.apply(this, arguments);
}

let pgQuery: PG7QueryReturnValue = query.apply(this, args);
if (pgQuery) {
if (pgQuery instanceof EventEmitter) {
api.wrapEmitter(pgQuery);
Expand All @@ -176,19 +180,12 @@ const plugin: Plugin = [
// well.
pgQuery = pgQuery.then(
(res) => {
if (api.enhancedDatabaseReportingEnabled()) {
span.addLabel('row_count', res.rowCount);
span.addLabel('oid', res.oid);
span.addLabel('rows', res.rows);
span.addLabel('fields', res.fields);
}
maybePopulateLabelsFromOutputs(span, null, res);
span.endSpan();
return res;
},
(err) => {
if (api.enhancedDatabaseReportingEnabled()) {
span.addLabel('error', err);
}
maybePopulateLabelsFromOutputs(span, err);
span.endSpan();
throw err;
});
Expand Down
2 changes: 1 addition & 1 deletion src/plugins/types/index.d.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -43,7 +43,7 @@ declare namespace pg_6 {
// https://github.com/brianc/node-postgres/blob/v6.4.2/lib/client.js#L355
type QueryReturnValue = (
pg_7.QueryConfig &
{ callback?: (err: Error, res: pg_7.QueryResult) => void }
{ callback?: (err: Error|null, res?: pg_7.QueryResult) => void }
) & (({ submit: Function } & Readable) | (pg_7.Query & PromiseLike<any>));

class Client {
Expand Down

0 comments on commit c070d4a

Please sign in to comment.