Skip to content

Commit

Permalink
feat(pg): add requireParentTrace option
Browse files Browse the repository at this point in the history
  • Loading branch information
@ryan-codaio
ryan-codaio committed Sep 16, 2022
1 parent 6584fe2 commit f0dc26d
Show file tree
Hide file tree
Showing 4 changed files with 102 additions and 45 deletions.
57 changes: 27 additions & 30 deletions plugins/node/opentelemetry-instrumentation-pg/src/instrumentation.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -24,7 +24,6 @@ import {
diag,
trace,
Span,
SpanKind,
SpanStatusCode,
} from '@opentelemetry/api';
import * as pgTypes from 'pg';
Expand All @@ -46,6 +45,7 @@ import {
DbSystemValues,
} from '@opentelemetry/semantic-conventions';
import { VERSION } from './version';
import { startSpan } from './utils';

const PG_POOL_COMPONENT = 'pg-pool';

Expand Down Expand Up @@ -129,21 +129,15 @@ export class PgInstrumentation extends InstrumentationBase {
this: pgTypes.Client,
callback?: PgErrorCallback
) {
const span = plugin.tracer.startSpan(
`${PgInstrumentation.COMPONENT}.connect`,
{
kind: SpanKind.CLIENT,
attributes: {
[SemanticAttributes.DB_SYSTEM]: DbSystemValues.POSTGRESQL,
[SemanticAttributes.DB_NAME]: this.database,
[SemanticAttributes.NET_PEER_NAME]: this.host,
[SemanticAttributes.DB_CONNECTION_STRING]:
utils.getConnectionString(this),
[SemanticAttributes.NET_PEER_PORT]: this.port,
[SemanticAttributes.DB_USER]: this.user,
},
}
);
const span = startSpan(plugin.tracer, plugin.getConfig(), `${PgInstrumentation.COMPONENT}.connect`, {
[SemanticAttributes.DB_SYSTEM]: DbSystemValues.POSTGRESQL,
[SemanticAttributes.DB_NAME]: this.database,
[SemanticAttributes.NET_PEER_NAME]: this.host,
[SemanticAttributes.DB_CONNECTION_STRING]:
utils.getConnectionString(this),
[SemanticAttributes.NET_PEER_PORT]: this.port,
[SemanticAttributes.DB_USER]: this.user,
});

if (callback) {
const parentSpan = trace.getSpan(context.active());
Expand Down Expand Up @@ -187,7 +181,12 @@ export class PgInstrumentation extends InstrumentationBase {
params
);
} else {
span = utils.handleTextQuery.call(this, plugin.tracer, query);
span = utils.handleTextQuery.call(
this,
plugin.tracer,
plugin.getConfig() as PgInstrumentationConfig,
query
);
}
} else if (typeof args[0] === 'object') {
const queryConfig = args[0] as NormalizedQueryConfig;
Expand All @@ -201,6 +200,7 @@ export class PgInstrumentation extends InstrumentationBase {
return utils.handleInvalidQuery.call(
this,
plugin.tracer,
plugin.getConfig() as PgInstrumentationConfig,
original,
...args
);
Expand Down Expand Up @@ -285,19 +285,16 @@ export class PgInstrumentation extends InstrumentationBase {
return function connect(this: PgPoolExtended, callback?: PgPoolCallback) {
const connString = utils.getConnectionString(this.options);
// setup span
const span = plugin.tracer.startSpan(`${PG_POOL_COMPONENT}.connect`, {
kind: SpanKind.CLIENT,
attributes: {
[SemanticAttributes.DB_SYSTEM]: DbSystemValues.POSTGRESQL,
[SemanticAttributes.DB_NAME]: this.options.database, // required
[SemanticAttributes.NET_PEER_NAME]: this.options.host, // required
[SemanticAttributes.DB_CONNECTION_STRING]: connString, // required
[SemanticAttributes.NET_PEER_PORT]: this.options.port,
[SemanticAttributes.DB_USER]: this.options.user,
[AttributeNames.IDLE_TIMEOUT_MILLIS]:
this.options.idleTimeoutMillis,
[AttributeNames.MAX_CLIENT]: this.options.maxClient,
},
const span = startSpan(plugin.tracer, plugin.getConfig(), `${PG_POOL_COMPONENT}.connect`, {
[SemanticAttributes.DB_SYSTEM]: DbSystemValues.POSTGRESQL,
[SemanticAttributes.DB_NAME]: this.options.database, // required
[SemanticAttributes.NET_PEER_NAME]: this.options.host, // required
[SemanticAttributes.DB_CONNECTION_STRING]: connString, // required
[SemanticAttributes.NET_PEER_PORT]: this.options.port,
[SemanticAttributes.DB_USER]: this.options.user,
[AttributeNames.IDLE_TIMEOUT_MILLIS]:
this.options.idleTimeoutMillis,
[AttributeNames.MAX_CLIENT]: this.options.maxClient,
});

if (callback) {
Expand Down
3 changes: 3 additions & 0 deletions plugins/node/opentelemetry-instrumentation-pg/src/types.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -40,6 +40,9 @@ export interface PgInstrumentationConfig extends InstrumentationConfig {
* @default undefined
*/
responseHook?: PgInstrumentationExecutionResponseHook;

/** Require that is a parent span to create new spans. Defaults to false. */
requireParentSpan?: boolean;
}

export type PostgresCallback = (err: Error, res: object) => unknown;
Expand Down
58 changes: 43 additions & 15 deletions plugins/node/opentelemetry-instrumentation-pg/src/utils.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -15,11 +15,15 @@
*/

import {
context,
trace,
Span,
SpanStatusCode,
Tracer,
SpanKind,
diag,
INVALID_SPAN_CONTEXT,
Attributes,
} from '@opentelemetry/api';
import { AttributeNames } from './enums/AttributeNames';
import {
Expand Down Expand Up @@ -58,19 +62,41 @@ export function getConnectionString(params: PgClientConnectionParams) {
return `postgresql://${host}:${port}/${database}`;
}

// Private helper function to start a span
function pgStartSpan(tracer: Tracer, client: PgClientExtended, name: string) {
const jdbcString = getConnectionString(client.connectionParameters);
export function startSpan(
tracer: Tracer,
instrumentationConfig: PgInstrumentationConfig,
name: string,
attributes: Attributes,
): Span {
// If a parent span is required but not present, use a noop span to propagate
// context without recording it. Adapted from opentelemetry-instrumentation-http:
// https://github.com/open-telemetry/opentelemetry-js/blob/597ea98e58a4f68bcd9aec5fd283852efe444cd6/experimental/packages/opentelemetry-instrumentation-http/src/http.ts#L660
const currentSpan = trace.getSpan(context.active());
if (instrumentationConfig.requireParentSpan && currentSpan === undefined) {
return trace.wrapSpanContext(INVALID_SPAN_CONTEXT);
}

return tracer.startSpan(name, {
kind: SpanKind.CLIENT,
attributes: {
[SemanticAttributes.DB_NAME]: client.connectionParameters.database, // required
[SemanticAttributes.DB_SYSTEM]: DbSystemValues.POSTGRESQL, // required
[SemanticAttributes.DB_CONNECTION_STRING]: jdbcString, // required
[SemanticAttributes.NET_PEER_NAME]: client.connectionParameters.host, // required
[SemanticAttributes.NET_PEER_PORT]: client.connectionParameters.port,
[SemanticAttributes.DB_USER]: client.connectionParameters.user,
},
attributes,
});
}

// Private helper function to start a span
function pgStartSpan(
client: PgClientExtended,
tracer: Tracer,
instrumentationConfig: PgInstrumentationConfig,
name: string,
) {
const jdbcString = getConnectionString(client.connectionParameters);
return startSpan(tracer, instrumentationConfig, name, {
[SemanticAttributes.DB_NAME]: client.connectionParameters.database, // required
[SemanticAttributes.DB_SYSTEM]: DbSystemValues.POSTGRESQL, // required
[SemanticAttributes.DB_CONNECTION_STRING]: jdbcString, // required
[SemanticAttributes.NET_PEER_NAME]: client.connectionParameters.host, // required
[SemanticAttributes.NET_PEER_PORT]: client.connectionParameters.port,
[SemanticAttributes.DB_USER]: client.connectionParameters.user,
});
}

Expand All @@ -84,7 +110,7 @@ export function handleConfigQuery(
// Set child span name
const queryCommand = getCommandFromText(queryConfig.name || queryConfig.text);
const name = PgInstrumentation.BASE_SPAN_NAME + ':' + queryCommand;
const span = pgStartSpan(tracer, this, name);
const span = pgStartSpan(this, tracer, instrumentationConfig, name);

// Set attributes
if (queryConfig.text) {
Expand Down Expand Up @@ -118,7 +144,7 @@ export function handleParameterizedQuery(
// Set child span name
const queryCommand = getCommandFromText(query);
const name = PgInstrumentation.BASE_SPAN_NAME + ':' + queryCommand;
const span = pgStartSpan(tracer, this, name);
const span = pgStartSpan(this, tracer, instrumentationConfig, name);

// Set attributes
span.setAttribute(SemanticAttributes.DB_STATEMENT, query);
Expand All @@ -133,12 +159,13 @@ export function handleParameterizedQuery(
export function handleTextQuery(
this: PgClientExtended,
tracer: Tracer,
instrumentationConfig: PgInstrumentationConfig,
query: string
) {
// Set child span name
const queryCommand = getCommandFromText(query);
const name = PgInstrumentation.BASE_SPAN_NAME + ':' + queryCommand;
const span = pgStartSpan(tracer, this, name);
const span = pgStartSpan(this, tracer, instrumentationConfig, name);

// Set attributes
span.setAttribute(SemanticAttributes.DB_STATEMENT, query);
Expand All @@ -153,11 +180,12 @@ export function handleTextQuery(
export function handleInvalidQuery(
this: PgClientExtended,
tracer: Tracer,
instrumentationConfig: PgInstrumentationConfig,
originalQuery: typeof pgTypes.Client.prototype.query,
...args: unknown[]
) {
let result;
const span = pgStartSpan(tracer, this, PgInstrumentation.BASE_SPAN_NAME);
const span = pgStartSpan(this, tracer, instrumentationConfig, PgInstrumentation.BASE_SPAN_NAME);
try {
result = originalQuery.apply(this, args as never);
} catch (e) {
Expand Down
29 changes: 29 additions & 0 deletions plugins/node/opentelemetry-instrumentation-pg/test/pg.test.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -640,4 +640,33 @@ describe('pg', () => {
});
});
});

describe('Instrumentation with requireParentSpan', () => {
beforeEach(() => {
const config: PgInstrumentationConfig = {
requireParentSpan: true,
};
instrumentation.setConfig(config);
memoryExporter.reset();
});

it('should not generate traces for connect() when requireParentSpan=true', async () => {
const connClient = new postgres.Client(CONFIG);
await connClient.connect();
const spans = memoryExporter.getFinishedSpans();
assert.strictEqual(spans.length, 0);
await connClient.end();
});

it('should not generate traces for client.query(text, callback) when requireParentSpan=true', done => {
client.query('SELECT NOW()', (err, res) => {
assert.strictEqual(err, null);
assert.ok(res);
const spans = memoryExporter.getFinishedSpans();
console.log(spans);
assert.strictEqual(spans.length, 0);
done();
});
});
});
});

0 comments on commit f0dc26d

Please sign in to comment.