Merge branch 'refs/heads/feature/transactions'

Browse Source
This commit is contained in:
Friedrich Lindenberg 2013-06-13 11:57:23 +02:00
commit 10dfcad00b
2 changed files with 102 additions and 38 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

69
dataset/persistence/database.py
View File

@ -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])
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

71
dataset/persistence/table.py
View File

@ -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,10 +39,10 @@ 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)
self.database._tables.pop(self.table.name, None)
self.table.drop(self.database.engine)
def _check_dropped(self):
if self._is_dropped:
@ -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)
connection=self.database.engine)
finally:
self.database._release()
def create_index(self, columns, name=None):
"""
@ -216,20 +225,22 @@ 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:
columns = [self.table.c[c] for c in columns]
idx = Index(name, *columns)
idx.create(self.database.engine)
except:
idx = None
self.indexes[name] = idx
return idx
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
def find_one(self, **_filter):
"""
@ -239,10 +250,10 @@ class Table(object):
row = table.find_one(country='United States')
"""
self._check_dropped()
res = list(self.find(_limit=1, **_filter))
if not len(res):
return None
return res[0]
args = self._args_to_clause(_filter)
query = self.table.select(whereclause=args, limit=1)
rp = self.database.executable.execute(query)
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):
"""