1 files changed, 135 insertions, 53 deletions

diff --git a/bitbake/lib/bb/persist_data.py b/bitbake/lib/bb/persist_data.py
index bef7018614..1a6319f949 100644
--- a/bitbake/lib/bb/persist_data.py
+++ b/bitbake/lib/bb/persist_data.py
@@ -29,6 +29,7 @@ import warnings
 from bb.compat import total_ordering
 from collections import Mapping
 import sqlite3
+import contextlib
 
 sqlversion = sqlite3.sqlite_version_info
 if sqlversion[0] < 3 or (sqlversion[0] == 3 and sqlversion[1] < 3):
@@ -45,75 +46,154 @@ if hasattr(sqlite3, 'enable_shared_cache'):
 
 @total_ordering
 class SQLTable(collections.MutableMapping):
+    class _Decorators(object):
+        @staticmethod
+        def retry(f):
+            """
+            Decorator that restarts a function if a database locked sqlite
+            exception occurs.
+            """
+            def wrap_func(self, *args, **kwargs):
+                count = 0
+                while True:
+                    try:
+                        return f(self, *args, **kwargs)
+                    except sqlite3.OperationalError as exc:
+                        if 'is locked' in str(exc) and count < 500:
+                            count = count + 1
+                            self.connection.close()
+                            self.connection = connect(self.cachefile)
+                            continue
+                        raise
+            return wrap_func
+
+        @staticmethod
+        def transaction(f):
+            """
+            Decorator that starts a database transaction and creates a database
+            cursor for performing queries. If no exception is thrown, the
+            database results are commited. If an exception occurs, the database
+            is rolled back. In all cases, the cursor is closed after the
+            function ends.
+
+            Note that the cursor is passed as an extra argument to the function
+            after `self` and before any of the normal arguments
+            """
+            def wrap_func(self, *args, **kwargs):
+                # Context manager will COMMIT the database on success,
+                # or ROLLBACK on an exception
+                with self.connection:
+                    # Automatically close the cursor when done
+                    with contextlib.closing(self.connection.cursor()) as cursor:
+                        return f(self, cursor, *args, **kwargs)
+            return wrap_func
+
     """Object representing a table/domain in the database"""
     def __init__(self, cachefile, table):
         self.cachefile = cachefile
         self.table = table
-        self.cursor = connect(self.cachefile)
-
-        self._execute("CREATE TABLE IF NOT EXISTS %s(key TEXT, value TEXT);"
-                      % table)
-
-    def _execute(self, *query):
-        """Execute a query, waiting to acquire a lock if necessary"""
-        count = 0
-        while True:
-            try:
-                return self.cursor.execute(*query)
-            except sqlite3.OperationalError as exc:
-                if 'database is locked' in str(exc) and count < 500:
-                    count = count + 1
+        self.connection = connect(self.cachefile)
+
+        self._execute_single("CREATE TABLE IF NOT EXISTS %s(key TEXT, value TEXT);" % table)
+
+    @_Decorators.retry
+    @_Decorators.transaction
+    def _execute_single(self, cursor, *query):
+        """
+        Executes a single query and discards the results. This correctly closes
+        the database cursor when finished
+        """
+        cursor.execute(*query)
+
+    @_Decorators.retry
+    def _row_iter(self, f, *query):
+        """
+        Helper function that returns a row iterator. Each time __next__ is
+        called on the iterator, the provided function is evaluated to determine
+        the return value
+        """
+        class CursorIter(object):
+            def __init__(self, cursor):
+                self.cursor = cursor
+
+            def __iter__(self):
+                return self
+
+            def __next__(self):
+                row = self.cursor.fetchone()
+                if row is None:
                     self.cursor.close()
-                    self.cursor = connect(self.cachefile)
-                    continue
-                raise
+                    raise StopIteration
+                return f(row)
+
+            def __enter__(self):
+                return self
+
+            def __exit__(self, typ, value, traceback):
+                self.cursor.close()
+                return False
+
+        cursor = self.connection.cursor()
+        try:
+            cursor.execute(*query)
+            return CursorIter(cursor)
+        except:
+            cursor.close()
 
     def __enter__(self):
-        self.cursor.__enter__()
+        self.connection.__enter__()
         return self
 
     def __exit__(self, *excinfo):
-        self.cursor.__exit__(*excinfo)
-
-    def __getitem__(self, key):
-        data = self._execute("SELECT * from %s where key=?;" %
-                             self.table, [key])
-        for row in data:
+        self.connection.__exit__(*excinfo)
+
+    @_Decorators.retry
+    @_Decorators.transaction
+    def __getitem__(self, cursor, key):
+        cursor.execute("SELECT * from %s where key=?;" % self.table, [key])
+        row = cursor.fetchone()
+        if row is not None:
             return row[1]
         raise KeyError(key)
 
-    def __delitem__(self, key):
+    @_Decorators.retry
+    @_Decorators.transaction
+    def __delitem__(self, cursor, key):
         if key not in self:
             raise KeyError(key)
-        self._execute("DELETE from %s where key=?;" % self.table, [key])
+        cursor.execute("DELETE from %s where key=?;" % self.table, [key])
 
-    def __setitem__(self, key, value):
+    @_Decorators.retry
+    @_Decorators.transaction
+    def __setitem__(self, cursor, key, value):
         if not isinstance(key, str):
             raise TypeError('Only string keys are supported')
         elif not isinstance(value, str):
             raise TypeError('Only string values are supported')
 
-        data = self._execute("SELECT * from %s where key=?;" %
-                                   self.table, [key])
-        exists = len(list(data))
-        if exists:
-            self._execute("UPDATE %s SET value=? WHERE key=?;" % self.table,
-                          [value, key])
+        cursor.execute("SELECT * from %s where key=?;" % self.table, [key])
+        row = cursor.fetchone()
+        if row is not None:
+            cursor.execute("UPDATE %s SET value=? WHERE key=?;" % self.table, [value, key])
         else:
-            self._execute("INSERT into %s(key, value) values (?, ?);" %
-                          self.table, [key, value])
-
-    def __contains__(self, key):
-        return key in set(self)
-
-    def __len__(self):
-        data = self._execute("SELECT COUNT(key) FROM %s;" % self.table)
-        for row in data:
+            cursor.execute("INSERT into %s(key, value) values (?, ?);" % self.table, [key, value])
+
+    @_Decorators.retry
+    @_Decorators.transaction
+    def __contains__(self, cursor, key):
+        cursor.execute('SELECT * from %s where key=?;' % self.table, [key])
+        return cursor.fetchone() is not None
+
+    @_Decorators.retry
+    @_Decorators.transaction
+    def __len__(self, cursor):
+        cursor.execute("SELECT COUNT(key) FROM %s;" % self.table)
+        row = cursor.fetchone()
+        if row is not None:
             return row[0]
 
     def __iter__(self):
-        data = self._execute("SELECT key FROM %s;" % self.table)
-        return (row[0] for row in data)
+        return self._row_iter(lambda row: row[0], "SELECT key from %s;" % self.table)
 
     def __lt__(self, other):
         if not isinstance(other, Mapping):
@@ -122,25 +202,27 @@ class SQLTable(collections.MutableMapping):
         return len(self) < len(other)
 
     def get_by_pattern(self, pattern):
-        data = self._execute("SELECT * FROM %s WHERE key LIKE ?;" %
-                             self.table, [pattern])
-        return [row[1] for row in data]
+        return self._row_iter(lambda row: row[1], "SELECT * FROM %s WHERE key LIKE ?;" %
+                              self.table, [pattern])
 
     def values(self):
         return list(self.itervalues())
 
     def itervalues(self):
-        data = self._execute("SELECT value FROM %s;" % self.table)
-        return (row[0] for row in data)
+        return self._row_iter(lambda row: row[0], "SELECT value FROM %s;" %
+                              self.table)
 
     def items(self):
         return list(self.iteritems())
 
     def iteritems(self):
-        return self._execute("SELECT * FROM %s;" % self.table)
+        return self._row_iter(lambda row: (row[0], row[1]), "SELECT * FROM %s;" %
+                              self.table)
 
-    def clear(self):
-        self._execute("DELETE FROM %s;" % self.table)
+    @_Decorators.retry
+    @_Decorators.transaction
+    def clear(self, cursor):
+        cursor.execute("DELETE FROM %s;" % self.table)
 
     def has_key(self, key):
         return key in self
@@ -195,7 +277,7 @@ class PersistData(object):
         del self.data[domain][key]
 
 def connect(database):
-    connection = sqlite3.connect(database, timeout=5, isolation_level=None)
+    connection = sqlite3.connect(database, timeout=5)
     connection.execute("pragma synchronous = off;")
     connection.text_factory = str
     return connection