Support ping() method for steady connections.

Author: Cito

DBUtils/PersistentDB.py

@@ -176,7 +176,7 @@ def steady_connection(self):
         """Get a steady, non-persistent DB-API 2 connection."""
         return connect(self._creator,
             self._maxusage, self._setsession,
-            self._failures, self._closeable,
+            self._failures, self._closeable, None,
             *self._args, **self._kwargs)
 
     def connection(self, shareable=False):

DBUtils/PooledDB.py

@@ -246,8 +246,9 @@ def wait():
     def steady_connection(self):
         """Get a steady, unpooled DB-API 2 connection."""
         return connect(self._creator,
-            self._maxusage, self._setsession, self._failures, True,
-                *self._args, **self._kwargs)
+            self._maxusage, self._setsession,
+            self._failures, True, None,
+            *self._args, **self._kwargs)
 
     def connection(self, shareable=True):
         """Get a steady, cached DB-API 2 connection from the pool.

DBUtils/SteadyDB.py

@@ -104,7 +104,7 @@ class InvalidCursor(SteadyDBError):
 
 
 def connect(creator, maxusage=None, setsession=None, failures=None,
-        closeable=True, *args, **kwargs):
+        closeable=True, ping=None, *args, **kwargs):
     """A tough version of the connection constructor of a DB-API 2 module.
 
     creator: either an arbitrary function returning new DB-API 2 compliant
@@ -120,12 +120,14 @@ def connect(creator, maxusage=None, setsession=None, failures=None,
         (OperationalError, InternalError) is not adequate
     closeable: if this is set to false, then closing the connection will
         be silently ignored, but by default the connection can be closed
+    ping: determines when the connection should be checked with ping()
+        (0 = default = never, 1 = on cursor(), 2 = on execute(), 3 = both)
     args, kwargs: the parameters that shall be passed to the creator
         function or the connection constructor of the DB-API 2 module
 
     """
     return SteadyDBConnection(creator, maxusage, setsession, failures,
-        closeable, *args, **kwargs)
+        closeable, ping, *args, **kwargs)
 
 
 class SteadyDBConnection:
@@ -134,7 +136,7 @@ class SteadyDBConnection:
     version = __version__
 
     def __init__(self, creator, maxusage=None, setsession=None, failures=None,
-            closeable=True, *args, **kwargs):
+            closeable=True, ping=None, *args, **kwargs):
         """Create a "tough" DB-API 2 connection."""
         # basic initialization to make finalizer work
         self._con = None
@@ -175,6 +177,7 @@ def __init__(self, creator, maxusage=None, setsession=None, failures=None,
             raise TypeError("'failures' must be a tuple of exceptions.")
         self._failures = failures
         self._closeable = closeable
+        self._ping = ping or 0
         self._args, self._kwargs = args, kwargs
         self._store(self._create())
 
@@ -333,6 +336,10 @@ def rollback(self):
         """Rollback pending transaction."""
         self._con.rollback()
 
+    def ping(self, *args, **kwargs):
+        """Ping connection."""
+        return self._con.ping(*args, **kwargs)
+
     def _cursor(self, *args, **kwargs):
         """A "tough" version of the method cursor()."""
         # The args and kwargs are not part of the standard,
@@ -342,7 +349,18 @@ def _cursor(self, *args, **kwargs):
                 if self._usage >= self._maxusage:
                     # the connection was used too often
                     raise self._failure
-            cursor = self._con.cursor(*args, **kwargs) # try to get a cursor
+            con = self._con
+            if self._ping & 1:
+                try: # if possible, ping the connection
+                    ping = con.ping()
+                except (AttributeError, TypeError, ValueError):
+                    self._ping = 0 # ping() is not available
+                    ping = None
+                except Exception:
+                    ping = False
+                if ping is not None and not ping:
+                    raise self._failure # connection is dead
+            cursor = con.cursor(*args, **kwargs) # try to get a cursor
         except self._failures, error: # error in getting cursor
             try: # try to reopen the connection
                 con2 = self._create()
@@ -436,20 +454,31 @@ def _get_tough_method(self, name):
         """Return a "tough" version of the given cursor method."""
         def tough_method(*args, **kwargs):
             execute = name.startswith('execute')
+            con = self._con
             try:
-                if self._con._maxusage:
-                    if self._con._usage >= self._con._maxusage:
+                if con._maxusage:
+                    if con._usage >= con._maxusage:
                         # the connection was used too often
-                        raise self._con._failure
+                        raise con._failure
+                if con._ping & 2:
+                    try: # if possible, ping the connection
+                        ping = con._con.ping()
+                    except (AttributeError, TypeError, ValueError):
+                        self._ping = 0 # ping() is not available
+                        ping = None
+                    except Exception:
+                        ping = False
+                    if ping is not None and not ping:
+                        raise con._failure # connection is dead
                 if execute:
                     self._setsizes()
                 method = getattr(self._cursor, name)
                 result = method(*args, **kwargs) # try to execute
                 if execute:
                     self._clearsizes()
-            except self._con._failures, error: # execution error
+            except con._failures, error: # execution error
                 try:
-                    cursor2 = self._con._cursor(
+                    cursor2 = con._cursor(
                         *self._args, **self._kwargs) # open new cursor
                 except Exception:
                     pass
@@ -466,14 +495,14 @@ def tough_method(*args, **kwargs):
                     else:
                         self.close()
                         self._cursor = cursor2
-                        self._con._usage += 1
+                        con._usage += 1
                         return result
                     try:
                         cursor2.close()
                     except Exception:
                         pass
                 try: # try to reopen the connection
-                    con2 = self._con._create()
+                    con2 = con._create()
                 except Exception:
                     pass
                 else:
@@ -499,10 +528,10 @@ def tough_method(*args, **kwargs):
                             error = None
                         if use2:
                             self.close()
-                            self._con._close()
-                            self._con._store(con2)
+                            con._close()
+                            con._store(con2)
                             self._cursor = cursor2
-                            self._con._usage += 1
+                            con._usage += 1
                             if error:
                                 raise error # raise the other error
                             return result
@@ -516,7 +545,7 @@ def tough_method(*args, **kwargs):
                         pass
                 raise error # reraise the original error again
             else:
-                self._con._usage += 1
+                con._usage += 1
                 return result
         return tough_method
 

DBUtils/Tests/TestSteadyDB.py

@@ -36,6 +36,9 @@ def connect(database=None, user=None):
 
 class Connection:
 
+    has_ping = False
+    num_pings = 0
+
     def __init__(self, database=None, user=None):
         self.database = database
         self.user = user
@@ -45,6 +48,7 @@ def __init__(self, database=None, user=None):
         self.open_cursors = 0
         self.num_uses = 0
         self.num_queries = 0
+        self.num_pings = 0
         self.session = []
         self.valid = True
 
@@ -63,6 +67,14 @@ def commit(self):
     def rollback(self):
         self.session.append('rollback')
 
+    def ping(self):
+        cls = self.__class__
+        cls.num_pings += 1
+        if not cls.has_ping:
+            raise AttributeError
+        if not self.valid:
+            raise OperationalError
+
     def cursor(self, name=None):
         if not self.valid:
             raise InternalError
@@ -200,12 +212,23 @@ def test01_MockedDBConnection(self):
         self.assertRaises(InternalError, cursor.close)
         self.assertRaises(InternalError, cursor.execute, 'select test')
         self.assert_(db.valid)
+        self.assert_(not db.__class__.has_ping)
+        self.assertEqual(db.__class__.num_pings, 0)
+        self.assertRaises(AttributeError, db.ping)
+        self.assertEqual(db.__class__.num_pings, 1)
+        db.__class__.has_ping = True
+        self.assert_(db.ping() is None)
+        self.assertEqual(db.__class__.num_pings, 2)
         db.close()
         self.assert_(not db.valid)
         self.assertEqual(db.num_uses, 0)
         self.assertEqual(db.num_queries, 0)
         self.assertRaises(InternalError, db.close)
         self.assertRaises(InternalError, db.cursor)
+        self.assertRaises(OperationalError, db.ping)
+        self.assertEqual(db.__class__.num_pings, 3)
+        db.__class__.has_ping = False
+        db.__class__.num_pings = 0
 
     def test02_BrokenDBConnection(self):
         self.assertRaises(TypeError, SteadyDBConnection, None)
@@ -239,7 +262,7 @@ def test03_SteadyDBClose(self):
             self.assert_(not db._con.valid)
 
     def test04_SteadyDBConnection(self):
-        db = SteadyDBconnect(dbapi, 0, None, None, 1,
+        db = SteadyDBconnect(dbapi, 0, None, None, True, 0,
             'SteadyDBTestDB', user='SteadyDBTestUser')
         self.assert_(isinstance(db, SteadyDBConnection))
         self.assert_(hasattr(db, '_con'))
@@ -355,9 +378,9 @@ def test04_SteadyDBConnection(self):
             ['doit', 'commit', 'dont', 'rollback'])
 
     def test05_SteadyDBConnectionCreatorFunction(self):
-        db1 = SteadyDBconnect(dbapi, 0, None, None, 1,
+        db1 = SteadyDBconnect(dbapi, 0, None, None, True,
             'SteadyDBTestDB', user='SteadyDBTestUser')
-        db2 = SteadyDBconnect(connect, 0, None, None, 1,
+        db2 = SteadyDBconnect(connect, 0, None, None, True,
             'SteadyDBTestDB', user='SteadyDBTestUser')
         self.assertEqual(db1.dbapi(), db2.dbapi())
         self.assertEqual(db1.threadsafety(), db2.threadsafety())
@@ -520,6 +543,77 @@ def test10_SteadyDBConnectionSetSizes(self):
         result = cursor.fetchone()
         self.assertEqual(result, ([6, 42, 7], {None: 7, 3: 15, 9: 42}))
 
+    def test11_SteadyDBConnectionPing(self):
+        Connection.has_ping = False
+        Connection.num_pings = 0
+        db = SteadyDBconnect(dbapi)
+        db.cursor().execute('select test')
+        self.assertEqual(Connection.num_pings, 0)
+        db.close()
+        db.cursor().execute('select test')
+        self.assertEqual(Connection.num_pings, 0)
+        db = SteadyDBconnect(dbapi, ping=3)
+        db.cursor().execute('select test')
+        self.assertEqual(Connection.num_pings, 1)
+        db.close()
+        db.cursor().execute('select test')
+        self.assertEqual(Connection.num_pings, 1)
+        Connection.has_ping = True
+        db = SteadyDBconnect(dbapi)
+        db.cursor().execute('select test')
+        self.assertEqual(Connection.num_pings, 1)
+        db.close()
+        db.cursor().execute('select test')
+        self.assertEqual(Connection.num_pings, 1)
+        db = SteadyDBconnect(dbapi, ping=3)
+        db.cursor().execute('select test')
+        self.assertEqual(Connection.num_pings, 3)
+        db.close()
+        db.cursor().execute('select test')
+        self.assertEqual(Connection.num_pings, 5)
+        db = SteadyDBconnect(dbapi, ping=1)
+        self.assertEqual(Connection.num_pings, 5)
+        db.cursor()
+        self.assertEqual(Connection.num_pings, 6)
+        db.close()
+        cursor = db.cursor()
+        self.assertEqual(Connection.num_pings, 7)
+        cursor.execute('select test')
+        self.assertEqual(Connection.num_pings, 7)
+        db = SteadyDBconnect(dbapi, ping=2)
+        self.assertEqual(Connection.num_pings, 7)
+        db.cursor()
+        self.assertEqual(Connection.num_pings, 7)
+        db.close()
+        cursor = db.cursor()
+        self.assertEqual(Connection.num_pings, 7)
+        cursor.execute('select test')
+        self.assertEqual(Connection.num_pings, 8)
+        db.close()
+        cursor = db.cursor()
+        self.assertEqual(Connection.num_pings, 8)
+        cursor.execute('select test')
+        self.assertEqual(Connection.num_pings, 9)
+        db = SteadyDBconnect(dbapi, ping=3)
+        self.assertEqual(Connection.num_pings, 9)
+        db.cursor()
+        self.assertEqual(Connection.num_pings, 10)
+        db.close()
+        cursor = db.cursor()
+        self.assertEqual(Connection.num_pings, 11)
+        cursor.execute('select test')
+        self.assertEqual(Connection.num_pings, 12)
+        db.close()
+        cursor = db.cursor()
+        self.assertEqual(Connection.num_pings, 13)
+        cursor.execute('select test')
+        self.assertEqual(Connection.num_pings, 14)
+        db.close()
+        cursor.execute('select test')
+        self.assertEqual(Connection.num_pings, 16)
+        Connection.has_ping = False
+        Connection.num_pings = 0
+
 
 if __name__ == '__main__':
     unittest.main()