connect and close working for all connection types

dmanjunath · Jan 7, 2017 · 3d3bdae · 3d3bdae
1 parent 95f75cb
commit 3d3bdae
Show file tree

Hide file tree

Showing 2 changed files with 37 additions and 20 deletions.
diff --git a/README.md b/README.md
@@ -28,7 +28,7 @@ module.exports = redshiftClient;
 You can either initialize a raw one time connection and close it after a single query, or you can open a connection pool and leave it open while your application is running.
 
 ##### ***By default node-redshift uses connection pooling
-##### ***If you want to close a redshift connection by hand you have to use raw connection, you can't end a pool
+
 #### 
 ##### rawConnection
 Pass in the rawConnection parameter in the redshift instantiation options to specify a raw connection.
@@ -126,7 +126,7 @@ There are 4 functions supported by the ORM
 ```
 /**
  * create a new instance of object
- * @param  {Object}   data Object with keys/values to create in database. keys are column names, values are data
+ * @param  {Object or Array}   data Object/Array with keys/values to create in database. keys are column names, values are data
  * @param  {Function} cb   
  * @return {Object}        Object that's inserted into redshift
  */

diff --git a/lib/connection.js b/lib/connection.js
@@ -12,27 +12,16 @@ var Redshift = function(config, options){
 
     if(options && options.rawConnection){
       that.connectionType = 'client';
-      var client = new pg.Client(config);
+      var client = new pg.Client(that.config);
       that.client = client;
     }
     else{
       that.connectionType = 'pool';
       // use connection pooling from pg module
-      var pool = new pg.Pool(config);
-      pool.connect(function(err, client, done) {
-        if(err) {
-          done(client); //https://github.com/brianc/node-postgres/wiki/Example
-          throw err;
-        }
-        else {
-          // store the client instance to make queries with
-          that.client = client;
-          that.pool = pool;
-
-          // store done to call back so it can return connection back to pool
-          // https://github.com/brianc/node-postgres#client-pooling
-          that.done = done;
-        }
+      var pool = new pg.Pool(that.config);
+
+      _connectPool(pool, that, function(err){
+        if(err) throw err;
       });
     }
   }
@@ -41,19 +30,47 @@ var Redshift = function(config, options){
   }
 };
 
+function _connectPool(pool, that, callback){
+  pool.connect(function(err, client, done) {
+    if(err) {
+      done(client); //https://github.com/brianc/node-postgres/wiki/Example
+      callback(err);
+    }
+    else {
+      // store the client instance to make queries with
+      that.client = client;
+      that.pool = pool;
+
+      // store done to call back so it can return connection back to pool
+      // https://github.com/brianc/node-postgres#client-pooling
+      that.done = done;
+      callback();
+    }
+  });
+}
+
 // connection functions for client pooling
 Redshift.prototype.close = function(callback){
   var that = this; //store original context of this because it will change inside callbacks
   if(that.connectionType && that.connectionType === 'client' && that.client) that.client.end(callback);
   else if(that.connectionType && that.connectionType === 'pool' && that.pool) that.pool.end();
 };
 
+
 Redshift.prototype.connect = function(callback){
   var that = this;
-  if(that.connectionType && that.client){
+  if(that.connectionType && that.connectionType === 'client' && that.client){
     that.client.connect(callback);
   }
-  else callback(new Error("Can't connect to database manually without usuing manual connection parameter rawConnection. Please add rawConnection to your config."));
+  else if(that.connectionType === 'pool'){
+    var pool = new pg.Pool(that.config);
+
+    _connectPool(pool, that, function(err){
+      if(err) callback(err);
+      else callback();
+    });
+  }
+  else callback(new Error("Couldn't connect to redshift. Invalid connection type"));
 };
 
 module.exports = Redshift;