Merge pull request #497 from CartoDB/cancel-copy-to-client-disconnect

Cancel copy to upon client disconnect

.gitignore

@@ -12,3 +12,5 @@ test/test.log
 test/acceptance/oauth/venv/*
 coverage/
 npm-debug.log
+log/*.log
+yarn.lock

app/services/stream_copy.js

@@ -3,27 +3,59 @@ const PSQL = require('cartodb-psql');
 const copyTo = require('pg-copy-streams').to;
 const copyFrom = require('pg-copy-streams').from;
 const StreamCopyMetrics = require('./stream_copy_metrics');
+const { Client } = require('pg');
 
 module.exports = {
     to (res, sql, userDbParams, user, logger, cb) {
         let metrics = new StreamCopyMetrics(logger, 'copyto', sql, user);
 
         const pg = new PSQL(userDbParams);
-        pg.connect(function (err, client) {
+        pg.connect(function (err, client, done) {
             if (err) {
                 return cb(err);
             }
 
+            let responseEnded = false;
+            let connectionClosedByClient = false;
             const copyToStream = copyTo(sql);
             const pgstream = client.query(copyToStream);
-            pgstream
+
+            res
                 .on('error', err => {
                     pgstream.unpipe(res);
+                    done();
                     return cb(err);
                 })
+                .on('close', () => {
+                    if (!responseEnded) {
+                        connectionClosedByClient = true;
+                        // Cancel the running COPY TO query
+                        // See https://www.postgresql.org/docs/9.5/static/protocol-flow.html#PROTOCOL-COPY
+                        const runningClient = client;
+                        const cancelingClient = new Client(runningClient.connectionParameters);
+                        cancelingClient.cancel(runningClient, pgstream);
+
+                        const err = new Error('Connection closed by client');
+                        pgstream.unpipe(res);
+                        // see https://node-postgres.com/api/pool#releasecallback
+                        done(err);
+                        return cb(err);
+                    }
+                })
+                .on('end', () => responseEnded = true);
+
+            pgstream
+                .on('error', err => {
+                    if (!connectionClosedByClient) {
+                        pgstream.unpipe(res);
+                        done(err);
+                        return cb(err);
+                    }
+                })
                 .on('data', data => metrics.addSize(data.length))
                 .on('end', () => {
                     metrics.end(copyToStream.rowCount);
+                    done();
                     return cb(null, metrics);
                 })
                 .pipe(res);
@@ -34,7 +66,7 @@ module.exports = {
         let metrics = new StreamCopyMetrics(logger, 'copyfrom', sql, user, gzip);
         
         const pg = new PSQL(userDbParams);
-        pg.connect(function (err, client) {
+        pg.connect(function (err, client, done) {
             if (err) {
                 return cb(err);
             }
@@ -44,10 +76,12 @@ module.exports = {
             pgstream
                 .on('error', err => {
                     req.unpipe(pgstream);
+                    done();
                     return cb(err); 
                 })
                 .on('end', function () {
                     metrics.end(copyFromStream.rowCount);
+                    done();
                     return cb(null, metrics);
                 });
 
@@ -57,13 +91,15 @@ module.exports = {
                 .on('error', err => {
                     req.unpipe(pgstream);
                     pgstream.end();
+                    done();
                     return cb(err);
                 })
                 .on('close', () => {
                     if (!requestEnded) {
                         const connection = client.connection;
-                        connection.sendCopyFail();
+                        connection.sendCopyFail('CARTO SQL API: Connection closed by client');
                         req.unpipe(pgstream);
+                        done();
                         return cb(new Error('Connection closed by client'));
                     }
                 })

test/acceptance/copy-endpoints.js

@@ -4,6 +4,7 @@ const fs = require('fs');
 const querystring = require('querystring');
 const assert = require('../support/assert');
 const os = require('os');
+const { Client } = require('pg');
 
 const StatsClient = require('../../app/stats/client');
 if (global.settings.statsd) {
@@ -18,6 +19,20 @@ const server = require('../../app/server')(statsClient);
 
 
 describe('copy-endpoints', function() {
+    before(function(done) {
+        const client = new Client({
+            user: 'postgres',
+            host: 'localhost',
+            database: 'cartodb_test_user_1_db',
+            port: 5432,
+        });
+        client.connect();
+        client.query('TRUNCATE copy_endpoints_test', (err/*, res */) => {
+            client.end();
+            done(err);
+        });
+    });
+
     it('should work with copyfrom endpoint', function(done){
         assert.response(server, {
             url: "/api/v1/sql/copyfrom?" + querystring.stringify({
@@ -255,3 +270,63 @@ describe('copy-endpoints timeout', function() {
         });
     });
 });
+
+
+describe('copy-endpoints db connections', function() {
+    before(function() {
+        this.db_pool_size = global.settings.db_pool_size;
+        global.settings.db_pool_size = 1;
+    });
+
+    after(function() {
+        global.settings.db_pool_size = this.db_pool_size;
+    });
+
+    it('copyfrom', function(done) {
+        const query = "COPY copy_endpoints_test (id, name) FROM STDIN WITH (FORMAT CSV, DELIMITER ',', HEADER true)";
+        function doCopyFrom() {
+            return new Promise(resolve => {
+                assert.response(server, {
+                    url: "/api/v1/sql/copyfrom?" + querystring.stringify({
+                        q: query
+                    }),
+                    data: fs.createReadStream(__dirname + '/../support/csv/copy_test_table.csv'),
+                    headers: {host: 'vizzuality.cartodb.com'},
+                    method: 'POST'
+                },{}, function(err, res) {
+                    assert.ifError(err);
+                    const response = JSON.parse(res.body);
+                    assert.ok(response.time);
+                    resolve();
+                });
+            });
+        }
+
+        Promise.all([doCopyFrom(), doCopyFrom(), doCopyFrom()]).then(function() {
+            done();
+        });
+    });
+
+    it('copyto', function(done) {
+        function doCopyTo() {
+            return new Promise(resolve => {
+                assert.response(server, {
+                    url: "/api/v1/sql/copyto?" + querystring.stringify({
+                        q: 'COPY copy_endpoints_test TO STDOUT',
+                        filename: '/tmp/output.dmp'
+                    }),
+                    headers: {host: 'vizzuality.cartodb.com'},
+                    method: 'GET'
+                },{}, function(err, res) {
+                    assert.ifError(err);
+                    assert.ok(res.body);
+                    resolve();
+                });
+            });
+        }
+
+        Promise.all([doCopyTo(), doCopyTo(), doCopyTo()]).then(function() {
+            done();
+        });
+    });
+});