integrate new pool into existing codebase

Author: brianc

lib/index.js

@@ -2,78 +2,43 @@ var EventEmitter = require('events').EventEmitter;
 var util = require('util');
 var Client = require(__dirname+'/client');
 var defaults =  require(__dirname + '/defaults');
-
-//external genericPool module
-var genericPool = require('generic-pool');
-
-//cache of existing client pools
-var pools = {};
+var pool = require(__dirname + '/pool');
+var types = require(__dirname + '/types');
+var Connection = require(__dirname + '/connection');
 
 var PG = function(clientConstructor) {
   EventEmitter.call(this);
-  this.Client = clientConstructor;
-  this.Connection = require(__dirname + '/connection');
-  this.Query = clientConstructor.Query;
   this.defaults = defaults;
+  this.Client = pool.Client = clientConstructor;
+  this.Query = this.Client.Query;
+  this.pools = pool;
+  this.types = types;
+  this.Connection = Connection;
 };
 
 util.inherits(PG, EventEmitter);
 
 PG.prototype.end = function() {
-  Object.keys(pools).forEach(function(name) {
-    var pool = pools[name];
+  var self = this;
+  Object.keys(self.pools.all).forEach(function(key) {
+    var pool = self.pools.all[key];
     pool.drain(function() {
       pool.destroyAllNow();
     });
   });
 };
 
 PG.prototype.connect = function(config, callback) {
-  var self = this;
-  var c = config;
-  var cb = callback;
-  //allow for no config to be passed
-  if(typeof c === 'function') {
-    cb = c;
-    c = defaults;
+  if(typeof config == "function") {
+    callback = config;
+    config = null;
+  }
+  var pool = this.pools.getOrCreate(config);
+  pool.connect(callback);
+  if(!pool.listeners('error').length) {
+    //propagate errors up to pg object
+    pool.on('error', this.emit.bind(this, 'error'));
   }
-
-  //get unique pool name even if object was used as config
-  var poolName = typeof(c) === 'string' ? c : c.user+c.host+c.port+c.database;
-  var pool = pools[poolName];
-
-  if(pool) { return pool.acquire(cb); }
-
-  pool = pools[poolName] = genericPool.Pool({
-    name: poolName,
-    create: function(callback) {
-      var client = new self.Client(c);
-      client.connect(function(err) {
-        if(err) { return callback(err); }
-
-        //handle connected client background errors by emitting event
-        //via the pg object and then removing errored client from the pool
-        client.on('error', function(e) {
-          self.emit('error', e, client);
-          pool.destroy(client);
-        });
-
-        callback(null, client);
-      });
-
-      client.on('drain', function() {
-        pool.release(client);
-      });
-    },
-    destroy: function(client) {
-      client.end();
-    },
-    max: defaults.poolSize,
-    idleTimeoutMillis: defaults.poolIdleTimeout,
-    reapIntervalMillis: defaults.reapIntervalMillis,
-    log: defaults.poolLog
-  });
-  return pool.acquire(cb);
 };
 
 // cancel the query runned by the given client
@@ -96,4 +61,3 @@ module.exports.__defineGetter__("native", function() {
   return module.exports.native;
 });
 
-module.exports.types = require('./types');

lib/pool.js

@@ -3,56 +3,61 @@ var EventEmitter = require('events').EventEmitter;
 var defaults = require(__dirname + '/defaults');
 var genericPool = require('generic-pool');
 
-//takes the same config structure as client
-var createPool = function(clientConfig) {
-  clientConfig = clientConfig || {};
-  var name = JSON.stringify(clientConfig);
-  var pool = createPool.all[name];
-  if(pool) {
-    return pool;
-  }
-  pool = genericPool.Pool({
-    name: name,
-    max: defaults.poolSize,
-    idleTimeoutMillis: defaults.poolIdleTimeout,
-    reapIntervalMillis: defaults.reapIntervalMillis,
-    log: defaults.poolLog,
-    create: function(cb) {
-      var client = new createPool.Client(clientConfig);
-      client.connect(function(err) {
-        if(err) return cb(err, null);
+var pools = {
+  //dictionary of all key:pool pairs
+  all: {},
+  //reference to the client constructor - can override in tests or for require('pg').native
+  Client: require(__dirname + '/client'),
+  getOrCreate: function(clientConfig) {
+    clientConfig = clientConfig || {};
+    var name = JSON.stringify(clientConfig);
+    var pool = pools.all[name];
+    if(pool) {
+      return pool;
+    }
+    pool = genericPool.Pool({
+      name: name,
+      max: defaults.poolSize,
+      idleTimeoutMillis: defaults.poolIdleTimeout,
+      reapIntervalMillis: defaults.reapIntervalMillis,
+      log: defaults.poolLog,
+      create: function(cb) {
+        var client = new pools.Client(clientConfig);
+        client.connect(function(err) {
+          if(err) return cb(err, null);
 
-        //handle connected client background errors by emitting event
-        //via the pg object and then removing errored client from the pool
-        client.on('error', function(e) {
-          pool.emit('error', e, client);
-          pool.destroy(client);
-        });
+          //handle connected client background errors by emitting event
+          //via the pg object and then removing errored client from the pool
+          client.on('error', function(e) {
+            pool.emit('error', e, client);
+            pool.destroy(client);
+          });
 
-        return cb(null, client);
-      });
-    },
-    destroy: function(client) {
-      client.end();
-    }
-  });
-  createPool.all[name] = pool;
-  //mixin EventEmitter to pool
-  EventEmitter.call(pool);
-  for(var key in EventEmitter.prototype) {
-    if(EventEmitter.prototype.hasOwnProperty(key)) {
-      pool[key] = EventEmitter.prototype[key];
+          return cb(null, client);
+        });
+      },
+      destroy: function(client) {
+        client.end();
+      }
+    });
+    pools.all[name] = pool;
+    //mixin EventEmitter to pool
+    EventEmitter.call(pool);
+    for(var key in EventEmitter.prototype) {
+      if(EventEmitter.prototype.hasOwnProperty(key)) {
+        pool[key] = EventEmitter.prototype[key];
+      }
     }
+    //monkey-patch with connect method
+    pool.connect = function(cb) {
+      pool.acquire(function(err, client) {
+        if(err)  return cb(err, null, function() {/*NOOP*/});
+        //support both 2 (old) and 3 arguments
+        (cb.length > 2 ? newConnect : oldConnect)(pool, client, cb);
+      });
+    };
+    return pool;
   }
-  //monkey-patch with connect method
-  pool.connect = function(cb) {
-    pool.acquire(function(err, client) {
-      if(err)  return cb(err, null, function() {/*NOOP*/});
-      //support both 2 (old) and 3 arguments
-      (cb.length > 2 ? newConnect : oldConnect)(pool, client, cb);
-    });
-  };
-  return pool;
 };
 
 //the old connect method of the pool
@@ -62,12 +67,15 @@ var createPool = function(clientConfig) {
 //a bunch of problems, but for backwards compatibility
 //we're leaving it in
 var alarmDuration = 5000;
-var errorMessage = ['A client has been checked out from the pool for longer than ' + alarmDuration + ' ms.',
-'You might have a leak!',
-'You should use the following new way to check out clients','pg.connect(function(err, client, done)) {',
-'  //do something',
-'  done();  //call done() to signal you are finished with the client',
-'}'].join(require('os').EOL);
+var errorMessage = [
+  'A client has been checked out from the pool for longer than ' + alarmDuration + ' ms.',
+  'You might have a leak!',
+  'You should use the following new way to check out clients','pg.connect(function(err, client, done)) {',
+  '  //do something',
+  '  done();  //call done() to signal you are finished with the client',
+  '}'
+].join(require('os').EOL);
+
 var oldConnect = function(pool, client, cb) {
   var tid = setTimeout(function() {
     console.error(errorMessage);
@@ -90,10 +98,4 @@ var newConnect = function(pool, client, cb) {
   });
 };
 
-//list of all created pools
-createPool.all = {};
-
-//reference to client constructor
-createPool.Client = require(__dirname + '/client');
-
-module.exports = createPool;
+module.exports = pools;

test/integration/connection-pool/error-tests.js

@@ -17,6 +17,7 @@ helper.pg.connect(helper.config, assert.success(function(client) {
       assert.ok(error);
       assert.ok(brokenClient);
       assert.equal(client.id, brokenClient.id);
+      client.emit('drain');
       helper.pg.end();
     });
     //kill the connection from client

test/integration/connection-pool/optional-config-tests.js

@@ -10,5 +10,11 @@ helper.pg.defaults.poolSize = 1;
 
 helper.pg.connect(assert.calls(function(err, client) {
   assert.isNull(err);
-  client.end();
+  client.query('SELECT NOW()');
+  client.once('drain', function() {
+    setTimeout(function() {
+      helper.pg.end();
+      
+    }, 10);
+  });
 }));

test/integration/connection-pool/test-helper.js

@@ -9,7 +9,7 @@ helper.testPoolSize = function(max) {
     for(var i = 0; i < max; i++) {
       helper.pg.poolSize = 10;
       test("connection  #" + i + " executes", function() {
-        helper.pg.connect(helper.config, function(err, client) {
+        helper.pg.connect(helper.config, function(err, client, done) {
           assert.isNull(err);
           client.query("select * from person", function(err, result) {
             assert.lengthIs(result.rows, 26)
@@ -19,7 +19,8 @@ helper.testPoolSize = function(max) {
           })
           var query = client.query("SELECT * FROM NOW()")
           query.on('end',function() {
-            sink.add()
+            sink.add();
+            done();
           })
         })
       })