CartoDB-SQL-API/app/services/stream_copy.js

81 lines
2.6 KiB
JavaScript
Raw Normal View History

2018-05-25 23:42:30 +08:00
const zlib = require('zlib');
const PSQL = require('cartodb-psql');
const copyTo = require('pg-copy-streams').to;
const copyFrom = require('pg-copy-streams').from;
2018-05-26 00:47:41 +08:00
const StreamCopyMetrics = require('./stream_copy_metrics');
2018-05-25 23:42:30 +08:00
module.exports = {
2018-05-26 00:47:41 +08:00
to (res, sql, userDbParams, logger, cb) {
2018-05-26 00:50:56 +08:00
let metrics = new StreamCopyMetrics(logger, 'copyto', sql);
2018-05-25 23:42:30 +08:00
const pg = new PSQL(userDbParams);
pg.connect(function (err, client) {
if (err) {
return cb(err);
}
const copyToStream = copyTo(sql);
const pgstream = client.query(copyToStream);
pgstream
.on('error', err => {
pgstream.unpipe(res);
return cb(err);
})
2018-05-26 00:47:41 +08:00
.on('data', data => metrics.addSize(data.length))
2018-05-25 23:42:30 +08:00
.on('end', () => {
2018-05-26 00:47:41 +08:00
metrics.end(copyToStream.rowCount);
2018-05-25 23:42:30 +08:00
return cb(null, metrics);
})
.pipe(res);
});
},
2018-05-26 00:47:41 +08:00
from (req, sql, userDbParams, gzip, logger, cb) {
2018-05-26 00:50:56 +08:00
let metrics = new StreamCopyMetrics(logger, 'copyfrom', sql, gzip);
2018-05-25 23:42:30 +08:00
const pg = new PSQL(userDbParams);
pg.connect(function (err, client) {
if (err) {
return cb(err);
}
let copyFromStream = copyFrom(sql);
const pgstream = client.query(copyFromStream);
pgstream
.on('error', err => {
req.unpipe(pgstream);
2018-05-29 22:19:53 +08:00
return cb(err);
2018-05-25 23:42:30 +08:00
})
.on('end', function () {
2018-05-26 00:47:41 +08:00
metrics.end(copyFromStream.rowCount);
return cb(null, metrics);
2018-05-25 23:42:30 +08:00
});
let requestEnded = false;
req
.on('error', err => {
req.unpipe(pgstream);
pgstream.end();
return cb(err);
})
.on('close', () => {
if (!requestEnded) {
const connection = client.connection;
2018-05-31 22:41:22 +08:00
connection.sendCopyFail('CARTO SQL API: Connection closed by client');
2018-05-25 23:42:30 +08:00
req.unpipe(pgstream);
return cb(new Error('Connection closed by client'));
}
})
.on('data', data => metrics.size += data.length)
2018-05-29 22:19:53 +08:00
.on('end', () => requestEnded = true);
if (gzip) {
req.pipe(zlib.createGunzip()).pipe(pgstream);
} else {
req.pipe(pgstream);
}
2018-05-25 23:42:30 +08:00
});
}
2018-05-26 00:50:56 +08:00
};