Skip to content

Commit

Permalink
fix: ensure correct run-context for pg instrumentation (#2506)
Browse files Browse the repository at this point in the history 
Refs: #2430
  • Loading branch information
@trentm
trentm authored Jan 4, 2022
1 parent 291be94 commit f8f6f78
Show file tree
Hide file tree
Showing 4 changed files with 106 additions and 79 deletions.
2 changes: 2 additions & 0 deletions CHANGELOG.asciidoc
View file
Original file line number Diff line number Diff line change
Expand Up @@ -45,6 +45,8 @@ Notes:
AWS span. This change also ensures captured errors from failing client
commands are a child of the AWS span. ({issues}2430[#2430])
* Fixes for run context handling for 'pg' instrumentation. ({issues}2430[#2430])
[[release-notes-3.26.0]]
==== 3.26.0 2021/12/07
Expand Down
16 changes: 7 additions & 9 deletions examples/trace-pg.js
View file
Original file line number Diff line number Diff line change
Expand Up @@ -3,9 +3,8 @@
//
// By default this will use a Postgres on localhost with user 'postgres'.
// You can use:
// npm run docker:start
// to start a Postgres container (and other containers used for testing of
// this project).
// npm run docker:start postgres
// to start a Postgres container. Then `npm run docker:stop` to stop it.

const apm = require('../').start({ // elastic-apm-node
serviceName: 'example-trace-pg'
Expand All @@ -27,12 +26,11 @@ client.connect(function (err) {
// an HTTP server, we manually start a transaction. More details at:
// https://www.elastic.co/guide/en/apm/agent/nodejs/current/custom-transactions.html
apm.startTransaction('t1')
client.query('SELECT $1::text as message', ['Hello world!'], (err, res) => {
if (err) {
console.log('[t1] Failure: err is', err)
} else {
console.log('[t1] Success: message is %s', res.rows[0].message)
}
client.query('SELECT $1::text as message', ['hi'], (err, res) => {
console.log('[t1] err=%s res=%s', err && err.message, !err && res.rows[0].message)
})
client.query('SELECT $1::text as message', ['bye'], (err, res) => {
console.log('[t1] err=%s res=%s', err && err.message, !err && res.rows[0].message)
apm.endTransaction()
})

Expand Down
138 changes: 69 additions & 69 deletions lib/instrumentation/modules/pg.js
View file
Original file line number Diff line number Diff line change
Expand Up @@ -10,12 +10,15 @@ var symbols = require('../../symbols')
var { getDBDestination } = require('../context')

module.exports = function (pg, agent, { version, enabled }) {
if (!enabled) {
return pg
}
if (!semver.satisfies(version, '>=4.0.0 <9.0.0')) {
agent.logger.debug('pg version %s not supported - aborting...', version)
return pg
}

patchClient(pg.Client, 'pg.Client', agent, enabled)
patchClient(pg.Client, 'pg.Client', agent)

// Trying to access the pg.native getter will trigger and log the warning
// "Cannot find module 'pg-native'" to STDERR if the module isn't installed.
Expand All @@ -29,7 +32,7 @@ module.exports = function (pg, agent, { version, enabled }) {
pg.__defineGetter__('native', function () {
var native = getter()
if (native && native.Client) {
patchClient(native.Client, 'pg.native.Client', agent, enabled)
patchClient(native.Client, 'pg.native.Client', agent)
}
return native
})
Expand All @@ -38,86 +41,83 @@ module.exports = function (pg, agent, { version, enabled }) {
return pg
}

function patchClient (Client, klass, agent, enabled) {
if (!enabled) return

function patchClient (Client, klass, agent) {
agent.logger.debug('shimming %s.prototype.query', klass)
shimmer.wrap(Client.prototype, 'query', wrapQuery)

function wrapQuery (orig, name) {
return function wrappedFunction (sql) {
var span = agent.startSpan('SQL', 'db', 'postgresql', 'query')
var id = span && span.transaction.id

if (sql && typeof sql.text === 'string') sql = sql.text

agent.logger.debug('intercepted call to %s.prototype.%s %o', klass, name, { id: id, sql: sql })
agent.logger.debug('intercepted call to %s.prototype.%s', klass, name)
const ins = agent._instrumentation
const span = ins.createSpan('SQL', 'db', 'postgresql', 'query')
if (!span) {
return orig.apply(this, arguments)
}

if (span) {
// get connection parameters from Client
let host, port
if (typeof this.connectionParameters === 'object') {
({ host, port } = this.connectionParameters)
}
span.setDestinationContext(getDBDestination(span, host, port))
let sqlText = sql
if (sql && typeof sql.text === 'string') {
sqlText = sql.text
}
if (typeof sqlText === 'string') {
span.setDbContext({ statement: sqlText, type: 'sql' })
span.name = sqlSummary(sqlText)
} else {
agent.logger.debug('unable to parse sql form pg module (type: %s)', typeof sqlText)
}

var args = arguments
var index = args.length - 1
var cb = args[index]
// Get connection parameters from Client.
let host, port
if (typeof this.connectionParameters === 'object') {
({ host, port } = this.connectionParameters)
}
span.setDestinationContext(getDBDestination(span, host, port))

if (this[symbols.knexStackObj]) {
span.customStackTrace(this[symbols.knexStackObj])
this[symbols.knexStackObj] = null
}
if (this[symbols.knexStackObj]) {
span.customStackTrace(this[symbols.knexStackObj])
this[symbols.knexStackObj] = null
}

if (Array.isArray(cb)) {
index = cb.length - 1
cb = cb[index]
}
let index = arguments.length - 1
let cb = arguments[index]
if (Array.isArray(cb)) {
index = cb.length - 1
cb = cb[index]
}

if (typeof sql === 'string') {
span.setDbContext({ statement: sql, type: 'sql' })
span.name = sqlSummary(sql)
const spanRunContext = ins.currRunContext().enterSpan(span)
const onQueryEnd = ins.bindFunctionToRunContext(spanRunContext, (_err) => {
agent.logger.debug('intercepted end of %s.prototype.%s', klass, name)
span.end()
})

if (typeof cb === 'function') {
arguments[index] = ins.bindFunction((err, res) => {
onQueryEnd(err)
return cb(err, res)
})
return orig.apply(this, arguments)
} else {
var queryOrPromise = orig.apply(this, arguments)

// It is important to prefer `.on` to `.then` for pg <7 >=6.3.0, because
// `query.then` is broken in those versions. See
// https://github.com/brianc/node-postgres/commit/b5b49eb895727e01290e90d08292c0d61ab86322#r23267714
if (typeof queryOrPromise.on === 'function') {
queryOrPromise.on('end', onQueryEnd)
queryOrPromise.on('error', onQueryEnd)
if (queryOrPromise instanceof EventEmitter) {
ins.bindEmitter(queryOrPromise)
}
} else if (typeof queryOrPromise.then === 'function') {
queryOrPromise.then(
() => { onQueryEnd() },
onQueryEnd
)
} else {
agent.logger.debug('unable to parse sql form pg module (type: %s)', typeof sql)
}

const onQueryEnd = (_err) => {
agent.logger.debug('intercepted end of %s.prototype.%s %o', klass, name, { id: id })
span.end()
agent.logger.debug('ERROR: unknown pg query type: %s', typeof queryOrPromise)
}

if (typeof cb === 'function') {
args[index] = agent._instrumentation.bindFunction((err, res) => {
onQueryEnd(err)
return cb(err, res)
})
return orig.apply(this, arguments)
} else {
var queryOrPromise = orig.apply(this, arguments)

// It is import to prefer `.on` to `.then` for pg <7 >=6.3.0, because
// `query.then` is broken in those versions. See
// https://github.com/brianc/node-postgres/commit/b5b49eb895727e01290e90d08292c0d61ab86322#r23267714
if (typeof queryOrPromise.on === 'function') {
queryOrPromise.on('end', onQueryEnd)
queryOrPromise.on('error', onQueryEnd)
if (queryOrPromise instanceof EventEmitter) {
agent._instrumentation.bindEmitter(queryOrPromise)
}
} else if (typeof queryOrPromise.then === 'function') {
queryOrPromise.then(
() => { onQueryEnd() },
onQueryEnd
)
} else {
agent.logger.debug('ERROR: unknown pg query type: %s %o', typeof queryOrPromise, { id: id })
}

return queryOrPromise
}
} else {
return orig.apply(this, arguments)
return queryOrPromise
}
}
}
Expand Down
Loading

0 comments on commit f8f6f78

Please sign in to comment.