Merge branch 'refs/heads/feature/transactions'

dataset/persistence/database.py

@@ -1,5 +1,5 @@
 import logging
-from threading import RLock
+import threading
 
 from sqlalchemy import create_engine
 from migrate.versioning.util import construct_engine
@@ -22,7 +22,8 @@ class Database(object):
         if url.startswith('postgres'):
             kw['poolclass'] = NullPool
         engine = create_engine(url, **kw)
-        self.lock = RLock()
+        self.lock = threading.RLock()
+        self.local = threading.local()
         self.url = url
         self.engine = construct_engine(engine)
         self.metadata = MetaData()
@@ -31,6 +32,49 @@ class Database(object):
             self.metadata.reflect(self.engine)
         self._tables = {}
 
+    @property
+    def executable(self):
+        """ The current connection or engine against which statements
+        will be executed. """
+        if hasattr(self.local, 'connection'):
+            return self.local.connection
+        return self.engine
+
+    def _acquire(self):
+        self.lock.acquire()
+
+    def _release(self):
+        if not hasattr(self.local, 'tx'):
+            self.lock.release()
+        else:
+            self.local.must_release = True
+
+    def begin(self):
+        """ Enter a transaction explicitly. No data will be written
+        until the transaction has been committed. """
+        if not hasattr(self.local, 'connection'):
+            self.local.connection = self.engine.connect()
+        if not hasattr(self.local, 'tx'):
+            self.local.tx = self.local.connection.begin()
+
+    def commit(self):
+        """ Commit the current transaction, making all statements executed
+        since the transaction was begun permanent. """
+        self.local.tx.commit()
+        del self.local.tx
+        if not hasattr(self.local, 'must_release'):
+            self.lock.release()
+            del self.local.must_release
+
+    def rollback(self):
+        """ Roll back the current transaction, discarding all statements
+        executed since the transaction was begun. """
+        self.local.tx.rollback()
+        del self.local.tx
+        if not hasattr(self.local, 'must_release'):
+            self.lock.release()
+            del self.local.must_release
+
     @property
     def tables(self):
         """ Get a listing of all tables that exist in the database.
@@ -51,7 +95,8 @@ class Database(object):
 
             table = db.create_table('population')
         """
-        with self.lock:
+        self._acquire()
+        try:
             log.debug("Creating table: %s on %r" % (table_name, self.engine))
             table = SQLATable(table_name, self.metadata)
             col = Column('id', Integer, primary_key=True)
@@ -59,6 +104,8 @@ class Database(object):
             table.create(self.engine)
             self._tables[table_name] = table
             return Table(self, table)
+        finally:
+            self._release()
 
     def load_table(self, table_name):
         """
@@ -72,11 +119,14 @@ class Database(object):
 
             table = db.load_table('population')
         """
-        with self.lock:
+        self._acquire()
+        try:
             log.debug("Loading table: %s on %r" % (table_name, self))
             table = SQLATable(table_name, self.metadata, autoload=True)
             self._tables[table_name] = table
             return Table(self, table)
+        finally:
+            self._release()
 
     def get_table(self, table_name):
         """
@@ -90,13 +140,16 @@ class Database(object):
             # you can also use the short-hand syntax:
             table = db['population']
         """
-        with self.lock:
         if table_name in self._tables:
             return Table(self, self._tables[table_name])
+        self._acquire()
+        try:
             if self.engine.has_table(table_name):
                 return self.load_table(table_name)
             else:
                 return self.create_table(table_name)
+        finally: 
+            self._release()
 
     def __getitem__(self, table_name):
         return self.get_table(table_name)
@@ -113,7 +166,7 @@ class Database(object):
             for row in res:
                 print row['user'], row['c']
         """
-        return ResultIter(self.engine.execute(query))
+        return ResultIter(self.executable.execute(query))
 
     def __repr__(self):
         return '<Database(%s)>' % self.url

dataset/persistence/table.py

@@ -15,7 +15,7 @@ log = logging.getLogger(__name__)
 class Table(object):
 
     def __init__(self, database, table):
-        self.indexes = {}
+        self.indexes = dict([(i.name, i) for i in table.indexes])
         self.database = database
         self.table = table
         self._is_dropped = False
@@ -39,8 +39,8 @@ class Table(object):
         dropping the table. If you want to re-create the table, make
         sure to get a fresh instance from the :py:class:`Database <dataset.Database>`.
         """
+        self.database._acquire()
         self._is_dropped = True
-        with self.database.lock:
         self.database._tables.pop(self.table.name, None)
         self.table.drop(self.database.engine)
 
@@ -67,7 +67,7 @@ class Table(object):
         self._check_dropped()
         if ensure:
             self._ensure_columns(row, types=types)
-        res = self.database.engine.execute(self.table.insert(row))
+        res = self.database.executable.execute(self.table.insert(row))
         return res.lastrowid
 
     def insert_many(self, rows, chunk_size=1000, ensure=True, types={}):
@@ -133,7 +133,7 @@ class Table(object):
         try:
             filters = self._args_to_clause(dict(clause))
             stmt = self.table.update(filters, row)
-            rp = self.database.engine.execute(stmt)
+            rp = self.database.executable.execute(stmt)
             return rp.rowcount > 0
         except KeyError:
             return False
@@ -151,7 +151,13 @@ class Table(object):
         if ensure:
             self.create_index(keys)
 
-        if not self.update(row, keys, ensure=ensure, types=types):
+        filters = {}
+        for key in keys:
+            filters[key] = row.get(key)
+
+        if self.find_one(**filters) is not None:
+            self.update(row, keys, ensure=ensure, types=types)
+        else:
             self.insert(row, ensure=ensure, types=types)
 
     def delete(self, **_filter):
@@ -171,7 +177,7 @@ class Table(object):
             stmt = self.table.delete(q)
         else:
             stmt = self.table.delete()
-        self.database.engine.execute(stmt)
+        self.database.executable.execute(stmt)
 
     def _ensure_columns(self, row, types={}):
         for column in set(row.keys()) - set(self.table.columns.keys()):
@@ -202,11 +208,14 @@ class Table(object):
             table.create_column('created_at', sqlalchemy.DateTime)
         """
         self._check_dropped()
-        with self.database.lock:
+        self.database._acquire()
+        try:
             if name not in self.table.columns.keys():
                 col = Column(name, type)
                 col.create(self.table,
                        connection=self.database.engine)
+        finally:
+            self.database._release()
 
     def create_index(self, columns, name=None):
         """
@@ -216,18 +225,20 @@ class Table(object):
             table.create_index(['name', 'country'])
         """
         self._check_dropped()
-        with self.database.lock:
         if not name:
             sig = abs(hash('||'.join(columns)))
             name = 'ix_%s_%s' % (self.table.name, sig)
         if name in self.indexes:
             return self.indexes[name]
         try:
+            self.database._acquire()
             columns = [self.table.c[c] for c in columns]
             idx = Index(name, *columns)
             idx.create(self.database.engine)
         except:
             idx = None
+        finally:
+            self.database._release()
         self.indexes[name] = idx
         return idx
 
@@ -239,10 +250,10 @@ class Table(object):
             row = table.find_one(country='United States')
         """
         self._check_dropped()
-        res = list(self.find(_limit=1, **_filter))
+        args = self._args_to_clause(_filter)
-        if not len(res):
+        query = self.table.select(whereclause=args, limit=1)
-            return None
+        rp = self.database.executable.execute(query)
-        return res[0]
+        return rp.fetchone()
 
     def _args_to_order_by(self, order_by):
         if order_by[0] == '-':
@@ -288,7 +299,7 @@ class Table(object):
 
         # query total number of rows first
         count_query = self.table.count(whereclause=args, limit=_limit, offset=_offset)
-        rp = self.database.engine.execute(count_query)
+        rp = self.database.executable.execute(count_query)
         total_row_count = rp.fetchone()[0]
 
         if _step is None or _step is False or _step == 0:
@@ -311,7 +322,7 @@ class Table(object):
                 break
             queries.append(self.table.select(whereclause=args, limit=qlimit,
                                              offset=qoffset, order_by=order_by))
-        return ResultIter((self.database.engine.execute(q) for q in queries))
+        return ResultIter((self.database.executable.execute(q) for q in queries))
 
     def __len__(self):
         """