From d1ac31c105ebcaa07735745dd43cb40d29b432e1 Mon Sep 17 00:00:00 2001
From: "matthew.blasius" <matthew.blasius@mandiant.com>
Date: Thu, 5 Nov 2015 16:47:53 -0500
Subject: [PATCH] Support a close callback when closing the stream

---
 index.js      |  3 ++-
 test/close.js | 20 +++++++++++++++++++-
 2 files changed, 21 insertions(+), 2 deletions(-)

diff --git a/index.js b/index.js
index 4be25f1b..e1965092 100644
--- a/index.js
+++ b/index.js
@@ -32,10 +32,11 @@ for(var key in Cursor.prototype) {
   }
 }
 
-QueryStream.prototype.close = function() {
+QueryStream.prototype.close = function(cb) {
   this._closing = true
   var self = this
   Cursor.prototype.close.call(this, function(err) {
+    if (cb) { cb(err); }
     if(err) return self.emit('error', err)
     process.nextTick(function() {
       self.push(null)
diff --git a/test/close.js b/test/close.js
index 35af30b5..c0a30096 100644
--- a/test/close.js
+++ b/test/close.js
@@ -34,7 +34,6 @@ helper('early close', function(client) {
   })
 })
 
-
 helper('should not throw errors after early close', function(client) {
   it('can be closed early without error', function(done) {
     var stream = new QueryStream('SELECT * FROM generate_series(0, 2000) num');
@@ -80,3 +79,22 @@ helper('should not throw errors after early close', function(client) {
     });
   });
 });
+
+helper('close callback', function (client) {
+  it('notifies an optional callback when the conneciton is closed', function (done) {
+    var stream = new QueryStream('SELECT * FROM generate_series(0, $1) num', [10], {batchSize: 2, highWaterMark: 2});
+    var query = client.query(stream);
+    query.once('readable', function() { // only reading once
+      query.read();
+    });
+    query.once('readable', function() {
+      query.close(function () {
+        // nothing to assert.  This test will time out if the callback does not work.
+        done();
+      });
+    });
+    query.on('close', function () {
+      assert(false, "close event should not fire"); // no close event because we did not read to the end of the stream.
+    });
+  });
+});