commit 2a99f95e6bd7357c6adbd2e5f4f09b24d169f770 Author: Matthew Finkel Matthew.Finkel@gmail.com Date: Fri Mar 21 02:48:46 2014 +0000
Return a database handle via generator
Define a new idiom for retrieving an instance of bridgedb.Storage.Database
Use: with bridgedb.Storage.getDB() as db: db.<do stuff with the database> --- lib/bridgedb/Storage.py | 156 ++++++++++++++++++++++++++++++++++++++++++++--- 1 file changed, 149 insertions(+), 7 deletions(-)
diff --git a/lib/bridgedb/Storage.py b/lib/bridgedb/Storage.py index fdeea23..d9d8b29 100644 --- a/lib/bridgedb/Storage.py +++ b/lib/bridgedb/Storage.py @@ -9,10 +9,15 @@ import binascii import sqlite3 import time import hashlib +from contextlib import contextmanager +from contextlib import GeneratorContextManager +from functools import wraps from ipaddr import IPAddress, IPv6Address, IPv4Address +import sys
import bridgedb.Stability as Stability from bridgedb.Stability import BridgeHistory +import threading
toHex = binascii.b2a_hex fromHex = binascii.a2b_hex @@ -201,13 +206,14 @@ class BridgeData: self.first_seen = first_seen self.last_seen = last_seen
-class Database: +class Database(object): def __init__(self, sqlite_fname, db_fname=None): if db_fname is None: self._conn = openDatabase(sqlite_fname) else: self._conn = openOrConvertDatabase(sqlite_fname, db_fname) self._cur = self._conn.cursor() + self.sqlite_fname = sqlite_fname
def commit(self): self._conn.commit() @@ -216,6 +222,7 @@ class Database: self._conn.rollback()
def close(self): + #print "Closing DB" self._cur.close() self._conn.close()
@@ -498,12 +505,147 @@ def openOrConvertDatabase(sqlite_file, db_file): conn.commit() return conn
-_THE_DB = None +class DBGeneratorContextManager(GeneratorContextManager): + """Overload __exit__() so we can call the generator many times""" + def __exit__(self, type, value, traceback): + """Significantly based on contextlib.py""" + if type is None: + try: + self.gen.next() + except StopIteration: + return + return + else: + if value is None: + # Need to force instantiation so we can reliably + # tell if we get the same exception back + value = type() + try: + self.gen.throw(type, value, traceback) + raise RuntimeError("generator didn't stop after throw()") + except StopIteration, exc: + # Suppress the exception *unless* it's the same exception that + # was passed to throw(). This prevents a StopIteration + # raised inside the "with" statement from being suppressed + return exc is not value + except: + # only re-raise if it's *not* the exception that was + # passed to throw(), because __exit__() must not raise + # an exception unless __exit__() itself failed. But throw() + # has to raise the exception to signal propagation, so this + # fixes the impedance mismatch between the throw() protocol + # and the __exit__() protocol. + # + if sys.exc_info()[1] is not value: + raise +def contextmanager(func): + @wraps(func) + def helper(*args, **kwds): + return DBGeneratorContextManager(func(*args, **kwds)) + return helper + +_DB_FNAME = None +_LOCK = None +_LOCKED = False +_OPENED_DB = None +_TOC = None + +def clearGlobalDB(): + """Start from scratch""" + global _DB_FNAME + global _LOCK + global _LOCKED + global _OPENED_DB + global _TOC + + _DB_FNAME = None + _LOCK = None + _LOCKED = False + _OPENED_DB = None + _TOC = None + +def checkAndConvertDB(sqlite_file, db_file): + openOrConvertDatabase(sqlite_file, db_file).close() + +def setDBFilename(sqlite_fname): + global _DB_FNAME + _DB_FNAME = sqlite_fname + +@contextmanager +def getDB(block=True): + """Generator: Return a usable database handler + + Always return a :class:`bridgedb.Storage.Database` that is + usable within the current thread. If a connection already exists + and it was created by the current thread, then return the + associated :class:`bridgedb.Storage.Database` instance. Otherwise, + create a new instance, blocking until the existing connection + is closed, if applicable. + + Note: This is a blocking call, be careful about deadlocks! + + :rtype: :class:`bridgedb.Storage.Database` + :returns: An instance of :class:`bridgedb.Storage.Database` used to + query the database + """ + global _DB_FNAME + global _LOCK + global _LOCKED + global _OPENED_DB + global _TOC + + refcount = 0 + + if not _LOCK: + _LOCK = threading.RLock() + this_thread = threading.current_thread().ident
-def setGlobalDB(db): - global _THE_DB - _THE_DB = db + #print "_DB_FNAME: %s, _LOCK: %s, _LOCKED: %s, _OPENED_DB: %s, _TOC: %s" % \ + # (_DB_FNAME, _LOCK, _LOCKED, _OPENED_DB, _TOC) + + try: + if _TOC == this_thread: + refcount += 1 + #print "yielding db from earlier" + logging.debug("Refcount: %d" % refcount) + yield _OPENED_DB + logging.debug("Checking lock status") + logging.debug("lock _LOCKED: %s" % _LOCKED) + locked = _LOCK.acquire(False) + logging.debug("lock aquired: %s" % locked) + if locked: _LOCK.release() + if _LOCK.acquire(block): + db = Database(_DB_FNAME) + assert not _LOCKED + _LOCKED = True + _OPENED_DB = db + _TOC = this_thread + refcount += 1 + #print "yielding db of type: %s" % type(db) + #print "_DB_FNAME: %s, _LOCK: %s, _LOCKED: %s, _OPENED_DB: %s, _TOC: %s" % \ + # (_DB_FNAME, _LOCK, _LOCKED, _OPENED_DB, _TOC) + logging.debug("yielding db") + yield db + else: + #print "yielding False" + yield False + finally: + logging.debug("Refcount is %d in finally" % refcount) + if refcount == 1: + #assert _LOCKED + _LOCK.release() + assert _LOCK.acquire() + _LOCK.release() + _LOCKED = False + _OPENED_DB = None + _TOC = None + db.close() + logging.debug("Refcount resulted in closed") + else: + refcount -= 1
-def getDB(): - return _THE_DB +def closeDB(): + next(getDB(), None)
+def dbIsLocked(): + return _LOCKED