Skip to content

Commit 8f7748e

Browse files
dszakallasgergelyke
dszakallas
authored and
gergelyke
committed
feat(pg): Postgres support
1 parent de162a3 commit 8f7748e

14 files changed

Lines changed: 517 additions & 12 deletions

File tree

‎lib/agent/index.js‎

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -85,18 +85,19 @@ Agent.prototype.getConfig = function () {
8585
Agent.prototype.serverReceive = function (data) {
8686
this.totalRequestCount++
8787
var spanId = data.spanId
88-
var originTime = data.originTime
8988
var span = this.openSpan(data)
9089
var parentId
9190

9291
if (!isNaN(data.parentId)) {
9392
parentId = parseInt(data.parentId, 10)
9493
}
9594

95+
var transportDelay = data.time - data.originTime
96+
9697
this.incomingEdgeMetrics.report({
9798
serviceKey: data.parentId,
9899
protocol: data.protocol,
99-
transportDelay: data.time - data.originTime
100+
transportDelay: transportDelay
100101
})
101102

102103
span.events.push({
@@ -107,7 +108,7 @@ Agent.prototype.serverReceive = function (data) {
107108
endpoint: data.url,
108109
method: data.method,
109110
parent: parentId,
110-
originTime: originTime
111+
originTime: data.originTime
111112
}
112113
})
113114
}

‎lib/agent/index.spec.js‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -56,7 +56,7 @@ describe('The Trace agent', function () {
5656
})
5757

5858
it('does server receive', function () {
59-
var originTime = time - 1000
59+
var originTime = String(time - 1000)
6060

6161
agent.serverReceive({
6262
id: transactionId,
@@ -71,7 +71,7 @@ describe('The Trace agent', function () {
7171
})
7272

7373
expect(incomingEdgeMetrics.report).to.have.been.calledWith({
74-
transportDelay: time - originTime,
74+
transportDelay: 1000,
7575
serviceKey: parentId,
7676
protocol: 'http'
7777
})

‎lib/agent/metrics/incomingEdge/index.spec.js‎

Lines changed: 22 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -69,6 +69,28 @@ describe('The IncomingEdgeMetrics module', function () {
6969
}])
7070
})
7171

72+
it('does not break when there\'s no edge', function () {
73+
var ISOString = 'date-string'
74+
var collectorApi = {
75+
sendIncomingEdgeMetrics: this.sandbox.spy()
76+
}
77+
78+
var edgeMetrics = EdgeMetrics.create({
79+
collectorApi: collectorApi,
80+
config: {
81+
collectInterval: 1
82+
}
83+
})
84+
85+
this.sandbox.stub(Date.prototype, 'toISOString', function () {
86+
return ISOString
87+
})
88+
89+
edgeMetrics.sendMetrics()
90+
91+
expect(collectorApi.sendIncomingEdgeMetrics).to.not.have.been.called
92+
})
93+
7294
it('works with root', function () {
7395
var ISOString = 'date-string'
7496
var collectorApi = {

‎lib/consts.js‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@ module.exports = {
33
ERROR: '1'
44
},
55
PROTOCOLS: {
6+
POSTGRES: 'pg',
67
HTTP: 'http',
78
MONGODB: 'mongodb'
89
},

‎lib/instrumentations/core/http/server.js‎

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -12,7 +12,7 @@ function wrapListener (listener, agent, mustCollectStore) {
1212
var ignoreHeaders = agent.getConfig().ignoreHeaders
1313

1414
return function (request, response) {
15-
var serverRecieveTime
15+
var serverReceiveTime
1616

1717
var headers = request.headers
1818
var spanId = headers['x-span-id']
@@ -47,14 +47,14 @@ function wrapListener (listener, agent, mustCollectStore) {
4747
}
4848

4949
var method = request.method
50-
serverRecieveTime = microtime.now()
50+
serverReceiveTime = microtime.now()
5151

5252
var collectorDataBag = {
5353
id: requestId,
5454
spanId: spanId,
5555
host: headers.host,
5656
url: util.formatDataUrl(requestUrl.pathname),
57-
time: serverRecieveTime,
57+
time: serverReceiveTime,
5858
method: method,
5959
protocol: consts.PROTOCOLS.HTTP,
6060
parentId: headers['x-parent'],
@@ -69,7 +69,7 @@ function wrapListener (listener, agent, mustCollectStore) {
6969
* @method instrumentedFinish
7070
*/
7171
function instrumentedFinish () {
72-
var responseTime = serverSendTime - serverRecieveTime
72+
var responseTime = serverSendTime - serverReceiveTime
7373

7474
var collectorDataBag = {
7575
mustCollect: mustCollectStore[requestId],
@@ -109,7 +109,7 @@ function wrapListener (listener, agent, mustCollectStore) {
109109
}
110110

111111
response.setHeader('x-server-send', serverSendTime)
112-
response.setHeader('x-server-receive', serverRecieveTime)
112+
response.setHeader('x-server-receive', serverReceiveTime)
113113

114114
if (spanId) {
115115
debug('trace event (ss); request: %s, x-span-id header has been set to: %s', requestId, spanId)

‎lib/instrumentations/index.js‎

Lines changed: 9 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -17,7 +17,8 @@ var INSTRUMENTED_LIBS = [
1717
'koa',
1818
'express',
1919
// knex and bookshelf does some black magic, so we have to do this :(
20-
'bluebird.main'
20+
'bluebird.main',
21+
'pg'
2122
]
2223

2324
var CORE_LIBS = [
@@ -31,7 +32,13 @@ function instrument (shortName, fileName, nodule, agent) {
3132
return debug('not found %s', fileName)
3233
}
3334

34-
newNodule = require(fileName)(nodule, agent)
35+
var pkg
36+
try {
37+
pkg = require(path.join(shortName, 'package.json'))
38+
} catch (err) {
39+
debug('cannot load package.json for %s: %s', shortName, err.message)
40+
}
41+
newNodule = require(fileName)(nodule, agent, pkg)
3542

3643
return newNodule
3744
}

‎lib/instrumentations/pg.js‎

Lines changed: 62 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,62 @@
1+
var semver = require('semver')
2+
var debug = require('debug')('risingstack/trace')
3+
var consts = require('../consts')
4+
var Shimmer = require('../utils/shimmer')
5+
var utils = require('./utils')
6+
7+
function tryParseSql (raw) {
8+
if (typeof raw !== 'string') {
9+
return
10+
}
11+
var matches = /\s*(select|update|insert|delete)/i.exec(raw)
12+
return matches ? matches[1] : undefined
13+
}
14+
15+
function wrapNative (native, agent, version) {
16+
// We support it only versions gte 4.0.0
17+
if (!version) {
18+
debug('trace: warning: Cannot determine postgres version. ' +
19+
'(No package.json?) Native queries won\'t be instrumented.'
20+
)
21+
} else if (!semver.satisfies(version, '>= 4.0.0')) {
22+
debug('trace: warning: You are using node-postgres version ' +
23+
version + ', (<4.0.0) so native queries won\'t be instrumented.'
24+
)
25+
} else {
26+
wrapClient(native, agent)
27+
}
28+
}
29+
30+
function wrapClient (pg, agent) {
31+
Shimmer.wrap(pg.Client.prototype, 'pg.Client.prototype', 'query', function (original) {
32+
return function () {
33+
var host = this.host
34+
var url = this.database
35+
var args = Array.prototype.slice.apply(arguments)
36+
37+
return utils.wrapQuery.call(this, original, args, agent, {
38+
protocol: consts.PROTOCOLS.POSTGRES,
39+
url: url,
40+
host: host,
41+
method: tryParseSql(args[0])
42+
})
43+
}
44+
})
45+
}
46+
47+
module.exports = function wrap (pg, agent, pkg) {
48+
var version = pkg ? pkg.version : undefined
49+
var original = pg.__lookupGetter__('native')
50+
if (original) {
51+
delete pg.native
52+
pg.__defineGetter__('native', function getNative () {
53+
var tmp = original()
54+
wrapNative(tmp, agent, version)
55+
return tmp
56+
})
57+
}
58+
59+
wrapClient(pg, agent)
60+
61+
return pg
62+
}

‎lib/instrumentations/pg.spec.e2e.js‎

Lines changed: 152 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,152 @@
1+
'use strict'
2+
3+
var expect = require('chai').expect
4+
var wrap = require('./pg')
5+
var Shimmer = require('../utils/shimmer')
6+
7+
function fakeAgent (sandbox) {
8+
return {
9+
generateSpanId: function () { return 'fakeSpanId' },
10+
getMicrotime: function () { return 42 },
11+
getTransactionId: function () { return 'fakeTransactionId' },
12+
clientSend: sandbox.spy(),
13+
clientReceive: sandbox.spy(),
14+
CLIENT_SEND: 'fakeSend'
15+
}
16+
}
17+
18+
describe('pg module wrapper', function () {
19+
beforeEach(function () {
20+
Shimmer.unwrapAll()
21+
})
22+
23+
it('should instrument JS query', function (done) {
24+
var agent = fakeAgent(this.sandbox)
25+
var pg = wrap(require('pg'), agent)
26+
27+
var conString = 'postgres://localhost/postgres'
28+
var qryString = 'SELECT 1 AS "one"'
29+
30+
var client = new pg.Client(conString)
31+
client.connect(function (err) {
32+
if (err) {
33+
console.error(err)
34+
throw err
35+
}
36+
client.query(qryString, function (err, result) {
37+
if (err) {
38+
console.error(err)
39+
throw err
40+
}
41+
expect(result.rows[0].one).to.eql(1)
42+
expect(agent.clientReceive).to.have.been.calledWith({
43+
host: 'localhost',
44+
id: 'fakeTransactionId',
45+
method: 'SELECT',
46+
mustCollect: undefined,
47+
protocol: 'pg',
48+
responseTime: 0,
49+
spanId: 'fakeSpanId',
50+
status: 0,
51+
statusCode: 200,
52+
time: 42,
53+
url: 'postgres'
54+
})
55+
client.end()
56+
done()
57+
})
58+
59+
expect(agent.clientSend).to.have.been.called
60+
expect(agent.clientSend).to.have.been.calledWith({
61+
id: 'fakeTransactionId',
62+
spanId: 'fakeSpanId',
63+
host: 'localhost',
64+
time: 42,
65+
method: 'SELECT',
66+
type: 'fakeSend',
67+
url: 'postgres'
68+
})
69+
})
70+
})
71+
72+
it('should instrument JS query w/o callback', function (done) {
73+
var agent = fakeAgent(this.sandbox)
74+
agent.clientReceive = this.sandbox.spy(function (options) {
75+
done()
76+
})
77+
var pg = wrap(require('pg'), agent)
78+
79+
var conString = 'postgres://localhost/postgres'
80+
var qryString = 'SELECT 1 AS "one"'
81+
82+
var client = new pg.Client(conString)
83+
client.connect(function (err) {
84+
if (err) {
85+
console.error(err)
86+
throw err
87+
}
88+
client.query(qryString)
89+
90+
expect(agent.clientSend).to.have.been.called
91+
expect(agent.clientSend).to.have.been.calledWith({
92+
id: 'fakeTransactionId',
93+
spanId: 'fakeSpanId',
94+
host: 'localhost',
95+
time: 42,
96+
method: 'SELECT',
97+
type: 'fakeSend',
98+
url: 'postgres'
99+
})
100+
})
101+
})
102+
103+
it('should instrument native query', function (done) {
104+
var agent = fakeAgent(this.sandbox)
105+
var pkg = require('pg/package.json')
106+
var pg = wrap(require('pg'), agent, pkg).native
107+
108+
var conString = 'postgres://localhost/postgres'
109+
var qryString = 'SELECT 1 AS "one"'
110+
111+
var client = new pg.Client(conString)
112+
client.connect(function (err) {
113+
if (err) {
114+
console.error(err)
115+
throw err
116+
}
117+
client.query(qryString, function (err, result) {
118+
if (err) {
119+
console.error(err)
120+
throw err
121+
}
122+
expect(result.rows[0].one).to.eql(1)
123+
expect(agent.clientReceive).to.have.been.calledWith({
124+
host: 'localhost',
125+
id: 'fakeTransactionId',
126+
method: 'SELECT',
127+
mustCollect: undefined,
128+
protocol: 'pg',
129+
responseTime: 0,
130+
spanId: 'fakeSpanId',
131+
status: 0,
132+
statusCode: 200,
133+
time: 42,
134+
url: 'postgres'
135+
})
136+
client.end()
137+
done()
138+
})
139+
140+
expect(agent.clientSend).to.have.been.called
141+
expect(agent.clientSend).to.have.been.calledWith({
142+
id: 'fakeTransactionId',
143+
spanId: 'fakeSpanId',
144+
host: 'localhost',
145+
time: 42,
146+
method: 'SELECT',
147+
type: 'fakeSend',
148+
url: 'postgres'
149+
})
150+
})
151+
})
152+
})

0 commit comments

Comments
 (0)