Make Query & NativeQuery implement the promise interface

brianc · brianc · commit a8435f6ddbc1 · 2016-06-10T16:00:13.000-05:00
diff --git a/lib/native/query.js b/lib/native/query.js
@@ -26,6 +26,23 @@ var NativeQuery = module.exports = function(native) {
 
 util.inherits(NativeQuery, EventEmitter);
 
+NativeQuery.prototype.then = function(callback) {
+  return this.promise().then(callback);
+};
+
+NativeQuery.prototype.catch = function(callback) {
+  return this.promise().catch(callback);
+};
+
+NativeQuery.prototype.promise = function() {
+  if (this._promise) return this._promise;
+  this._promise = new Promise(function(resolve, reject) {
+    this.once('end', resolve);
+    this.once('error', reject);
+  }.bind(this));
+  return this._promise;
+};
+
 NativeQuery.prototype.handleError = function(err) {
   var self = this;
   //copy pq error fields into the error object
diff --git a/lib/query.js b/lib/query.js
@@ -26,11 +26,29 @@ var Query = function(config, values, callback) {
   this._result = new Result(config.rowMode, config.types);
   this.isPreparedStatement = false;
   this._canceledDueToError = false;
+  this._promise = null;
   EventEmitter.call(this);
 };
 
 util.inherits(Query, EventEmitter);
 
+Query.prototype.then = function(callback) {
+  return this.promise().then(callback);
+};
+
+Query.prototype.catch = function(callback) {
+  return this.promise().catch(callback);
+};
+
+Query.prototype.promise = function() {
+  if (this._promise) return this._promise;
+  this._promise = new Promise(function(resolve, reject) {
+    this.once('end', resolve);
+    this.once('error', reject);
+  }.bind(this));
+  return this._promise;
+};
+
 Query.prototype.requiresPreparation = function() {
   //named queries must always be prepared
   if(this.name) { return true; }
@@ -52,14 +70,13 @@ Query.prototype.requiresPreparation = function() {
 //metadata used when parsing row results
 Query.prototype.handleRowDescription = function(msg) {
   this._result.addFields(msg.fields);
+  this._accumulateRows = this.callback || !this.listenerCount('row');
 };
 
 Query.prototype.handleDataRow = function(msg) {
   var row = this._result.parseRow(msg.fields);
   this.emit('row', row, this._result);
-
-  //if there is a callback collect rows
-  if(this.callback) {
+  if (this._accumulateRows) {
     this._result.addRow(row);
   }
 };
diff --git a/test/integration/client/query-as-promise-tests.js b/test/integration/client/query-as-promise-tests.js
@@ -0,0 +1,28 @@
+var helper = require(__dirname + '/../test-helper');
+var pg = helper.pg;
+
+process.on('unhandledRejection', function(e) {
+  //console.error(e, e.stack)
+  //process.exit(1)
+})
+
+pg.connect(helper.config, assert.success(function(client, done) {
+  client.query('SELECT $1::text as name', ['foo'])
+    .then(function(result) {
+      assert.equal(result.rows[0].name, 'foo')
+      return client
+    })
+    .then(function(client) {
+      client.query('ALKJSDF')
+        .catch(function(e) {
+          assert(e instanceof error)
+        })
+    })
+
+  client.query('SELECT 1 as num')
+    .then(function(result) {
+      assert.equal(result.rows[0].num, 1)
+      done()
+      pg.end()
+    })
+}))
diff --git a/test/integration/client/type-coercion-tests.js b/test/integration/client/type-coercion-tests.js
@@ -202,14 +202,14 @@ helper.pg.connect(helper.config, assert.calls(function(err, client, done) {
 if(!helper.config.binary) {
   test("postgres date type", function() {
     var client = helper.client();
-    var testDate = new Date (2010, 9, 31);
+    var testDate = new Date(2010, 9, 31);
     client.on('error', function(err) {
       console.log(err);
       client.end();
     });
     client.query("SELECT $1::date", [testDate], assert.calls(function(err, result){
       assert.isNull(err);
-      assert.strictEqual(result.rows[0].date.toString(), testDate.toString());
+      assert.strictEqual(result.rows[0].date.toString(), new Date(Date.UTC(2010, 9, 31)).toString());
     }));
     client.on('drain', client.end.bind(client));
   });
