Implemented multiquery jobs for Batch API

batch/index.js

@@ -2,6 +2,7 @@
 
 var redis = require('redis');
 var JobRunner = require('./job_runner');
+var QueryRunner = require('./query_runner');
 var JobCanceller = require('./job_canceller');
 var JobQueuePool = require('./job_queue_pool');
 var JobSubscriber = require('./job_subscriber');
@@ -20,7 +21,8 @@ module.exports = function batchFactory (metadataBackend) {
     var userIndexer = new UserIndexer(metadataBackend);
     var jobBackend = new JobBackend(metadataBackend, jobQueue, jobPublisher, userIndexer);
     var userDatabaseMetadataService = new UserDatabaseMetadataService(metadataBackend);
-    var jobRunner = new JobRunner(jobBackend, userDatabaseMetadataService);
+    var queryRunner = new QueryRunner();
+    var jobRunner = new JobRunner(jobBackend, jobQueue, queryRunner, userDatabaseMetadataService);
     var jobCanceller = new JobCanceller(metadataBackend, userDatabaseMetadataService, jobBackend);
 
     return new Batch(jobSubscriber, jobQueuePool, jobRunner, jobCanceller);

batch/job_backend.js

@@ -17,6 +17,16 @@ JobBackend.prototype.create = function (username, sql, host, callback) {
     var self = this;
     var job_id = uuid.v4();
     var now = new Date().toISOString();
+
+    if (Array.isArray(sql)) {
+        for (var i = 0; i < sql.length; i++) {
+            sql[i] = {
+                query: sql[i],
+                status: 'pending'
+            };
+        }
+    }
+
     var redisParams = [
         this.redisPrefix + job_id,
         'user', username,
@@ -193,15 +203,22 @@ JobBackend.prototype.get = function (job_id, callback) {
     });
 };
 
-JobBackend.prototype.setRunning = function (job, callback) {
+JobBackend.prototype.setRunning = function (job, index, callback) {
     var self = this;
     var now = new Date().toISOString();
     var redisParams = [
         this.redisPrefix + job.job_id,
         'status', 'running',
-        'updated_at', now
+        'updated_at', now,
     ];
 
+    if (!callback) {
+        callback = index;
+    } else if (index || index === 0) {
+        job.query[index].status = 'running';
+        redisParams = redisParams.concat('query', JSON.stringify(job.query));
+    }
+
     this.metadataBackend.redisCmd(this.db, 'HMSET', redisParams, function (err) {
         if (err) {
             return callback(err);
@@ -211,7 +228,7 @@ JobBackend.prototype.setRunning = function (job, callback) {
     });
 };
 
-JobBackend.prototype.setPending = function (job, callback) {
+JobBackend.prototype.setPending = function (job, index, callback) {
     var self = this;
     var now = new Date().toISOString();
     var redisKey = this.redisPrefix + job.job_id;
@@ -221,6 +238,13 @@ JobBackend.prototype.setPending = function (job, callback) {
         'updated_at', now
     ];
 
+    if (!callback) {
+        callback = index;
+    } else if (index || index === 0) {
+        job.query[index].status = 'pending';
+        redisParams = redisParams.concat('query', JSON.stringify(job.query));
+    }
+
     this.metadataBackend.redisCmd(this.db, 'HMSET', redisParams ,  function (err) {
         if (err) {
             return callback(err);
@@ -230,7 +254,7 @@ JobBackend.prototype.setPending = function (job, callback) {
     });
 };
 
-JobBackend.prototype.setDone = function (job, callback) {
+JobBackend.prototype.setDone = function (job, index, callback) {
     var self = this;
     var now = new Date().toISOString();
     var redisKey = this.redisPrefix + job.job_id;
@@ -240,6 +264,13 @@ JobBackend.prototype.setDone = function (job, callback) {
         'updated_at', now
     ];
 
+    if (!callback) {
+        callback = index;
+    } else if (index || index === 0) {
+        job.query[index].status = 'done';
+        redisParams = redisParams.concat('query', JSON.stringify(job.query));
+    }
+
     this.metadataBackend.redisCmd(this.db, 'HMSET', redisParams ,  function (err) {
         if (err) {
             return callback(err);
@@ -255,7 +286,30 @@ JobBackend.prototype.setDone = function (job, callback) {
     });
 };
 
-JobBackend.prototype.setFailed = function (job, error, callback) {
+JobBackend.prototype.setJobPendingAndQueryDone = function (job, index, callback) {
+    var self = this;
+    var now = new Date().toISOString();
+    var redisKey = this.redisPrefix + job.job_id;
+
+    job.query[index].status = 'done';
+
+    var redisParams = [
+        redisKey,
+        'status', 'pending',
+        'updated_at', now,
+        'query', JSON.stringify(job.query)
+    ];
+
+    this.metadataBackend.redisCmd(this.db, 'HMSET', redisParams ,  function (err) {
+        if (err) {
+            return callback(err);
+        }
+
+        self.get(job.job_id, callback);
+    });
+};
+
+JobBackend.prototype.setFailed = function (job, error, index, callback) {
     var self = this;
     var now = new Date().toISOString();
     var redisKey = this.redisPrefix + job.job_id;
@@ -266,6 +320,13 @@ JobBackend.prototype.setFailed = function (job, error, callback) {
         'updated_at', now
     ];
 
+    if (!callback) {
+        callback = index;
+    } else if (index || index === 0) {
+        job.query[index].status = 'failed';
+        redisParams = redisParams.concat('query', JSON.stringify(job.query));
+    }
+
     this.metadataBackend.redisCmd(this.db, 'HMSET', redisParams , function (err) {
         if (err) {
             return callback(err);
@@ -281,7 +342,7 @@ JobBackend.prototype.setFailed = function (job, error, callback) {
     });
 };
 
-JobBackend.prototype.setCancelled = function (job, callback) {
+JobBackend.prototype.setCancelled = function (job, index, callback) {
     var self = this;
     var now = new Date().toISOString();
     var redisKey = this.redisPrefix + job.job_id;
@@ -291,6 +352,13 @@ JobBackend.prototype.setCancelled = function (job, callback) {
         'updated_at', now
     ];
 
+    if (!callback) {
+        callback = index;
+    } else if (index || index === 0) {
+        job.query[index].status = 'cancelled';
+        redisParams = redisParams.concat('query', JSON.stringify(job.query));
+    }
+
     this.metadataBackend.redisCmd(this.db, 'HMSET', redisParams ,  function (err) {
         if (err) {
             return callback(err);

batch/job_canceller.js

@@ -8,6 +8,16 @@ function JobCanceller(metadataBackend, userDatabaseMetadataService, jobBackend)
     this.jobBackend = jobBackend;
 }
 
+function getIndexOfRunningQuery(job) {
+    if (Array.isArray(job.query)) {
+        for (var i = 0; i < job.query.length; i++) {
+            if (job.query[i].status === 'running') {
+                return i;
+            }
+        }
+    }
+}
+
 JobCanceller.prototype.cancel = function (job_id, callback) {
     var self = this;
 
@@ -36,7 +46,9 @@ JobCanceller.prototype.cancel = function (job_id, callback) {
                     return callback(err);
                 }
 
-                self.jobBackend.setCancelled(job, callback);
+                var queryIndex = getIndexOfRunningQuery(job);
+
+                self.jobBackend.setCancelled(job, queryIndex, callback);
             });
         });
     });

batch/job_runner.js

@@ -1,12 +1,40 @@
 'use strict';
 
 var errorCodes = require('../app/postgresql/error_codes').codeToCondition;
-var PSQL = require('cartodb-psql');
-var queue = require('queue-async');
 
+function getNextQuery(job) {
+    if (!Array.isArray(job.query)) {
+        return {
+            query: job.query
+        };
+    }
 
-function JobRunner(jobBackend, userDatabaseMetadataService) {
+    for (var i = 0; i < job.query.length; i++) {
+        if (job.query[i].status === 'pending') {
+            return {
+                index: i,
+                query: job.query[i].query
+            };
+        }
+    }
+}
+
+function isLastQuery(job, index) {
+    if (!Array.isArray(job.query)) {
+        return true;
+    }
+
+    if (index >= (job.query.length -1)) {
+        return true;
+    }
+
+    return false;
+}
+
+function JobRunner(jobBackend, jobQueue, queryRunner,userDatabaseMetadataService) {
     this.jobBackend = jobBackend;
+    this.jobQueue = jobQueue;
+    this.queryRunner = queryRunner;
     this.userDatabaseMetadataService = userDatabaseMetadataService;
 }
 
@@ -19,107 +47,74 @@ JobRunner.prototype.run = function (job_id, callback) {
         }
 
         if (job.status !== 'pending') {
-            var error = new Error('Cannot run job ' + job.job_id + ' due to its status is ' + job.status);
+            var invalidJobStatusError = new Error([
-            error.name = 'InvalidJobStatus';
+                'Cannot run job',
-            return callback(error);
+                job.job_id,
+                'due to its status is',
+                job.status
+            ].join(' '));
+            invalidJobStatusError.name = 'InvalidJobStatus';
+            return callback(invalidJobStatusError);
         }
 
-        self.userDatabaseMetadataService.getUserMetadata(job.user, function (err, userDatabaseMetadata) {
+        var query = getNextQuery(job);
+
+        if (!query) {
+            var queryNotFoundError = new Error([
+                'Cannot run job',
+                job.job_id,
+                ', there is no query to run'
+            ].join(' '));
+            queryNotFoundError.name = 'QueryNotFound';
+            return callback(queryNotFoundError);
+        }
+
+        self.jobBackend.setRunning(job, query.index, function (err, job) {
             if (err) {
                 return callback(err);
             }
 
-            self.jobBackend.setRunning(job, function (err, job) {
+            self._run(job, query, callback);
-                if (err) {
-                    return callback(err);
-                }
-
-                self._runInSeries(job, userDatabaseMetadata, callback);
-            });
         });
     });
 };
 
-JobRunner.prototype._runInSeries = function(job, userDatabaseMetadata, callback) {
+JobRunner.prototype._run = function (job, query, callback) {
     var self = this;
-    var jobQueue = queue(1); // performs in series
+    self.userDatabaseMetadataService.getUserMetadata(job.user, function (err, userDatabaseMetadata) {
-    var isMultiQuery = true;
-
-    if (!Array.isArray(job.query)) {
-        isMultiQuery = false;
-        job.query = [ job.query ];
-    }
-
-    for (var i = 0; i < job.query.length; i++) {
-        jobQueue.defer(this._run.bind(this), job, userDatabaseMetadata, i, isMultiQuery);
-    }
-
-    jobQueue.await(function (err) {
-        if (err) {
-            // if query has been cancelled then it's going to get the current job status saved by query_canceller
-            if (errorCodes[err.code.toString()] === 'query_canceled') {
-                return self.jobBackend.get(job.job_id, callback);
-            }
-
-            return self.jobBackend.setFailed(job, err, callback);
-        }
-
-        self.jobBackend.setDone(job, callback);
-    });
-};
-
-JobRunner.prototype._run = function (job, userDatabaseMetadata, index, isMultiQuery, callback) {
-    this._query(job, userDatabaseMetadata, index, function (err, result) {
-        var note = '';
-
-        if (err && isMultiQuery) {
-            if (index > 0) {
-                note = '; previous queries have finished successfully';
-            }
-
-            if (index < (job.query.length - 1)) {
-                note += (note ? ' and ' : '; ') + 'later queries were omitted';
-            }
-
-            err.message = 'error on query ' + (index + 1) +': ' + err.message + note;
-
-            return callback(err);
-        }
-
         if (err) {
             return callback(err);
         }
 
-        callback(null, result);
+        self.queryRunner.run(job.job_id, query.query, userDatabaseMetadata, function (err /*, result */) {
-    });
-};
-
-JobRunner.prototype._query = function (job, userDatabaseMetadata, index, callback) {
-    var self = this;
-
-    var pg = new PSQL(userDatabaseMetadata, {}, { destroyOnError: true });
-
-    pg.query('SET statement_timeout=0', function (err) {
-        if(err) {
-            return self.jobBackend.setFailed(job, err, callback);
-        }
-
-        // mark query to allow to users cancel their queries
-        var sql = '/* ' + job.job_id + ' */ ' + job.query[index];
-
-        pg.eventedQuery(sql, function (err, query) {
             if (err) {
+                // if query has been cancelled then it's going to get the current
+                // job status saved by query_canceller
+                if (errorCodes[err.code.toString()] === 'query_canceled') {
+                    return self.jobBackend.get(job.job_id, callback);
+                }
+
                 return self.jobBackend.setFailed(job, err, callback);
             }
 
-            query.on('error', callback);
+            if (isLastQuery(job, query.index)) {
+                console.log('set done', query.index);
+                return self.jobBackend.setDone(job, query.index, callback);
+            }
 
-            query.on('end', function (result) {
+
-                // only if result is present then query is done sucessfully otherwise an error has happened
+            self.jobBackend.setJobPendingAndQueryDone(job, query.index, function (err, job) {
-                // and it was handled by error listener
+                if (err) {
-                if (result) {
+                    return callback(err);
-                    callback(null, result);
                 }
+
+                self.jobQueue.enqueue(job.job_id, userDatabaseMetadata.host, function (err){
+                    if (err) {
+                        return callback(err);
+                    }
+
+                    callback(null, job);
+                });
             });
         });
     });

batch/query_runner.js

@@ -0,0 +1,40 @@
+'use strict';
+
+var PSQL = require('cartodb-psql');
+
+function QueryRunner() {
+}
+
+QueryRunner.prototype.run = function (job_id, sql, userDatabaseMetadata, callback) {
+
+    var pg = new PSQL(userDatabaseMetadata, {}, { destroyOnError: true });
+
+    pg.query('SET statement_timeout=0', function (err) {
+        if(err) {
+            return callback(err);
+        }
+
+        // mark query to allow to users cancel their queries
+        sql = '/* ' + job_id + ' */ ' + sql;
+
+        pg.eventedQuery(sql, function (err, query) {
+            if (err) {
+                return callback(err);
+            }
+
+            query.on('error', callback);
+
+            query.on('end', function (result) {
+                // only if result is present then query is done sucessfully otherwise an error has happened
+                // and it was handled by error listener
+                if (result) {
+                    callback(null, result);
+                }
+            });
+        });
+    });
+
+};
+
+
+module.exports = QueryRunner;

test/acceptance/batch.test.js

@@ -198,16 +198,25 @@ describe('batch module', function() {
     it('should perform job with array of select', function (done) {
         var queries = ['select * from private_table', 'select * from private_table'];
 
+
         createJob(queries, function (err, job) {
             if (err) {
                 return done(err);
             }
 
-            batch.on('job:done', function (job_id) {
+            var queriesDone = 0;
+
+            var checkJobDone = function (job_id) {
                 if (job_id === job.job_id) {
-                    done();
+                    queriesDone += 1;
+                    if (queriesDone === queries.length) {
+                        done();
+                    }
                 }
-            });
+            };
+
+            batch.on('job:done', checkJobDone);
+            batch.on('job:pending', checkJobDone);
         });
     });
 

test/acceptance/job.use-case-9.test.js

@@ -51,8 +51,8 @@ describe('Use case 9: modify a pending multiquery job', function() {
             method: 'POST',
             data: querystring.stringify({
                 query: [
-                    "SELECT * FROM untitle_table_4",
+                    "select pg_sleep(3)",
-                    "select pg_sleep(3)"
+                    "SELECT * FROM untitle_table_4"
                 ]
             })
         }, {