From 1db9b3ec3dd62aa8dc5de247531f0c6cee0d16ab Mon Sep 17 00:00:00 2001
From: Dan Robinson <dan@heapanalytics.com>
Date: Sun, 6 Apr 2014 22:02:24 -0700
Subject: [PATCH] Adds handling for errors after initial response.

Includes a test.
---
 copy-to.js      | 10 +++++++---
 test/copy-to.js | 23 +++++++++++++++++++++++
 2 files changed, 30 insertions(+), 3 deletions(-)

diff --git a/copy-to.js b/copy-to.js
index 3ead141..404406b 100644
--- a/copy-to.js
+++ b/copy-to.js
@@ -68,10 +68,14 @@ CopyStreamQuery.prototype._transform = function(chunk, enc, cb) {
   }
   while((chunk.length - offset) > 5) {
     var messageCode = chunk[offset]
-    //complete
-    if(messageCode == code.c) {
+    //complete or error
+    if(messageCode == code.c || messageCode == code.E) {
       this._detach()
-      this.connection.stream.unshift(chunk.slice(offset + 5))
+      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();
     }
diff --git a/test/copy-to.js b/test/copy-to.js
index 8e91be7..0545630 100644
--- a/test/copy-to.js
+++ b/test/copy-to.js
@@ -58,3 +58,26 @@ var testLeak = function(rounds) {
 }
 
 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 array_agg(e) FROM (SELECT generate_series(1, 100000000) AS e) t) TO STDOUT"
+
+  var runStream = function(callback) {
+    var stream = fromClient.query(copy(txt))
+    stream.on('data', function(data) {
+      // Just throw away the data.
+    })
+    stream.on('error', callback)
+  }
+
+  runStream(function(err) {
+    assert.notEqual(err, null)
+    expectedMessage = 'invalid memory alloc request size'
+    assert.notEqual(err.toString().indexOf(expectedMessage), -1, 'Error message should mention memory alloc request size.')
+    fromClient.end()
+  })
+}
+
+testInternalPostgresError()