From d3be101fb042639ab2fdee89df3f79b98c7492da Mon Sep 17 00:00:00 2001
From: Charmander <~@charmander.me>
Date: Tue, 21 Nov 2017 14:01:08 -0800
Subject: [PATCH] Allow a custom type to be used for Client promises

Matches the Pool API.
---
 lib/client.js                                 | 12 ++++------
 lib/native/client.js                          | 12 ++++------
 package.json                                  |  1 +
 .../client/query-as-promise-tests.js          | 24 ++++++++++++++++++-
 4 files changed, 34 insertions(+), 15 deletions(-)

diff --git a/lib/client.js b/lib/client.js
index 432443793..625aad682 100644
--- a/lib/client.js
+++ b/lib/client.js
@@ -31,6 +31,7 @@ var Client = function (config) {
 
   var c = config || {}
 
+  this._Promise = c.Promise || global.Promise
   this._types = new TypeOverrides(c.types)
   this._ending = false
   this._connecting = false
@@ -232,7 +233,7 @@ Client.prototype.connect = function (callback) {
     return
   }
 
-  return new Promise((resolve, reject) => {
+  return new this._Promise((resolve, reject) => {
     this._connect((error) => {
       if (error) {
         reject(error)
@@ -409,12 +410,9 @@ Client.prototype.query = function (config, values, callback) {
   } else {
     query = new Query(config, values, callback)
     if (!query.callback) {
-      let resolveOut, rejectOut
-      result = new Promise((resolve, reject) => {
-        resolveOut = resolve
-        rejectOut = reject
+      result = new this._Promise((resolve, reject) => {
+        query.callback = (err, res) => err ? reject(err) : resolve(res)
       })
-      query.callback = (err, res) => err ? rejectOut(err) : resolveOut(res)
     }
   }
 
@@ -458,7 +456,7 @@ Client.prototype.end = function (cb) {
   if (cb) {
     this.connection.once('end', cb)
   } else {
-    return new Promise((resolve) => {
+    return new this._Promise((resolve) => {
       this.connection.once('end', resolve)
     })
   }
diff --git a/lib/native/client.js b/lib/native/client.js
index b18ff6ffa..c88bfb12e 100644
--- a/lib/native/client.js
+++ b/lib/native/client.js
@@ -25,6 +25,7 @@ var Client = module.exports = function (config) {
   EventEmitter.call(this)
   config = config || {}
 
+  this._Promise = config.Promise || global.Promise
   this._types = new TypeOverrides(config.types)
 
   this.native = new Native({
@@ -121,7 +122,7 @@ Client.prototype.connect = function (callback) {
     return
   }
 
-  return new Promise((resolve, reject) => {
+  return new this._Promise((resolve, reject) => {
     this._connect((error) => {
       if (error) {
         reject(error)
@@ -156,7 +157,7 @@ Client.prototype.query = function (config, values, callback) {
     query = new NativeQuery(config, values, callback)
     if (!query.callback) {
       let resolveOut, rejectOut
-      result = new Promise((resolve, reject) => {
+      result = new this._Promise((resolve, reject) => {
         resolveOut = resolve
         rejectOut = reject
       })
@@ -196,11 +197,8 @@ Client.prototype.end = function (cb) {
   }
   var result
   if (!cb) {
-    var resolve, reject
-    cb = (err) => err ? reject(err) : resolve()
-    result = new global.Promise(function (res, rej) {
-      resolve = res
-      reject = rej
+    result = new this._Promise(function (resolve, reject) {
+      cb = (err) => err ? reject(err) : resolve()
     })
   }
   this.native.end(function () {
diff --git a/package.json b/package.json
index 2e7097235..16fc6c7d6 100644
--- a/package.json
+++ b/package.json
@@ -29,6 +29,7 @@
   },
   "devDependencies": {
     "async": "0.9.0",
+    "bluebird": "3.5.2",
     "co": "4.6.0",
     "eslint": "4.2.0",
     "eslint-config-standard": "10.2.1",
diff --git a/test/integration/client/query-as-promise-tests.js b/test/integration/client/query-as-promise-tests.js
index 625a33392..803b89099 100644
--- a/test/integration/client/query-as-promise-tests.js
+++ b/test/integration/client/query-as-promise-tests.js
@@ -1,4 +1,5 @@
 'use strict'
+var bluebird = require('bluebird')
 var helper = require(__dirname + '/../test-helper')
 var pg = helper.pg
 
@@ -7,10 +8,10 @@ process.on('unhandledRejection', function (e) {
   process.exit(1)
 })
 
-const pool = new pg.Pool()
 const suite = new helper.Suite()
 
 suite.test('promise API', (cb) => {
+  const pool = new pg.Pool()
   pool.connect().then((client) => {
     client.query('SELECT $1::text as name', ['foo'])
       .then(function (result) {
@@ -31,3 +32,24 @@ suite.test('promise API', (cb) => {
       })
   })
 })
+
+suite.test('promise API with configurable promise type', (cb) => {
+  const client = new pg.Client({ Promise: bluebird })
+  const connectPromise = client.connect()
+  assert(connectPromise instanceof bluebird, 'Client connect() returns configured promise')
+
+  connectPromise
+    .then(() => {
+      const queryPromise = client.query('SELECT 1')
+      assert(queryPromise instanceof bluebird, 'Client query() returns configured promise')
+
+      return queryPromise.then(() => {
+        client.end(cb)
+      })
+    })
+    .catch((error) => {
+      process.nextTick(() => {
+        throw error
+      })
+    })
+});