Skip to content

Commit 1a2d926

Browse files
xiao-lixmayurkale22
xiao-lix
and
mayurkale22
committed
feat(plugin): pg-pool plugin implementation (#501)
* feat: pg-pool plugin implementation * feat: pg-pool plugin implementation * fix: linting * fix: add attributes for span & add tests for pool.query() * fix: add span.setStatus * chore: address comments * fix: linting Co-authored-by: Mayur Kale <[email protected]>
1 parent 38ce1c3 commit 1a2d926

File tree

8 files changed

+613
-0
lines changed

8 files changed

+613
-0
lines changed

‎packages/opentelemetry-plugin-postgres/opentelemetry-plugin-pg-pool/src/enums.ts

+36
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,36 @@
1+
/*!
2+
* Copyright 2019, OpenTelemetry Authors
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* https://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
17+
export enum AttributeNames {
18+
// required by https://github.com/open-telemetry/opentelemetry-specification/blob/master/specification/data-semantic-conventions.md#databases-client-calls
19+
COMPONENT = 'component',
20+
DB_TYPE = 'db.type',
21+
DB_INSTANCE = 'db.instance',
22+
DB_STATEMENT = 'db.statement',
23+
PEER_ADDRESS = 'peer.address',
24+
PEER_HOSTNAME = 'peer.host',
25+
26+
// optional
27+
DB_USER = 'db.user',
28+
PEER_PORT = 'peer.port',
29+
PEER_IPV4 = 'peer.ipv4',
30+
PEER_IPV6 = 'peer.ipv6',
31+
PEER_SERVICE = 'peer.service',
32+
33+
// PG-POOL specific -- not specified by spec
34+
IDLE_TIMEOUT_MILLIS = 'idle.timeout.millis',
35+
MAX_CLIENT = 'max',
36+
}

‎packages/opentelemetry-plugin-postgres/opentelemetry-plugin-pg-pool/src/index.ts

+17
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,17 @@
1+
/*!
2+
* Copyright 2019, OpenTelemetry Authors
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* https://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
17+
export * from './pg-pool';

‎packages/opentelemetry-plugin-postgres/opentelemetry-plugin-pg-pool/src/pg-pool.ts

+117
Original file line numberDiff line numberDiff line change
@@ -13,3 +13,120 @@
1313
* See the License for the specific language governing permissions and
1414
* limitations under the License.
1515
*/
16+
17+
import { BasePlugin } from '@opentelemetry/core';
18+
import { CanonicalCode, SpanKind } from '@opentelemetry/types';
19+
import { AttributeNames } from './enums';
20+
import * as shimmer from 'shimmer';
21+
import * as pgPoolTypes from 'pg-pool';
22+
import {
23+
PostgresPoolPluginOptions,
24+
PgPoolCallback,
25+
PgPoolExtended,
26+
} from './types';
27+
import * as utils from './utils';
28+
29+
export class PostgresPoolPlugin extends BasePlugin<typeof pgPoolTypes> {
30+
protected _config: PostgresPoolPluginOptions;
31+
32+
static readonly COMPONENT = 'pg-pool';
33+
static readonly DB_TYPE = 'sql';
34+
35+
readonly supportedVersions = ['2.*'];
36+
37+
constructor(readonly moduleName: string) {
38+
super();
39+
this._config = {};
40+
}
41+
42+
protected patch(): typeof pgPoolTypes {
43+
shimmer.wrap(
44+
this._moduleExports.prototype,
45+
'connect',
46+
this._getPoolConnectPatch() as never
47+
);
48+
49+
return this._moduleExports;
50+
}
51+
52+
protected unpatch(): void {
53+
shimmer.unwrap(this._moduleExports.prototype, 'connect');
54+
}
55+
56+
private _getPoolConnectPatch() {
57+
const plugin = this;
58+
return (originalConnect: typeof pgPoolTypes.prototype.connect) => {
59+
plugin._logger.debug(
60+
`Patching ${PostgresPoolPlugin.COMPONENT}.prototype.connect`
61+
);
62+
return function connect(this: PgPoolExtended, callback?: PgPoolCallback) {
63+
const jdbcString = utils.getJDBCString(this.options);
64+
// setup span
65+
const span = plugin._tracer.startSpan(
66+
`${PostgresPoolPlugin.COMPONENT}.connect`,
67+
{
68+
kind: SpanKind.CLIENT,
69+
parent: plugin._tracer.getCurrentSpan() || undefined,
70+
attributes: {
71+
[AttributeNames.COMPONENT]: PostgresPoolPlugin.COMPONENT, // required
72+
[AttributeNames.DB_TYPE]: PostgresPoolPlugin.DB_TYPE, // required
73+
[AttributeNames.DB_INSTANCE]: this.options.database, // required
74+
[AttributeNames.PEER_HOSTNAME]: this.options.host, // required
75+
[AttributeNames.PEER_ADDRESS]: jdbcString, // required
76+
[AttributeNames.PEER_PORT]: this.options.port,
77+
[AttributeNames.DB_USER]: this.options.user,
78+
[AttributeNames.IDLE_TIMEOUT_MILLIS]: this.options
79+
.idleTimeoutMillis,
80+
[AttributeNames.MAX_CLIENT]: this.options.maxClient,
81+
},
82+
}
83+
);
84+
85+
if (callback) {
86+
const parentSpan = plugin._tracer.getCurrentSpan();
87+
callback = utils.patchCallback(span, callback) as PgPoolCallback;
88+
// If a parent span exists, bind the callback
89+
if (parentSpan) {
90+
callback = plugin._tracer.bind(callback);
91+
}
92+
}
93+
94+
const connectResult: unknown = originalConnect.call(
95+
this,
96+
callback as never
97+
);
98+
99+
// No callback was provided, return a promise instead
100+
if (connectResult instanceof Promise) {
101+
const connectResultPromise = connectResult as Promise<unknown>;
102+
return plugin._tracer.bind(
103+
connectResultPromise
104+
.then((result: any) => {
105+
// Resturn a pass-along promise which ends the span and then goes to user's orig resolvers
106+
return new Promise((resolve, _) => {
107+
span.setStatus({ code: CanonicalCode.OK });
108+
span.end();
109+
resolve(result);
110+
});
111+
})
112+
.catch((error: Error) => {
113+
return new Promise((_, reject) => {
114+
span.setStatus({
115+
code: CanonicalCode.UNKNOWN,
116+
message: error.message,
117+
});
118+
span.end();
119+
reject(error);
120+
});
121+
})
122+
);
123+
}
124+
125+
// Else a callback was provided, so just return the result
126+
return connectResult;
127+
};
128+
};
129+
}
130+
}
131+
132+
export const plugin = new PostgresPoolPlugin(PostgresPoolPlugin.COMPONENT);

‎packages/opentelemetry-plugin-postgres/opentelemetry-plugin-pg-pool/src/types.ts

+39
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,39 @@
1+
/*!
2+
* Copyright 2019, OpenTelemetry Authors
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* https://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
17+
import * as pgTypes from 'pg';
18+
import * as pgPoolTypes from 'pg-pool';
19+
20+
export interface PostgresPoolPluginOptions {}
21+
22+
export type PgPoolCallback = (
23+
err: Error,
24+
client: any,
25+
done: (release?: any) => void
26+
) => void;
27+
28+
export interface PgPoolOptionsParams {
29+
database: string;
30+
host: string;
31+
port: number;
32+
user: string;
33+
idleTimeoutMillis: number; // the minimum amount of time that an object may sit idle in the pool before it is eligible for eviction due to idle time
34+
maxClient: number; // maximum size of the pool
35+
}
36+
37+
export interface PgPoolExtended extends pgPoolTypes<pgTypes.Client> {
38+
options: PgPoolOptionsParams;
39+
}

‎packages/opentelemetry-plugin-postgres/opentelemetry-plugin-pg-pool/src/utils.ts

+45
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,45 @@
1+
/*!
2+
* Copyright 2019, OpenTelemetry Authors
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* https://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
17+
import { Span, CanonicalCode } from '@opentelemetry/types';
18+
import { PgPoolOptionsParams, PgPoolCallback, PgPoolExtended } from './types';
19+
20+
export function getJDBCString(params: PgPoolOptionsParams) {
21+
const host = params.host || 'localhost'; // postgres defaults to localhost
22+
const port = params.port || 5432; // postgres defaults to port 5432
23+
const database = params.database || '';
24+
return `jdbc:postgresql://${host}:${port}/${database}`;
25+
}
26+
27+
export function patchCallback(span: Span, cb: PgPoolCallback): PgPoolCallback {
28+
return function patchedCallback(
29+
this: PgPoolExtended,
30+
err: Error,
31+
res: object,
32+
done: any
33+
) {
34+
if (err) {
35+
span.setStatus({
36+
code: CanonicalCode.UNKNOWN,
37+
message: err.message,
38+
});
39+
} else if (res) {
40+
span.setStatus({ code: CanonicalCode.OK });
41+
}
42+
span.end();
43+
cb.call(this, err, res, done);
44+
};
45+
}

‎packages/opentelemetry-plugin-postgres/opentelemetry-plugin-pg-pool/test/assertionUtils.ts

+79
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,79 @@
1+
/*!
2+
* Copyright 2019, OpenTelemetry Authors
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* https://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
17+
import {
18+
SpanKind,
19+
Attributes,
20+
Event,
21+
Span,
22+
TimedEvent,
23+
} from '@opentelemetry/types';
24+
import * as assert from 'assert';
25+
import { ReadableSpan } from '@opentelemetry/tracing';
26+
import {
27+
hrTimeToMilliseconds,
28+
hrTimeToMicroseconds,
29+
} from '@opentelemetry/core';
30+
31+
export const assertSpan = (
32+
span: ReadableSpan,
33+
kind: SpanKind,
34+
attributes: Attributes,
35+
events: Event[]
36+
) => {
37+
assert.strictEqual(span.spanContext.traceId.length, 32);
38+
assert.strictEqual(span.spanContext.spanId.length, 16);
39+
assert.strictEqual(span.kind, kind);
40+
41+
// check all the AttributeNames fields
42+
Object.keys(span.attributes).forEach(key => {
43+
assert.deepStrictEqual(span.attributes[key], attributes[key]);
44+
});
45+
46+
assert.ok(span.endTime);
47+
assert.strictEqual(span.links.length, 0);
48+
49+
assert.ok(
50+
hrTimeToMicroseconds(span.startTime) < hrTimeToMicroseconds(span.endTime)
51+
);
52+
assert.ok(hrTimeToMilliseconds(span.endTime) > 0);
53+
54+
// events
55+
assert.strictEqual(
56+
span.events.length,
57+
events.length,
58+
'Should contain same number of events'
59+
);
60+
span.events.forEach((_: TimedEvent, index: number) => {
61+
assert.deepStrictEqual(span.events[index], events[index]);
62+
});
63+
};
64+
65+
// Check if sourceSpan was propagated to targetSpan
66+
export const assertPropagation = (
67+
childSpan: ReadableSpan,
68+
parentSpan: Span
69+
) => {
70+
const targetSpanContext = childSpan.spanContext;
71+
const sourceSpanContext = parentSpan.context();
72+
assert.strictEqual(targetSpanContext.traceId, sourceSpanContext.traceId);
73+
assert.strictEqual(childSpan.parentSpanId, sourceSpanContext.spanId);
74+
assert.strictEqual(
75+
targetSpanContext.traceFlags,
76+
sourceSpanContext.traceFlags
77+
);
78+
assert.notStrictEqual(targetSpanContext.spanId, sourceSpanContext.spanId);
79+
};

0 commit comments

Comments
 (0)