Fix rowMode='array' in postgres patcher

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;

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,65 @@ describe('capturePostgres', function() {
       });
     });
   });
+
+
+
+  describe('#passPgParams', function() {
+    var postgres, query, queryObj, sandbox, segment, subsegment;
+
+    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);
+      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']);
+    });
+
+    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');
+    });
+  });
 });