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

61 lines
1.8 KiB
JavaScript
Raw Normal View History

2018-05-25 23:42:30 +08:00
const PSQL = require('cartodb-psql');
const copyTo = require('pg-copy-streams').to;
const copyFrom = require('pg-copy-streams').from;
const { Client } = require('pg');
2018-05-25 23:42:30 +08:00
const ACTION_TO = 'to';
const ACTION_FROM = 'from';
2018-06-14 02:12:21 +08:00
module.exports = class StreamCopy {
constructor(sql, userDbParams) {
2018-07-20 16:22:54 +08:00
const dbParams = Object.assign({}, userDbParams, {
port: global.settings.db_batch_port || userDbParams.port
});
this.pg = new PSQL(dbParams);
2018-06-12 22:56:18 +08:00
this.sql = sql;
2018-06-21 20:53:07 +08:00
this.stream = null;
2018-06-12 22:56:18 +08:00
}
static get ACTION_TO() {
return ACTION_TO;
}
static get ACTION_FROM() {
return ACTION_FROM;
}
getPGStream(action, cb) {
this.pg.connect((err, client, done) => {
2018-06-08 22:58:32 +08:00
if (err) {
return cb(err);
2018-06-08 22:58:32 +08:00
}
let streamMaker = action === ACTION_TO ? copyTo : copyFrom;
this.stream = streamMaker(this.sql);
2018-06-21 20:53:07 +08:00
const pgstream = client.query(this.stream);
2018-06-08 22:58:32 +08:00
pgstream
2018-06-14 06:54:03 +08:00
.on('end', () => done())
.on('error', err => done(err))
.on('cancelQuery', err => {
if(action === ACTION_TO) {
// See https://www.postgresql.org/docs/9.5/static/protocol-flow.html#PROTOCOL-COPY
const cancelingClient = new Client(client.connectionParameters);
cancelingClient.cancel(client, pgstream);
// see https://node-postgres.com/api/pool#releasecallback
done(err);
} else if (action === ACTION_FROM) {
client.connection.sendCopyFail('CARTO SQL API: Connection closed by client');
}
});
2018-06-08 22:50:12 +08:00
cb(null, pgstream);
2018-06-08 22:50:12 +08:00
});
2018-05-25 23:42:30 +08:00
}
2018-06-21 20:53:07 +08:00
getRowCount() {
return this.stream.rowCount;
}
2018-05-26 00:50:56 +08:00
};