Support context managers (fixes #33)

Author: Cito

Diff for: dbutils/pooled_db.py

@@ -114,6 +114,13 @@
     cur.close()  # or del cur
     db.close()  # or del db
 
+You can also use context managers for simpler code:
+
+    with pool.connection() as db:
+        with db.cursor as cur:
+            cur.execute(...)
+            res = cur.fetchone()
+
 Note that you need to explicitly start transactions by calling the
 begin() method.  This ensures that the connection will not be shared
 with other threads, that the transparent reopening will be suspended
@@ -440,6 +447,14 @@ def __del__(self):
         except:  # builtin Exceptions might not exist any more
             pass
 
+    def __enter__(self):
+        """Enter a runtime context for the connection."""
+        return self
+
+    def __exit__(self, *exc):
+        """Exit a runtime context for the connection."""
+        self.close()
+
 
 class SharedDBConnection:
     """Auxiliary class for shared connections."""
@@ -526,3 +541,11 @@ def __del__(self):
             self.close()
         except:  # builtin Exceptions might not exist any more
             pass
+
+    def __enter__(self):
+        """Enter a runtime context for the connection."""
+        return self
+
+    def __exit__(self, *exc):
+        """Exit a runtime context for the connection."""
+        self.close()

Diff for: dbutils/pooled_pg.py

@@ -83,6 +83,11 @@
     res = db.query(...).getresult()
     db.close()  # or del db
 
+You can also a context manager for simpler code:
+
+    with pool.connection() as db:
+        res = db.query(...).getresult()
+
 Note that you need to explicitly start transactions by calling the
 begin() method.  This ensures that the transparent reopening will be
 suspended until the end of the transaction, and that the connection will
@@ -293,3 +298,11 @@ def __del__(self):
             self.close()
         except:  # builtin Exceptions might not exist any more
             pass
+
+    def __enter__(self):
+        """Enter a runtime context for the connection."""
+        return self
+
+    def __exit__(self, *exc):
+        """Exit a runtime context for the connection."""
+        self.close()

Diff for: docs/changelog.html

@@ -12,14 +12,22 @@
 <h1 class="title">Changelog for DBUtils</h1>
 
 <section id="id1">
+<h2>2.0.2</h2>
+<p>DBUtils 2.0.2 was released on ...</p>
+<p>Changes:</p>
+<ul class="simple">
+<li><p>Allow using context managers for pooled connections.</p></li>
+</ul>
+</section>
+<section id="id2">
 <h2>2.0.1</h2>
 <p>DBUtils 2.0.1 was released on April 8, 2021.</p>
 <p>Changes:</p>
 <ul class="simple">
-<li><p>Avoid &quot;name Exception is not defined&quot; when exiting</p></li>
+<li><p>Avoid &quot;name Exception is not defined&quot; when exiting.</p></li>
 </ul>
 </section>
-<section id="id2">
+<section id="id3">
 <h2>2.0</h2>
 <p>DBUtils 2.0 was released on September 26, 2020.</p>
 <p>It is intended to be used with Python versions 2.7 and 3.5 to 3.9.</p>
@@ -35,7 +43,7 @@ <h2>2.0</h2>
 <li><p>This changelog has been compiled from the former release notes.</p></li>
 </ul>
 </section>
-<section id="id3">
+<section id="id4">
 <h2>1.4</h2>
 <p>DBUtils 1.4 was released on September 26, 2020.</p>
 <p>It is intended to be used with Python versions 2.7 and 3.5 to 3.9.</p>
@@ -46,7 +54,7 @@ <h2>1.4</h2>
 inside a transaction.</p></li>
 </ul>
 </section>
-<section id="id4">
+<section id="id5">
 <h2>1.3</h2>
 <p>DBUtils 1.3 was released on March 3, 2018.</p>
 <p>It is intended to be used with Python versions 2.6, 2.7 and 3.4 to 3.7.</p>
@@ -55,12 +63,12 @@ <h2>1.3</h2>
 <li><p>This version now supports context handlers for connections and cursors.</p></li>
 </ul>
 </section>
-<section id="id5">
+<section id="id6">
 <h2>1.2</h2>
 <p>DBUtils 1.2 was released on February 5, 2017.</p>
 <p>It is intended to be used with Python versions 2.6, 2.7 and 3.0 to 3.6.</p>
 </section>
-<section id="id6">
+<section id="id7">
 <h2>1.1.1</h2>
 <p>DBUtils 1.1.1 was released on February 4, 2017.</p>
 <p>It is intended to be used with Python versions 2.3 to 2.7.</p>
@@ -74,7 +82,7 @@ <h2>1.1.1</h2>
 <li><p>Fixed a problem when running under Jython (reported by Vitaly Kruglikov).</p></li>
 </ul>
 </section>
-<section id="id7">
+<section id="id8">
 <h2>1.1</h2>
 <p>DBUtils 1.1 was released on August 14, 2011.</p>
 <p>Improvements:</p>
@@ -103,7 +111,7 @@ <h2>1.1</h2>
 <li><p>Fixed some minor issues with the <span class="docutils literal">DBUtilsExample</span> for Webware.</p></li>
 </ul>
 </section>
-<section id="id8">
+<section id="id9">
 <h2>1.0</h2>
 <p>DBUtils 1.0 was released on November 29, 2008.</p>
 <p>It is intended to be used with Python versions 2.2 to 2.6.</p>
@@ -136,7 +144,7 @@ <h2>1.0</h2>
 the MySQLdb module (problem reported by Gregory Pinero).</p></li>
 </ul>
 </section>
-<section id="id9">
+<section id="id10">
 <h2>0.9.4</h2>
 <p>DBUtils 0.9.4 was released on July 7, 2007.</p>
 <p>This release fixes a problem in the destructor code and has been supplemented
@@ -145,7 +153,7 @@ <h2>0.9.4</h2>
 in the last release, since you can now pass custom creator functions
 for database connections instead of DB-API 2 modules.</p>
 </section>
-<section id="id10">
+<section id="id11">
 <h2>0.9.3</h2>
 <p>DBUtils 0.9.3 was released on May 21, 2007.</p>
 <p>Changes:</p>
@@ -160,7 +168,7 @@ <h2>0.9.3</h2>
 Added Chinese translation of the User's Guide, kindly contributed by gashero.</p></li>
 </ul>
 </section>
-<section id="id11">
+<section id="id12">
 <h2>0.9.2</h2>
 <p>DBUtils 0.9.2 was released on September 22, 2006.</p>
 <p>It is intended to be used with Python versions 2.2 to 2.5.</p>
@@ -170,7 +178,7 @@ <h2>0.9.2</h2>
 storage engine. Accordingly, renamed <span class="docutils literal">SolidPg</span> to <span class="docutils literal">SteadyPg</span>.</p></li>
 </ul>
 </section>
-<section id="id12">
+<section id="id13">
 <h2>0.9.1</h2>
 <p>DBUtils 0.9.1 was released on May 8, 2006.</p>
 <p>It is intended to be used with Python versions 2.2 to 2.4.</p>
@@ -184,7 +192,7 @@ <h2>0.9.1</h2>
 <li><p>Improved the documentation and added a User's Guide.</p></li>
 </ul>
 </section>
-<section id="id13">
+<section id="id14">
 <h2>0.8.1 - 2005-09-13</h2>
 <p>DBUtils 0.8.1 was released on September 13, 2005.</p>
 <p>It is intended to be used with Python versions 2.0 to 2.4.</p>

Diff for: docs/changelog.rst

@@ -1,14 +1,24 @@
 Changelog for DBUtils
 +++++++++++++++++++++
 
+2.0.2
+=====
+
+DBUtils 2.0.2 was released on ...
+
+Changes:
+
+* Allow using context managers for pooled connections.
+
+
 2.0.1
 =====
 
 DBUtils 2.0.1 was released on April 8, 2021.
 
 Changes:
 
-* Avoid "name Exception is not defined" when exiting
+* Avoid "name Exception is not defined" when exiting.
 
 2.0
 ===

Diff for: docs/main.de.html

@@ -428,6 +428,11 @@ <h3>PooledDB (pooled_db)</h3>
 res = cur.fetchone()
 cur.close()  # oder del cur
 db.close()  # oder del db</pre>
+<p>Sie können dies auch durch Verwendung von Kontext-Managern vereinfachen:</p>
+<pre class="literal-block">with pool.connection() as db:
+    with db.cursor as cur:
+        cur.execute(...)
+        res = cur.fetchone()</pre>
 <p>Bitte beachten Sie, dass Transaktionen explizit durch Aufruf der Methode
 <span class="docutils literal">begin()</span> eingeleiten werden müssen. Hierdurch wird sichergestellt,
 dass die Verbindung nicht mehr mit anderen Threads geteilt wird, dass das

Diff for: docs/main.de.rst

@@ -443,6 +443,14 @@ sie gebraucht werden, etwa so::
   cur.close()  # oder del cur
   db.close()  # oder del db
 
+Sie können dies auch durch Verwendung von Kontext-Managern vereinfachen::
+
+  with pool.connection() as db:
+      with db.cursor as cur:
+          cur.execute(...)
+          res = cur.fetchone()
+
+
 Bitte beachten Sie, dass Transaktionen explizit durch Aufruf der Methode
 ``begin()`` eingeleiten werden müssen. Hierdurch wird sichergestellt,
 dass die Verbindung nicht mehr mit anderen Threads geteilt wird, dass das

Diff for: docs/main.html

@@ -390,6 +390,11 @@ <h3>PooledDB (pooled_db)</h3>
 res = cur.fetchone()
 cur.close()  # or del cur
 db.close()  # or del db</pre>
+<p>You can also use context managers for simpler code:</p>
+<pre class="literal-block">with pool.connection() as db:
+    with db.cursor as cur:
+        cur.execute(...)
+        res = cur.fetchone()</pre>
 <p>Note that you need to explicitly start transactions by calling the
 <span class="docutils literal">begin()</span> method. This ensures that the connection will not be shared
 with other threads, that the transparent reopening will be suspended

Diff for: docs/main.rst

@@ -405,6 +405,13 @@ object stays alive as long as you are using it, like that::
   cur.close()  # or del cur
   db.close()  # or del db
 
+You can also use context managers for simpler code::
+
+  with pool.connection() as db:
+      with db.cursor as cur:
+          cur.execute(...)
+          res = cur.fetchone()
+
 Note that you need to explicitly start transactions by calling the
 ``begin()`` method. This ensures that the connection will not be shared
 with other threads, that the transparent reopening will be suspended

Diff for: tests/test_persistent_db.py

@@ -286,6 +286,14 @@ def test_failed_transaction(self):
         db._con.close()
         cursor.execute('select test')
 
+    def test_context_manager(self):
+        persist = PersistentDB(dbapi)
+        with persist.connection() as db:
+            with db.cursor() as cursor:
+                cursor.execute('select test')
+                r = cursor.fetchone()
+            self.assertEqual(r, 'test')
+
 
 if __name__ == '__main__':
     unittest.main()

Diff for: tests/test_persistent_pg.py

@@ -189,6 +189,12 @@ def test_failed_transaction(self):
         db._con.close()
         self.assertEqual(db.query('select test'), 'test')
 
+    def test_context_manager(self):
+        persist = PersistentPg()
+        with persist.connection() as db:
+            db.query('select test')
+            self.assertEqual(db.num_queries, 1)
+
 
 if __name__ == '__main__':
     unittest.main()

Diff for: tests/test_pooled_db.py

@@ -1181,6 +1181,33 @@ def test_reset_transaction(self):
         self.assertFalse(con._transaction)
         self.assertEqual(con._con.session, ['rollback'])
 
+    def test_context_manager(self):
+        pool = PooledDB(dbapi, 1, 1, 1)
+        con = pool._idle_cache[0]._con
+        with pool.connection() as db:
+            self.assertTrue(hasattr(db, '_shared_con'))
+            self.assertFalse(pool._idle_cache)
+            self.assertTrue(con.valid)
+            with db.cursor() as cursor:
+                self.assertEqual(con.open_cursors, 1)
+                cursor.execute('select test')
+                r = cursor.fetchone()
+            self.assertEqual(con.open_cursors, 0)
+            self.assertEqual(r, 'test')
+            self.assertEqual(con.num_queries, 1)
+        self.assertTrue(pool._idle_cache)
+        with pool.dedicated_connection() as db:
+            self.assertFalse(hasattr(db, '_shared_con'))
+            self.assertFalse(pool._idle_cache)
+            with db.cursor() as cursor:
+                self.assertEqual(con.open_cursors, 1)
+                cursor.execute('select test')
+                r = cursor.fetchone()
+            self.assertEqual(con.open_cursors, 0)
+            self.assertEqual(r, 'test')
+            self.assertEqual(con.num_queries, 2)
+        self.assertTrue(pool._idle_cache)
+
 
 class TestSharedDBConnection(unittest.TestCase):
 

Diff for: tests/test_pooled_pg.py

@@ -14,7 +14,7 @@
 
 from . import mock_pg  # noqa
 
-from dbutils.pooled_pg import PooledPg, InvalidConnection
+from dbutils.pooled_pg import PooledPg, InvalidConnection, TooManyConnections
 
 
 class TestPooledPg(unittest.TestCase):
@@ -304,6 +304,19 @@ def test_reset_transaction(self):
         self.assertEqual(con.session, [])
         self.assertEqual(con.num_queries, 0)
 
+    def test_context_manager(self):
+        pool = PooledPg(1, 1, 1)
+        with pool.connection() as db:
+            db_con = db._con._con
+            db.query('select test')
+            self.assertEqual(db_con.num_queries, 1)
+            self.assertRaises(TooManyConnections, pool.connection)
+        with pool.connection() as db:
+            db_con = db._con._con
+            db.query('select test')
+            self.assertEqual(db_con.num_queries, 2)
+            self.assertRaises(TooManyConnections, pool.connection)
+
 
 if __name__ == '__main__':
     unittest.main()