feat: postgresql responseHook support #528
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -14,7 +14,13 @@ | |
| * limitations under the License. | ||
| */ | ||
|
|
||
| import { | ||
| Span, | ||
| SpanStatusCode, | ||
| Tracer, | ||
| SpanKind, | ||
| diag, | ||
| } from '@opentelemetry/api'; | ||
| import { AttributeNames } from './enums/AttributeNames'; | ||
| import { | ||
| SemanticAttributes, | ||
|
|
@@ -27,9 +33,11 @@ import { | |
| PgClientConnectionParams, | ||
| PgPoolCallback, | ||
| PgPoolExtended, | ||
| PgInstrumentationConfig, | ||
| } from './types'; | ||
| import * as pgTypes from 'pg'; | ||
| import { PgInstrumentation } from './'; | ||
vmarchaud marked this conversation as resolved.
Show resolved
Hide resolved
|
||
| import { safeExecuteInTheMiddle } from '@opentelemetry/instrumentation'; | ||
|
|
||
| function arrayStringifyHelper(arr: Array<unknown>): string { | ||
| return '[' + arr.toString() + ']'; | ||
|
|
@@ -161,7 +169,30 @@ export function handleInvalidQuery( | |
| return result; | ||
| } | ||
|
|
||
| export function handleExecutionResult( | ||
| config: PgInstrumentationConfig, | ||
| span: Span, | ||
| pgResult: pgTypes.QueryResult | pgTypes.QueryArrayResult | unknown | ||
| ) { | ||
| if (typeof config.responseHook === 'function') { | ||
| safeExecuteInTheMiddle( | ||
vmarchaud marked this conversation as resolved.
Show resolved
Hide resolved
|
||
| () => { | ||
| config.responseHook!(span, { | ||
| data: pgResult as pgTypes.QueryResult | pgTypes.QueryArrayResult, | ||
| }); | ||
| }, | ||
| err => { | ||
| if (err) { | ||
| diag.error('Error running response hook', err); | ||
|
There was a problem hiding this comment. optional: I like to add a prefix to my logs so if it prints, there is context on where it's coming from. Something like: There was a problem hiding this comment. Great idea, I’d suggest implementing this (in a separate PR, of course) at the InstrumentationBase level so that it won't be necessary to manually add the prefix separately for each instrumentation class. WDYT? There was a problem hiding this comment. I now read the v0.21.0 release notes and looks like it was added there: There was a problem hiding this comment. It was indeed added on instrumentation but you should use There was a problem hiding this comment. @vmarchaud, @blumamir for this PR I recommend not adding this change. There was a problem hiding this comment. Yes, I agree. We tried to refactor our instrumentations as well and bumped into the same issue, and decided not to use this component logger after all. But I see there is no convention for the contrib instrumentations on that, so it's up to you. There was a problem hiding this comment. why not passing the component logger to the util class directly from class when you call the util method ? @blumamir @nata7che in worst scenario I would change the signature to be public from protected instead of trying to mimic the behaviour of component diag logger. But it should be possible to simply pass the logger from class to util |
||
| } | ||
| }, | ||
| true | ||
| ); | ||
| } | ||
| } | ||
|
|
||
| export function patchCallback( | ||
| instrumentationConfig: PgInstrumentationConfig, | ||
| span: Span, | ||
| cb: PostgresCallback | ||
| ): PostgresCallback { | ||
|
|
@@ -176,7 +207,10 @@ export function patchCallback( | |
| code: SpanStatusCode.ERROR, | ||
| message: err.message, | ||
| }); | ||
| } else { | ||
vmarchaud marked this conversation as resolved.
Show resolved
Hide resolved
|
||
| handleExecutionResult(instrumentationConfig, span, res); | ||
|
There was a problem hiding this comment. Shouldn't we also call the response hook in the There was a problem hiding this comment. The PgPool instrumentation added an additional patching to the There was a problem hiding this comment. I have to admit I'm not very familiar with this instrumentation and its behavior. |
||
| } | ||
|
|
||
| span.end(); | ||
| cb.call(this, err, res); | ||
| }; | ||
|
|
||
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -24,7 +24,11 @@ import { | |
| trace, | ||
| } from '@opentelemetry/api'; | ||
| import { BasicTracerProvider } from '@opentelemetry/tracing'; | ||
| import { | ||
| PgInstrumentation, | ||
| PgInstrumentationConfig, | ||
| PgResponseHookInformation, | ||
| } from '../src'; | ||
| import { AsyncHooksContextManager } from '@opentelemetry/context-async-hooks'; | ||
| import * as testUtils from '@opentelemetry/test-utils'; | ||
| import { | ||
|
|
@@ -95,6 +99,11 @@ const runCallbackTest = ( | |
| }; | ||
|
|
||
| describe('[email protected]', () => { | ||
| function create(config: PgInstrumentationConfig = {}) { | ||
| instrumentation.setConfig(config); | ||
| instrumentation.enable(); | ||
| } | ||
|
|
||
| let pool: pgPool<pg.Client>; | ||
| let contextManager: AsyncHooksContextManager; | ||
| let instrumentation: PgInstrumentation; | ||
|
|
@@ -130,6 +139,7 @@ describe('[email protected]', () => { | |
| if (testPostgresLocally) { | ||
| testUtils.cleanUpDocker('postgres'); | ||
| } | ||
|
|
||
| pool.end(() => { | ||
| done(); | ||
| }); | ||
|
|
@@ -288,5 +298,154 @@ describe('[email protected]', () => { | |
| assert.strictEqual(resNoPromise, undefined, 'No promise is returned'); | ||
| }); | ||
| }); | ||
|
|
||
| describe('when specifying a responseHook configuration', () => { | ||
| const dataAttributeName = 'pg_data'; | ||
| const query = 'SELECT 0::text'; | ||
| const events: TimedEvent[] = []; | ||
|
|
||
| describe('AND valid responseHook', () => { | ||
| const pgPoolattributes = { | ||
| ...DEFAULT_PGPOOL_ATTRIBUTES, | ||
| }; | ||
| const pgAttributes = { | ||
| ...DEFAULT_PG_ATTRIBUTES, | ||
| [SemanticAttributes.DB_STATEMENT]: query, | ||
| [dataAttributeName]: '{"rowCount":1}', | ||
| }; | ||
|
|
||
| beforeEach(async () => { | ||
| const config: PgInstrumentationConfig = { | ||
| enhancedDatabaseReporting: true, | ||
| responseHook: ( | ||
| span: Span, | ||
| responseInfo: PgResponseHookInformation | ||
| ) => | ||
| span.setAttribute( | ||
| dataAttributeName, | ||
| JSON.stringify({ rowCount: responseInfo?.data.rowCount }) | ||
| ), | ||
| }; | ||
|
|
||
| create(config); | ||
| }); | ||
|
|
||
| it('should attach response hook data to resulting spans for query with callback ', done => { | ||
| const parentSpan = provider | ||
| .getTracer('test-pg-pool') | ||
| .startSpan('test span'); | ||
| context.with(trace.setSpan(context.active(), parentSpan), () => { | ||
| const resNoPromise = pool.query(query, (err, result) => { | ||
| if (err) { | ||
vmarchaud marked this conversation as resolved.
Show resolved
Hide resolved
|
||
| return done(err); | ||
| } | ||
| runCallbackTest( | ||
| parentSpan, | ||
| pgPoolattributes, | ||
| events, | ||
| unsetStatus, | ||
| 2, | ||
| 0 | ||
| ); | ||
| runCallbackTest( | ||
| parentSpan, | ||
| pgAttributes, | ||
| events, | ||
| unsetStatus, | ||
| 2, | ||
| 1 | ||
| ); | ||
| done(); | ||
| }); | ||
| assert.strictEqual( | ||
| resNoPromise, | ||
| undefined, | ||
| 'No promise is returned' | ||
| ); | ||
| }); | ||
| }); | ||
|
|
||
| it('should attach response hook data to resulting spans for query returning a Promise', async () => { | ||
| const span = provider | ||
| .getTracer('test-pg-pool') | ||
| .startSpan('test span'); | ||
| await context.with( | ||
| trace.setSpan(context.active(), span), | ||
| async () => { | ||
| const result = await pool.query(query); | ||
| runCallbackTest( | ||
| span, | ||
| pgPoolattributes, | ||
| events, | ||
| unsetStatus, | ||
| 2, | ||
| 0 | ||
| ); | ||
| runCallbackTest(span, pgAttributes, events, unsetStatus, 2, 1); | ||
| assert.ok(result, 'pool.query() returns a promise'); | ||
|
There was a problem hiding this comment.
|
||
| } | ||
| ); | ||
| }); | ||
| }); | ||
|
|
||
| describe('AND invalid responseHook', () => { | ||
| const pgPoolattributes = { | ||
| ...DEFAULT_PGPOOL_ATTRIBUTES, | ||
| }; | ||
| const pgAttributes = { | ||
| ...DEFAULT_PG_ATTRIBUTES, | ||
| [SemanticAttributes.DB_STATEMENT]: query, | ||
| }; | ||
|
|
||
| beforeEach(async () => { | ||
| create({ | ||
| enhancedDatabaseReporting: true, | ||
| responseHook: ( | ||
| span: Span, | ||
| responseInfo: PgResponseHookInformation | ||
| ) => { | ||
| throw 'some kind of failure!'; | ||
| }, | ||
| }); | ||
| }); | ||
|
|
||
| it('should not do any harm when throwing an exception', done => { | ||
| const parentSpan = provider | ||
| .getTracer('test-pg-pool') | ||
| .startSpan('test span'); | ||
| context.with(trace.setSpan(context.active(), parentSpan), () => { | ||
| const resNoPromise = pool.query(query, (err, result) => { | ||
| if (err) { | ||
| return done(err); | ||
| } | ||
| assert.ok(result); | ||
|
|
||
| runCallbackTest( | ||
| parentSpan, | ||
| pgPoolattributes, | ||
| events, | ||
| unsetStatus, | ||
| 2, | ||
| 0 | ||
| ); | ||
| runCallbackTest( | ||
| parentSpan, | ||
| pgAttributes, | ||
| events, | ||
| unsetStatus, | ||
| 2, | ||
| 1 | ||
| ); | ||
| done(); | ||
| }); | ||
| assert.strictEqual( | ||
| resNoPromise, | ||
| undefined, | ||
| 'No promise is returned' | ||
| ); | ||
| }); | ||
| }); | ||
| }); | ||
| }); | ||
| }); | ||
| }); | ||
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
I understand that nested
new Promisewas here before but it really is not needed.You can just do the sync calls inside the then callback function and it behaves exactly the same.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
thanks @rauno56, I agree and also noticed it when writing this PR. However, and if that is alright with you, I'd prefer to leave this change to a different PR?