Bump version

* Expose row count after COPY FROM command

* correct conditional

.travis.yml

@@ -1,6 +1,17 @@
 language: node_js
 node_js:
-  - "0.10"
-  - "0.11"
+  - "0.12"
+  - "4"
+  - "5"
+  - "6"
+
+addons:
+  postgresql: "9.2"
+
+services:
+  - postgresql
+
+before_install:
+  - npm install npm --global
 env:
-  - PGUSER=postgres
+  - PGUSER=postgres PGDATABASE=postgres

README.md

@@ -1,5 +1,7 @@
 ## pg-copy-streams
 
+[![Build Status](https://travis-ci.org/brianc/node-pg-copy-streams.svg)](https://travis-ci.org/brianc/node-pg-copy-streams)
+
 COPY FROM / COPY TO for node-postgres.  Stream from one database to another, and stuff.
 
 ## how? what? huh?
@@ -38,7 +40,7 @@ var copyFrom = require('pg-copy-streams').from;
 
 pg.connect(function(err, client, done) {
   var stream = client.query(copyFrom('COPY my_table FROM STDIN'));
-  var fileStream = fs.createReadStream('some_file.tdv')
+  var fileStream = fs.createReadStream('some_file.tsv')
   fileStream.on('error', done);
   fileStream.pipe(stream).on('finish', done).on('error', done);
 });

copy-to.js

@@ -4,11 +4,11 @@ module.exports = function(txt, options) {
 
 var Transform = require('stream').Transform
 var util = require('util')
+var code = require('./message-formats')
 
 var CopyStreamQuery = function(text, options) {
   Transform.call(this, options)
   this.text = text
-  this._listeners = {}
   this._copyOutResponse = null
   this.rowCount = 0
 }
@@ -20,31 +20,16 @@ var eventTypes = ['close', 'data', 'end', 'error']
 CopyStreamQuery.prototype.submit = function(connection) {
   connection.query(this.text)
   this.connection = connection
-  var self = this
-  eventTypes.forEach(function(type) {
-    self._listeners[type] = connection.stream.listeners(type)
-    connection.stream.removeAllListeners(type)
-  })
+  this.connection.removeAllListeners('copyData')
   connection.stream.pipe(this)
 }
 
-var code = {
-  E: 69, //Error
-  H: 72, //CopyOutResponse
-  d: 0x64, //CopyData
-  c: 0x63 //CopyDone
+CopyStreamQuery.prototype._detach = function() {
+  this.connection.stream.unpipe(this)
+  // Unpipe can drop us out of flowing mode
+  this.connection.stream.resume()
 }
 
-CopyStreamQuery.prototype._detach = function() {
-  this.connection.stream.unpipe()
-  var self = this
-  eventTypes.forEach(function(type) {
-    self.connection.stream.removeAllListeners(type)
-    self._listeners[type].forEach(function(listener) {
-      self.connection.stream.on(type, listener)
-    })
-  })
-}
 
 CopyStreamQuery.prototype._transform = function(chunk, enc, cb) {
   var offset = 0
@@ -53,14 +38,13 @@ CopyStreamQuery.prototype._transform = function(chunk, enc, cb) {
   }
   if(!this._copyOutResponse) {
     this._copyOutResponse = true
-    if(chunk[0] == code.E) {
+    if(chunk[0] == code.ErrorResponse) {
       this._detach()
-      this.connection.stream.unshift(chunk)
       this.push(null)
       return cb();
     }
-    if(chunk[0] != code.H) {
-      this.emit('error', new Error('Expected copy out response'))
+    if(chunk[0] != code.CopyOutResponse) {
+      this.emit('error', new Error('Expected CopyOutResponse code (H)'))
     }
     var length = chunk.readUInt32BE(1)
     offset = 1
@@ -69,19 +53,14 @@ CopyStreamQuery.prototype._transform = function(chunk, enc, cb) {
   while((chunk.length - offset) > 5) {
     var messageCode = chunk[offset]
     //complete or error
-    if(messageCode == code.c || messageCode == code.E) {
+    if(messageCode == code.CopyDone || messageCode == code.ErrorResponse) {
       this._detach()
-      if (messageCode == code.c) {
-        this.connection.stream.unshift(chunk.slice(offset + 5))
-      } else {
-        this.connection.stream.unshift(chunk.slice(offset))
-      }
       this.push(null)
       return cb();
     }
     //something bad happened
-    if(messageCode != code.d) {
-      return this.emit('error', new Error('expected "d" (copydata message)'))
+    if(messageCode != code.CopyData) {
+      return this.emit('error', new Error('Expected CopyData code (d)'))
     }
     var length = chunk.readUInt32BE(offset + 1) - 4 //subtract length of UInt32
     //can we read the next row?
@@ -108,6 +87,9 @@ CopyStreamQuery.prototype.handleError = function(e) {
   this.emit('error', e)
 }
 
+CopyStreamQuery.prototype.handleCopyData = function(chunk) {
+}
+
 CopyStreamQuery.prototype.handleCommandComplete = function() {
 }
 

index.js

@@ -1,5 +1,4 @@
 var CopyToQueryStream = require('./copy-to')
-
 module.exports = {
   to: function(txt, options) {
     return new CopyToQueryStream(txt, options)
@@ -11,6 +10,7 @@ module.exports = {
 
 var Transform = require('stream').Transform
 var util = require('util')
+var code = require('./message-formats')
 
 var CopyStreamQuery = function(text, options) {
   Transform.call(this, options)
@@ -27,28 +27,21 @@ CopyStreamQuery.prototype.submit = function(connection) {
   connection.query(this.text)
 }
 
-var code = {
-  H: 72, //CopyOutResponse
-  d: 0x64, //CopyData
-  c: 0x63 //CopyDone
-}
 
-var copyDataBuffer = Buffer([code.d])
+var copyDataBuffer = Buffer([code.CopyData])
 CopyStreamQuery.prototype._transform = function(chunk, enc, cb) {
   this.push(copyDataBuffer)
   var lenBuffer = Buffer(4)
   lenBuffer.writeUInt32BE(chunk.length + 4, 0)
   this.push(lenBuffer)
   this.push(chunk)
-  this.rowCount++
   cb()
 }
 
 CopyStreamQuery.prototype._flush = function(cb) {
-  var finBuffer = Buffer([code.c, 0, 0, 0, 4])
+  var finBuffer = Buffer([code.CopyDone, 0, 0, 0, 4])
   this.push(finBuffer)
-  //never call this callback, do not close underlying stream
-  //cb()
+  cb()
 }
 
 CopyStreamQuery.prototype.handleError = function(e) {
@@ -56,10 +49,17 @@ CopyStreamQuery.prototype.handleError = function(e) {
 }
 
 CopyStreamQuery.prototype.handleCopyInResponse = function(connection) {
-  this.pipe(connection.stream)
+  this.pipe(connection.stream, { end: false })
 }
 
-CopyStreamQuery.prototype.handleCommandComplete = function() {
+CopyStreamQuery.prototype.handleCommandComplete = function(msg) {
+  // Parse affected row count as in
+  // https://github.com/brianc/node-postgres/blob/35e5567f86774f808c2a8518dd312b8aa3586693/lib/result.js#L37
+  var match = /COPY (\d+)/.exec((msg || {}).text)
+  if (match) {
+    this.rowCount = parseInt(match[1], 10)
+  }
+
   this.unpipe()
   this.emit('end')
 }

message-formats.js

@@ -0,0 +1,17 @@
+/**
+ * The COPY feature uses the following protocol codes.
+ * The codes for the most recent protocol version are documented on
+ * https://www.postgresql.org/docs/current/static/protocol-message-formats.html
+ * 
+ * The protocol flow itself is described on
+ * https://www.postgresql.org/docs/current/static/protocol-flow.html
+ */
+module.exports = {
+  ErrorResponse:    0x45,
+  CopyInResponse:   0x47,
+  CopyOutResponse:  0x48,
+  CopyBothResponse: 0x57,
+  CopyData:         0x64,
+  CopyDone:         0x63,
+  CopyFail:         0x66
+} 

package.json

@@ -1,6 +1,6 @@
 {
   "name": "pg-copy-streams",
-  "version": "0.3.0",
+  "version": "1.1.1",
   "description": "Low-Level COPY TO and COPY FROM streams for PostgreSQL in JavaScript using",
   "main": "index.js",
   "scripts": {
@@ -23,7 +23,7 @@
     "url": "https://github.com/brianc/node-pg-copy-streams/issues"
   },
   "devDependencies": {
-    "pg.js": "~2.8.1",
+    "pg": "~4.4.3",
     "concat-stream": "~1.1.0",
     "gonna": "0.0.0",
     "lodash": "~2.2.1",

test/binary.js

@@ -4,7 +4,7 @@ var gonna = require('gonna')
 var async = require('async')
 var concat = require('concat-stream')
 var _ = require('lodash')
-var pg = require('pg.js')
+var pg = require('pg')
 
 var from = require('../').from
 var to = require('../').to

test/copy-from.js

@@ -3,7 +3,7 @@ var gonna = require('gonna')
 
 var concat = require('concat-stream')
 var _ = require('lodash')
-var pg = require('pg.js')
+var pg = require('pg')
 
 var copy = require('../').from
 
@@ -30,10 +30,6 @@ var testRange = function(top) {
 
   var txt = 'COPY numbers FROM STDIN'
   var stream = fromClient.query(copy(txt))
-  var rowEmitCount = 0
-  stream.on('row', function() {
-    rowEmitCount++
-  })
   for(var i = 0; i < top; i++) {
     stream.write(Buffer('' + i + '\t' + i*10 + '\n'))
   }

test/copy-to.js

@@ -4,7 +4,7 @@ var gonna = require('gonna')
 var _ = require('lodash')
 var async = require('async')
 var concat = require('concat-stream')
-var pg = require('pg.js')
+var pg = require('pg')
 
 var copy = require('../').to
 
@@ -25,6 +25,8 @@ testConstruction()
 var testRange = function(top) {
   var fromClient = client()
   var txt = 'COPY (SELECT * from generate_series(0, ' + (top - 1) + ')) TO STDOUT'
+  var res;
+
 
   var stream = fromClient.query(copy(txt))
   var done = gonna('finish piping out', 1000, function() {
@@ -32,57 +34,43 @@ var testRange = function(top) {
   })
 
   stream.pipe(concat(function(buf) {
-    var res = buf.toString('utf8')
+    res = buf.toString('utf8')
+  }))
+
+  stream.on('end', function() {
     var expected = _.range(0, top).join('\n') + '\n'
     assert.equal(res, expected)
     assert.equal(stream.rowCount, top, 'should have rowCount ' + top + ' but got ' + stream.rowCount)
     done()
-  }))
+  });
 }
 
 testRange(10000)
 
-var testLeak = function(rounds) {
-  var fromClient = client()
-  var txt = 'COPY (SELECT 10) TO STDOUT'
-
-  var runStream = function(num, callback) {
-    var stream = fromClient.query(copy(txt))
-    stream.on('data', function(data) {
-      // Just throw away the data.
-    })
-    stream.on('end', callback)
-    stream.on('error', callback)
-  }
-
-  async.timesSeries(rounds, runStream, function(err) {
-    assert.equal(err, null)
-    assert.equal(fromClient.connection.stream.listeners('close').length, 0)
-    assert.equal(fromClient.connection.stream.listeners('data').length, 1)
-    assert.equal(fromClient.connection.stream.listeners('end').length, 2)
-    assert.equal(fromClient.connection.stream.listeners('error').length, 1)
-    fromClient.end()
-  })
-}
-
-testLeak(5)
-
 var testInternalPostgresError = function() {
-  var fromClient = client()
-  // This attempts to make an array that's too large, and should fail.
-  var txt = "COPY (SELECT asdlfsdf AS e) t) TO STDOUT"
+  var cancelClient = client()
+  var queryClient = client()
 
   var runStream = function(callback) {
-    var stream = fromClient.query(copy(txt))
+    var txt = "COPY (SELECT pg_sleep(10)) TO STDOUT"
+    var stream = queryClient.query(copy(txt))
     stream.on('data', function(data) {
       // Just throw away the data.
     })
     stream.on('error', callback)
+
+    setTimeout(function() {
+      var cancelQuery = "SELECT pg_cancel_backend(pid) FROM pg_stat_activity WHERE query ~ 'pg_sleep' AND NOT query ~ 'pg_cancel_backend'"
+      cancelClient.query(cancelQuery)
+    }, 50)
   }
 
   runStream(function(err) {
     assert.notEqual(err, null)
-    fromClient.end()
+    var expectedMessage = 'canceling statement due to user request'
+    assert.notEqual(err.toString().indexOf(expectedMessage), -1, 'Error message should mention reason for query failure.')
+    cancelClient.end()
+    queryClient.end()
   })
 }