diff --git a/packages/postgres/lib/postgres_p.js b/packages/postgres/lib/postgres_p.js index 62178faf..d05111d4 100644 --- a/packages/postgres/lib/postgres_p.js +++ b/packages/postgres/lib/postgres_p.js @@ -31,36 +31,40 @@ module.exports = function capturePostgres(pg) { return pg; }; -function resolveArguments(argsObj) { - var args = {}; - - if (argsObj && argsObj.length > 0) { - if (argsObj[0] instanceof Object) { - args.sql = argsObj[0].text; - args.values = argsObj[0].values; - args.callback = typeof argsObj[1] === 'function' ? argsObj[1] : (typeof argsObj[2] === 'function' ? argsObj[2] : argsObj[0].callback); +// From pg/lib/utils.js +function pgNormalizeQueryConfig(config, values, callback) { + // can take in strings or config objects + var argsObj = typeof config === 'string' ? { text: config } : config; + if (values) { + if (typeof values === 'function') { + argsObj.callback = values; } else { - args.sql = argsObj[0]; - args.values = typeof argsObj[1] !== 'function' ? argsObj[1] : null; - args.callback = typeof argsObj[1] === 'function' ? argsObj[1] : (typeof argsObj[2] === 'function' ? argsObj[2] : undefined); + argsObj.values = values; } - - args.segment = (argsObj[argsObj.length-1] != null && argsObj[argsObj.length-1].constructor && (argsObj[argsObj.length-1].constructor.name === 'Segment' || - argsObj[argsObj.length-1].constructor.name === 'Subsegment')) ? argsObj[argsObj.length-1] : null; } - - return args; + if (callback) { + argsObj.callback = callback; + } + return argsObj; } function captureQuery() { - var args = resolveArguments(arguments); - var parent = AWSXRay.resolveSegment(args.segment); + var lastArg = arguments[arguments.length-1]; + var parent = AWSXRay.resolveSegment( + (lastArg != null && lastArg.constructor && + (lastArg.constructor.name === 'Segment' || lastArg.constructor.name === 'Subsegment')) + ? lastArg + : null + ); if (!parent) { AWSXRay.getLogger().info('Failed to capture Postgres. Cannot resolve sub/segment.'); return this.__query.apply(this, arguments); } + + var args = pgNormalizeQueryConfig.apply(this, arguments) || {}; + var subsegment = parent.addNewSubsegment(this.database + '@' + this.host); subsegment.namespace = 'remote'; @@ -86,7 +90,7 @@ function captureQuery() { } } - var result = this.__query.call(this, args.sql, args.values, args.callback); + var result = this.__query.call(this, args); if (this._queryable && !this._ending) { var query; diff --git a/packages/postgres/test/unit/postgres_p.test.js b/packages/postgres/test/unit/postgres_p.test.js index be212d6c..362e9cb9 100644 --- a/packages/postgres/test/unit/postgres_p.test.js +++ b/packages/postgres/test/unit/postgres_p.test.js @@ -56,7 +56,7 @@ describe('capturePostgres', function() { postgres.__query = function(args, values, callback) { this._queryable = true; this.queryQueue = [ null, null, queryObj ]; - queryObj.callback = callback; + queryObj.callback = callback || (typeof args === 'object' ? args.callback : undefined); return queryObj; }; @@ -95,7 +95,7 @@ describe('capturePostgres', function() { sandbox.stub(AWSXRay, 'getNamespace').returns(session); - query.call(postgres, { sql: 'sql here', callback: function() {} }); + query.call(postgres, { text: 'sql here', callback: function() {} }); assert.equal(queryObj.callback.name, 'autoContext'); queryObj.callback(); @@ -222,7 +222,7 @@ describe('capturePostgres', function() { it('should capture the error via the event', function() { var stubClose = sandbox.stub(subsegment, 'close'); - query.call(postgres, { sql: 'sql here', values: [] }).then(function() { + query.call(postgres, { text: 'sql here', values: [] }).then(function() { assert.throws(function() { queryObj.emit('error', err); }); @@ -231,4 +231,67 @@ describe('capturePostgres', function() { }); }); }); + + + + describe('#passPgParams', function() { + var postgres, query, queryObj, sandbox, segment, subsegment, stubAddNew; + + before(function() { + postgres = { Client: { prototype: { + query: function () {}, + host: 'database.location', + database: 'myTestDb', + connectionParameters: { + user: 'mcmuls', + host: 'database.location', + port: '8080', + database: 'myTestDb' + } + }}}; + postgres = capturePostgres(postgres); + + query = postgres.Client.prototype.query; + postgres = postgres.Client.prototype; + }); + + beforeEach(function() { + segment = new Segment('test'); + subsegment = segment.addNewSubsegment('testSub'); + + queryObj = {}; + + postgres.__query = function(args) { + this._queryable = true; + this.queryQueue = [ null, null, queryObj ]; + Object.assign(queryObj, args); + + return queryObj; + }; + + sandbox = sinon.createSandbox(); + sandbox.stub(AWSXRay, 'getSegment').returns(segment); + stubAddNew = sandbox.stub(segment, 'addNewSubsegment').returns(subsegment); + sandbox.stub(AWSXRay, 'isAutomaticMode').returns(true); + }); + + afterEach(function() { + sandbox.restore(); + }); + + it('should pass down raw sql query', function() { + query.call(postgres, 'sql here', ['values']); + assert.equal(queryObj.text, 'sql here'); + assert.deepEqual(queryObj.values, ['values']); + assert(stubAddNew.calledOnce); + }); + + it('should pass down parameterized query', function() { + query.call(postgres, { text: 'sql here', values: ['values'], rowMode: 'array'}); + assert.equal(queryObj.text, 'sql here'); + assert.deepEqual(queryObj.values, ['values']); + assert.equal(queryObj.rowMode, 'array'); + assert(stubAddNew.calledOnce); + }); + }); });