[tor-commits] [bridgedb/master] Return a database handle via generator

isis at torproject.org isis at torproject.org
Sat Apr 19 17:02:42 UTC 2014


commit 2a99f95e6bd7357c6adbd2e5f4f09b24d169f770
Author: Matthew Finkel <Matthew.Finkel at 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
+
+ at 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





More information about the tor-commits mailing list