Add max_jobs table to STATEDB

1 files changed, 22 insertions, 0 deletions

diff --git a/lorrycontroller/statedb.py b/lorrycontroller/statedb.py
index a056dc5..7353b0f 100644
--- a/lorrycontroller/statedb.py
+++ b/lorrycontroller/statedb.py
@@ -127,6 +127,11 @@ class StateDB(object):
             'exit TEXT, '
             'output TEXT)')
 
+        # Table for holding max number of jobs running at once. If no
+        # rows, there is no limit. Otherwise, there is exactly one
+        # row.
+        c.execute('CREATE TABLE max_jobs (max_jobs INT)')
+
         # A table to give the current pretended time, if one is set.
         # This table is either empty, in which case time.time() is
         # used, or has one row, which is used for the current time.
@@ -521,3 +526,20 @@ class StateDB(object):
             return row[0]
         else:
             return time.time()
+
+    def get_max_jobs(self):
+        c = self.get_cursor()
+        c.execute('SELECT max_jobs FROM max_jobs')
+        row = c.fetchone()
+        if row:
+            return row[0]
+        return None
+
+    def set_max_jobs(self, max_jobs):
+        logging.debug('StateDB.set_max_jobs(%r) called', max_jobs)
+        assert self.in_transaction
+        c = self.get_cursor()
+        c.execute('DELETE FROM max_jobs')
+        if max_jobs is not None:
+            c.execute(
+                'INSERT INTO max_jobs (max_jobs) VALUES (?)', (max_jobs,))