Cancel with user statement_timeout limit from redis

Author: rochoa

batch/index.js

@@ -23,7 +23,7 @@ module.exports = function batchFactory (metadataBackend, redisPool, name, statsd
     var queryRunner = new QueryRunner(userDatabaseMetadataService);
     var jobCanceller = new JobCanceller(userDatabaseMetadataService);
     var jobService = new JobService(jobBackend, jobCanceller);
-    var jobRunner = new JobRunner(jobService, jobQueue, queryRunner, statsdClient);
+    var jobRunner = new JobRunner(jobService, jobQueue, queryRunner, metadataBackend, statsdClient);
     var logger = new BatchLogger(loggerPath);
 
     return new Batch(

batch/job_runner.js

@@ -5,10 +5,16 @@ var jobStatus = require('./job_status');
 var Profiler = require('step-profiler');
 var _ = require('underscore');
 
-function JobRunner(jobService, jobQueue, queryRunner, statsdClient) {
+var REDIS_LIMITS = {
+    DB: 5,
+    PREFIX: 'limits:batch:' // + username
+};
+
+function JobRunner(jobService, jobQueue, queryRunner, metadataBackend, statsdClient) {
     this.jobService = jobService;
     this.jobQueue = jobQueue;
     this.queryRunner = queryRunner;
+    this.metadataBackend = metadataBackend;
     this.statsdClient = statsdClient;
 }
 
@@ -62,7 +68,14 @@ JobRunner.prototype.getQueryStatementTimeout = function(username, callback) {
         timeout = global.settings.batch_query_timeout;
     }
 
-    return callback(null, timeout);
+    var batchLimitsKey = REDIS_LIMITS.PREFIX + username;
+    this.metadataBackend.redisCmd(REDIS_LIMITS.DB, 'HGET', [batchLimitsKey, 'timeout'], function(err, timeoutLimit) {
+        if (timeoutLimit !== null && Number.isFinite(+timeoutLimit)) {
+            timeout = +timeoutLimit;
+        }
+
+        return callback(null, timeout);
+    });
 };
 
 JobRunner.prototype._run = function (job, query, timeout, profiler, callback) {

test/acceptance/batch/batch-limits.test.js

@@ -0,0 +1,48 @@
+require('../../helper');
+
+var assert = require('../../support/assert');
+var BatchTestClient = require('../../support/batch-test-client');
+var JobStatus = require('../../../batch/job_status');
+var redisUtils = require('../../support/redis_utils');
+var metadataBackend = require('cartodb-redis')({ pool: redisUtils.getPool() });
+
+describe('batch query statement_timeout limit', function() {
+
+    before(function(done) {
+        this.batchTestClient = new BatchTestClient();
+        this.batchQueryTimeout = global.settings.batch_query_timeout;
+        global.settings.batch_query_timeout = 15000;
+        metadataBackend.redisCmd(5, 'HMSET', ['limits:batch:vizzuality', 'timeout', 100], done);
+    });
+
+    after(function(done) {
+        global.settings.batch_query_timeout = this.batchQueryTimeout;
+        redisUtils.clean('limits:batch:*', function() {
+            this.batchTestClient.drain(done);
+        }.bind(this));
+    });
+
+    function jobPayload(query) {
+        return {
+            query: query
+        };
+    }
+
+    it('should cancel with user statement_timeout limit', function (done) {
+        var payload = jobPayload('select pg_sleep(10)');
+        this.batchTestClient.createJob(payload, function(err, jobResult) {
+            if (err) {
+                return done(err);
+            }
+            jobResult.getStatus(function (err, job) {
+                if (err) {
+                    return done(err);
+                }
+                assert.equal(job.status, JobStatus.FAILED);
+                assert.ok(job.failed_reason.match(/statement.*timeout/));
+                return done();
+            });
+        });
+    });
+
+});

test/integration/batch/job_runner.test.js

@@ -39,10 +39,12 @@ var JOB = {
 };
 
 describe('job runner', function() {
-    var jobRunner = new JobRunner(jobService, jobQueue, queryRunner, statsdClient);
+    var jobRunner = new JobRunner(jobService, jobQueue, queryRunner, metadataBackend, statsdClient);
 
     after(function (done) {
-        redisUtils.clean('batch:*', done);
+        redisUtils.clean('batch:*', function() {
+            redisUtils.clean('limits:batch:*', done);
+        });
     });
 
     it('.run() should run a job', function (done) {