Merge branch 'pgcopy-stream' of github.com:CartoDB/CartoDB-SQL-API into pgcopy-stream-q
This commit is contained in:
Paul Ramsey
commit
a4654f82b2
@ -34,6 +34,7 @@ CopyController.prototype.route = function (app) {
|
||||
authorizationMiddleware(this.metadataBackend),
|
||||
connectionParamsMiddleware(this.userDatabaseService),
|
||||
timeoutLimitsMiddleware(this.metadataBackend),
|
||||
validateCopyQuery(),
|
||||
handleCopyFrom(),
|
||||
responseCopyFrom(),
|
||||
errorMiddleware()
|
||||
@ -48,6 +49,7 @@ CopyController.prototype.route = function (app) {
|
||||
authorizationMiddleware(this.metadataBackend),
|
||||
connectionParamsMiddleware(this.userDatabaseService),
|
||||
timeoutLimitsMiddleware(this.metadataBackend),
|
||||
validateCopyQuery(),
|
||||
handleCopyTo(this.statsClient),
|
||||
errorMiddleware()
|
||||
];
|
||||
@ -57,79 +59,43 @@ CopyController.prototype.route = function (app) {
|
||||
app.get(`${base_url}/sql/copyto`, copyToMiddlewares(RATE_LIMIT_ENDPOINTS_GROUPS.COPY_TO));
|
||||
};
|
||||
|
||||
CopyController.prototype.handleCopyTo = function (req, res, next) {
|
||||
const { q: sql } = req.query;
|
||||
const filename = req.query.filename || 'carto-sql-copyto.dmp';
|
||||
|
||||
if (!sql) {
|
||||
throw new Error("Parameter 'q' is missing");
|
||||
}
|
||||
function handleCopyTo (statsClient) {
|
||||
return function handleCopyToMiddleware (req, res, next) {
|
||||
const sql = req.query.q;
|
||||
const filename = req.query.filename || 'carto-sql-copyto.dmp';
|
||||
|
||||
// Only accept SQL that starts with 'COPY'
|
||||
if (!sql.toUpperCase().startsWith("COPY ")) {
|
||||
throw new Error("SQL must start with COPY");
|
||||
}
|
||||
let metrics = {
|
||||
size: 0,
|
||||
time: null,
|
||||
format: getFormatFromCopyQuery(sql),
|
||||
total_rows: null
|
||||
};
|
||||
|
||||
try {
|
||||
// Open pgsql COPY pipe and stream out to HTTP response
|
||||
const pg = new PSQL(res.locals.userDbParams);
|
||||
pg.connect(function (err, client) {
|
||||
if (err) {
|
||||
return next(err);
|
||||
}
|
||||
res.header("Content-Disposition", `attachment; filename=${encodeURIComponent(filename)}`);
|
||||
res.header("Content-Type", "application/octet-stream");
|
||||
|
||||
let copyToStream = copyTo(sql);
|
||||
const pgstream = client.query(copyToStream);
|
||||
try {
|
||||
const startTime = Date.now();
|
||||
|
||||
res.on('error', next);
|
||||
pgstream.on('error', next);
|
||||
pgstream.on('end', next);
|
||||
|
||||
res.setHeader("Content-Disposition", `attachment; filename=${encodeURIComponent(filename)}`);
|
||||
res.setHeader("Content-Type", "application/octet-stream");
|
||||
|
||||
pgstream.pipe(res);
|
||||
});
|
||||
} catch (err) {
|
||||
next(err);
|
||||
}
|
||||
|
||||
};
|
||||
|
||||
CopyController.prototype.handleCopyFrom = function (req, res, next) {
|
||||
const { q: sql } = req.query;
|
||||
|
||||
if (!sql) {
|
||||
return next(new Error("Parameter 'q' is missing, must be in URL or first field in POST"));
|
||||
}
|
||||
|
||||
// Only accept SQL that starts with 'COPY'
|
||||
if (!sql.toUpperCase().startsWith("COPY ")) {
|
||||
return next(new Error("SQL must start with COPY"));
|
||||
}
|
||||
|
||||
try {
|
||||
const start_time = Date.now();
|
||||
|
||||
// Connect and run the COPY
|
||||
const pg = new PSQL(res.locals.userDbParams);
|
||||
pg.connect(function (err, client) {
|
||||
if (err) {
|
||||
return next(err);
|
||||
}
|
||||
|
||||
let copyFromStream = copyFrom(sql);
|
||||
const pgstream = client.query(copyFromStream);
|
||||
pgstream.on('error', next);
|
||||
pgstream.on('end', function () {
|
||||
const end_time = Date.now();
|
||||
res.body = {
|
||||
time: (end_time - start_time) / 1000,
|
||||
total_rows: copyFromStream.rowCount
|
||||
};
|
||||
|
||||
return next();
|
||||
// Open pgsql COPY pipe and stream out to HTTP response
|
||||
const pg = new PSQL(res.locals.userDbParams);
|
||||
pg.connect(function (err, client) {
|
||||
if (err) {
|
||||
return next(err);
|
||||
}
|
||||
|
||||
const copyToStream = copyTo(sql);
|
||||
const pgstream = client.query(copyToStream);
|
||||
pgstream
|
||||
.on('error', next)
|
||||
.on('data', data => metrics.size += data.length)
|
||||
.on('end', () => {
|
||||
metrics.time = (Date.now() - startTime) / 1000;
|
||||
metrics.total_rows = copyToStream.rowCount;
|
||||
statsClient.set('copyTo', JSON.stringify(metrics));
|
||||
})
|
||||
.pipe(res);
|
||||
});
|
||||
} catch (err) {
|
||||
next(err);
|
||||
@ -139,16 +105,7 @@ CopyController.prototype.handleCopyFrom = function (req, res, next) {
|
||||
|
||||
function handleCopyFrom () {
|
||||
return function handleCopyFromMiddleware (req, res, next) {
|
||||
const { sql } = req.query;
|
||||
|
||||
if (!sql) {
|
||||
return next(new Error("Parameter 'sql' is missing, must be in URL or first field in POST"));
|
||||
}
|
||||
|
||||
// Only accept SQL that starts with 'COPY'
|
||||
if (!sql.toUpperCase().startsWith("COPY ")) {
|
||||
return next(new Error("SQL must start with COPY"));
|
||||
}
|
||||
const sql = req.query.q;
|
||||
|
||||
res.locals.copyFromSize = 0;
|
||||
|
||||
@ -193,27 +150,42 @@ function handleCopyFrom () {
|
||||
};
|
||||
}
|
||||
|
||||
function responseCopyFrom () {
|
||||
function responseCopyFrom (statsClient) {
|
||||
return function responseCopyFromMiddleware (req, res, next) {
|
||||
if (!res.body || !res.body.total_rows) {
|
||||
return next(new Error("No rows copied"));
|
||||
}
|
||||
|
||||
if (req.profiler) {
|
||||
const copyFromMetrics = {
|
||||
size: res.locals.copyFromSize, //bytes
|
||||
format: getFormatFromCopyQuery(req.query.sql),
|
||||
time: res.body.time, //seconds
|
||||
total_rows: res.body.total_rows,
|
||||
gzip: req.get('content-encoding') === 'gzip'
|
||||
};
|
||||
|
||||
req.profiler.add({ copyFrom: copyFromMetrics });
|
||||
res.header('X-SQLAPI-Profiler', req.profiler.toJSONString());
|
||||
}
|
||||
const metrics = {
|
||||
size: res.locals.copyFromSize, //bytes
|
||||
format: getFormatFromCopyQuery(req.query.q),
|
||||
time: res.body.time, //seconds
|
||||
total_rows: res.body.total_rows,
|
||||
gzip: req.get('content-encoding') === 'gzip'
|
||||
};
|
||||
|
||||
statsClient.set('copyFrom', JSON.stringify(metrics));
|
||||
|
||||
res.send(res.body);
|
||||
};
|
||||
}
|
||||
|
||||
function validateCopyQuery () {
|
||||
return function validateCopyQueryMiddleware (req, res, next) {
|
||||
const sql = req.query.q;
|
||||
|
||||
if (!sql) {
|
||||
next(new Error("SQL is missing"));
|
||||
}
|
||||
|
||||
// Only accept SQL that starts with 'COPY'
|
||||
if (!sql.toUpperCase().startsWith("COPY ")) {
|
||||
next(new Error("SQL must start with COPY"));
|
||||
}
|
||||
|
||||
next();
|
||||
};
|
||||
}
|
||||
|
||||
|
||||
module.exports = CopyController;
|
||||
|
||||
@ -32,16 +32,6 @@ describe('copy-endpoints', function() {
|
||||
assert.equal(!!response.time, true);
|
||||
assert.strictEqual(response.total_rows, 6);
|
||||
|
||||
assert.ok(res.headers['x-sqlapi-profiler']);
|
||||
const headers = JSON.parse(res.headers['x-sqlapi-profiler']);
|
||||
assert.ok(headers.copyFrom);
|
||||
const metrics = headers.copyFrom;
|
||||
assert.equal(metrics.size, 57);
|
||||
assert.equal(metrics.format, 'CSV');
|
||||
assert.equal(metrics.time, response.time);
|
||||
assert.equal(metrics.total_rows, response.total_rows);
|
||||
assert.equal(metrics.gzip, false);
|
||||
|
||||
done();
|
||||
});
|
||||
});
|
||||
@ -96,7 +86,7 @@ describe('copy-endpoints', function() {
|
||||
assert.deepEqual(
|
||||
JSON.parse(res.body),
|
||||
{
|
||||
error:["Parameter 'q' is missing, must be in URL or first field in POST"]
|
||||
error:["SQL is missing"]
|
||||
}
|
||||
);
|
||||
done();
|
||||
@ -125,6 +115,25 @@ describe('copy-endpoints', function() {
|
||||
});
|
||||
});
|
||||
|
||||
it('should fail with copyto endpoint and without sql', function(done){
|
||||
assert.response(server, {
|
||||
url: "/api/v1/sql/copyto?" + querystring.stringify({
|
||||
filename: '/tmp/output.dmp'
|
||||
}),
|
||||
headers: {host: 'vizzuality.cartodb.com'},
|
||||
method: 'GET'
|
||||
},{}, function(err, res) {
|
||||
assert.ifError(err);
|
||||
assert.deepEqual(
|
||||
JSON.parse(res.body),
|
||||
{
|
||||
error:["SQL is missing"]
|
||||
}
|
||||
);
|
||||
done();
|
||||
});
|
||||
});
|
||||
|
||||
it('should work with copyfrom and gzip', function(done){
|
||||
assert.response(server, {
|
||||
url: "/api/v1/sql/copyfrom?" + querystring.stringify({
|
||||
@ -142,16 +151,6 @@ describe('copy-endpoints', function() {
|
||||
assert.equal(!!response.time, true);
|
||||
assert.strictEqual(response.total_rows, 6);
|
||||
|
||||
assert.ok(res.headers['x-sqlapi-profiler']);
|
||||
const headers = JSON.parse(res.headers['x-sqlapi-profiler']);
|
||||
assert.ok(headers.copyFrom);
|
||||
const metrics = headers.copyFrom;
|
||||
assert.equal(metrics.size, 57);
|
||||
assert.equal(metrics.format, 'CSV');
|
||||
assert.equal(metrics.time, response.time);
|
||||
assert.equal(metrics.total_rows, response.total_rows);
|
||||
assert.equal(metrics.gzip, true);
|
||||
|
||||
done();
|
||||
});
|
||||
});
|
||||
|
||||
Loading…
Reference in New Issue
Block a user