From bafcb8051fa543bd9b47db9c1b8d10dea01354fe Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Simon=20Mart=C3=ADn?= <oleurud@gmail.com>
Date: Thu, 21 Jun 2018 13:39:33 +0200
Subject: [PATCH] StreamCopy getResults refactor - changing get Result to
 getRowCount - ensuring to and from independecy in StreamCopy

---
 app/controllers/copy_controller.js |  9 +++++----
 app/services/stream_copy.js        | 24 ++++++++++++++++++++----
 2 files changed, 25 insertions(+), 8 deletions(-)

diff --git a/app/controllers/copy_controller.js b/app/controllers/copy_controller.js
index a6eb3dbf..feb12b22 100644
--- a/app/controllers/copy_controller.js
+++ b/app/controllers/copy_controller.js
@@ -80,11 +80,11 @@ function handleCopyTo (logger) {
                     .on('data', data => metrics.addSize(data.length))
                     .on('error', (err) => {
                         metrics.end(null, err);
-                        pgstream.unpipe();
+                        pgstream.unpipe(res);
 
                         return next(err);
                     })
-                    .on('end', () => metrics.end(streamCopy.getResult()))
+                    .on('end', () => metrics.end( streamCopy.getRowCount(StreamCopy.ACTION_TO) ))
                     .pipe(res)
                     .on('close', () => {
                         const err = new Error('Connection closed by client');
@@ -125,9 +125,10 @@ function handleCopyFrom (logger) {
                     })
                     .on('close', () => {
                         const err = new Error('Connection closed by client');
+                        pgstream.emit('cancelQuery', err);
+
                         metrics.end(null, err);
-                        const connection = client.connection;
-                        connection.sendCopyFail('CARTO SQL API: Connection closed by client');
+
                         req.unpipe(pgstream);
                         done();
                         next(err);
diff --git a/app/services/stream_copy.js b/app/services/stream_copy.js
index 3d59aa3b..789a9f64 100644
--- a/app/services/stream_copy.js
+++ b/app/services/stream_copy.js
@@ -3,6 +3,9 @@ const copyTo = require('pg-copy-streams').to;
 const copyFrom = require('pg-copy-streams').from;
 const { Client } = require('pg');
 
+const ACTION_TO = 'to';
+const ACTION_FROM = 'from';
+
 module.exports = class StreamCopy {
     constructor(sql, userDbParams) {
         this.pg = new PSQL(userDbParams);
@@ -13,6 +16,14 @@ module.exports = class StreamCopy {
         this.copyFromStream;
     }
 
+    static get ACTION_TO() {
+        return ACTION_TO;
+    }
+
+    static get ACTION_FROM() {
+        return ACTION_FROM;
+    }
+
     to(cb) {
         this.pg.connect((err, client, done) => {
             if (err) {
@@ -49,16 +60,21 @@ module.exports = class StreamCopy {
 
             pgstream
                 .on('end', () => done())
-                .on('error', err => done(err));
+                .on('error', err => done(err))
+                .on('cancelQuery', err => {
+                    client.connection.sendCopyFail('CARTO SQL API: Connection closed by client');
+                });
 
             cb(null, pgstream, copyFromStream, client, done);
         });
     }
 
-    getResult() {
-        if (this.copyToStream) {
+    getRowCount(action = ACTION_TO) {
+        if (action === ACTION_TO && this.copyToStream) {
             return this.copyToStream.rowCount;
-        } else if (this.copyFromStream) {
+        }
+
+        if (action === ACTION_FROM && this.copyFromStream) {
             return this.copyFromStream.rowCount;
         }
     }