Make better use of transactions

It's no longer feasible to treat each STATEDB operation as
an individual transaction, so we'll make it explicit.

1 files changed, 60 insertions, 41 deletions

diff --git a/lorry-controller-webapp b/lorry-controller-webapp
index efd7cfa..803a7b5 100755
--- a/lorry-controller-webapp
+++ b/lorry-controller-webapp
@@ -16,6 +16,7 @@
 # 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA.
 
 
+import contextlib
 import errno
 import glob
 import logging
@@ -66,6 +67,7 @@ class StateDB(object):
     def __init__(self, filename):
         self._filename = filename
         self._conn = None
+        self._transaction_cursor = None
 
     def _open(self):
         if self._conn is None:
@@ -92,28 +94,47 @@ class StateDB(object):
 
         self._conn.commit()
 
+    @property
+    def in_transaction(self):
+        return self._transaction_cursor is not None
+
+    @contextlib.contextmanager
+    def transaction(self):
+        assert not self.in_transaction
+        self._open()
+        self._transaction_cursor = self._conn.cursor()
+        self._transaction_cursor.execute('BEGIN TRANSACTION')
+        yield
+        self._conn.commit()
+        self._transaction_cursor = None
+
+    def get_cursor(self):
+        '''Return transaction cursor, or a new one.'''
+        if self._transaction_cursor is None:
+            self._open()
+            return self._conn.cursor()
+        else:
+            return self._transaction_cursor
+
     def get_running_queue(self):
         logging.debug('StateDB.get_running_queue called')
-        self._open()
-        c = self._conn.cursor()
+        c = self.get_cursor()
         for (running,) in c.execute('SELECT running FROM running_queue'):
             return bool(running)
 
     def set_running_queue(self, new_status):
         logging.debug('StateDB.set_running_queue(%r) called', new_status)
+        assert self.in_transaction
         if new_status:
             new_value = 1
         else:
             new_value = 0
-        self._open()
-        c = self._conn.cursor()
-        c.execute('UPDATE running_queue SET running = ?', str(new_value))
-        self._conn.commit()
+        self.get_cursor().execute(
+            'UPDATE running_queue SET running = ?', str(new_value))
 
     def get_lorry_info(self, path):
         logging.debug('StateDB.get_lorry_info(path=%r) called', path)
-        self._open()
-        c = self._conn.cursor()
+        c = self.get_cursor()
         c.execute(
             'SELECT path, text, generated, due, interval, running_job '
             'FROM lorries WHERE path IS ?',
@@ -132,8 +153,7 @@ class StateDB(object):
 
     def get_all_lorries_info(self):
         logging.debug('StateDB.get_all_lorries_info called')
-        self._open()
-        c = self._conn.cursor()
+        c = self.get_cursor()
         c.execute(
             'SELECT path, text, generated, due, interval, running_job '
             'FROM lorries '
@@ -152,8 +172,7 @@ class StateDB(object):
 
     def get_lorries_paths(self):
         logging.debug('StateDB.get_lorries_paths called')
-        self._open()
-        c = self._conn.cursor()
+        c = self.get_cursor()
         return [
             row[0]
             for row in c.execute('SELECT path FROM lorries ORDER BY due')]
@@ -174,30 +193,26 @@ class StateDB(object):
         assert generated is not None
         assert due is not None
         assert interval is not None
+        assert self.in_transaction
 
-        self._open()
-        c = self._conn.cursor()
-        c.execute('BEGIN TRANSACTION')
+        c = self.get_cursor()
         c.execute(
             'INSERT OR REPLACE INTO lorries '
             '(path, text, generated, due, interval, running_job) '
             'VALUES (?, ?, ?, ?, ?, ?)',
             (path, text, generated, due, interval, None))
-        self._conn.commit()
 
     def remove_lorry(self, path):
         logging.debug('StateDB.remove_lorry(%r) called', path)
-        self._open()
-        c = self._conn.cursor()
-        c.execute('BEGIN TRANSACTION')
+        assert self.in_transaction
+        c = self.get_cursor()
         c.execute('DELETE FROM lorries WHERE path IS ?', (path,))
-        self._conn.commit()
 
     def set_running_job(self, path, job_id):
         logging.debug(
             'StateDB.set_running_job(%r, %r) called', path, job_id)
-        self._open()
-        c = self._conn.cursor()
+        assert self.in_transaction
+        c = self.get_cursor()
         c.execute(
             'UPDATE lorries SET running_job=? WHERE path=?',
             (job_id, path))
@@ -302,20 +317,21 @@ class ReadConfiguration(LorryControllerRoute):
 
         conf_obj = self.read_config_file()
 
-        existing = set(self.statedb.get_lorries_paths())
-        logging.debug('existing at beginning: %r', existing)
+        with self.statedb.transaction():
+            existing = set(self.statedb.get_lorries_paths())
+            logging.debug('existing at beginning: %r', existing)
 
-        for section in conf_obj:
-            if section['type'] == 'lorries':
-                added = self.add_matching_lorries_to_statedb(section)
-                logging.debug('added: %r', added)
-                existing = existing.difference(added)
+            for section in conf_obj:
+                if section['type'] == 'lorries':
+                    added = self.add_matching_lorries_to_statedb(section)
+                    logging.debug('added: %r', added)
+                    existing = existing.difference(added)
 
-        logging.debug('existing at end: %r', existing)
-        for path in existing:
-            self.statedb.remove_lorry(path)
+            logging.debug('existing at end: %r', existing)
+            for path in existing:
+                self.statedb.remove_lorry(path)
 
-        return 'Configuration has been updated.'
+            return 'Configuration has been updated.'
 
     @property
     def config_file_name(self):
@@ -420,7 +436,8 @@ class StartQueue(LorryControllerRoute):
 
     def run(self, **kwargs):
         logging.debug('%s %s called', self.http_method, self.path)
-        self.statedb.set_running_queue(1)
+        with self.statedb.transaction():
+            self.statedb.set_running_queue(1)
         return 'Queue set to run'
 
 
@@ -431,7 +448,8 @@ class StopQueue(LorryControllerRoute):
 
     def run(self, **kwargs):
         logging.debug('%s %s called', self.http_method, self.path)
-        self.statedb.set_running_queue(0)
+        with self.statedb.transaction():
+            self.statedb.set_running_queue(0)
         return 'Queue set to not run'
 
 
@@ -442,12 +460,13 @@ class GiveMeJob(LorryControllerRoute):
 
     def run(self, **kwargs):
         logging.debug('%s %s called', self.http_method, self.path)
-        lorry_infos = self.statedb.get_all_lorries_info()
-        for lorry_info in lorry_infos:
-            if lorry_info['running_job'] is None:
-                path = lorry_info['path']
-                self.statedb.set_running_job(path, 1)
-                return { 'job-id':  path}
+        with self.statedb.transaction():
+            lorry_infos = self.statedb.get_all_lorries_info()
+            for lorry_info in lorry_infos:
+                if lorry_info['running_job'] is None:
+                    path = lorry_info['path']
+                    self.statedb.set_running_job(path, 1)
+                    return { 'job-id':  path}
         return { 'job-id': None }