Close cursor during __exit__; #18

Author: chadwhitacre

postgres.py

@@ -505,11 +505,13 @@ def __enter__(self):
         """Get a connection from the pool.
         """
         self.conn = self.pool.getconn()
-        return self.conn.cursor(*self.a, **self.kw)
+        self.cursor = self.conn.cursor(*self.a, **self.kw)
+        return self.cursor
 
     def __exit__(self, *exc_info):
         """Put our connection back in the pool.
         """
+        self.cursor.close()
         self.pool.putconn(self.conn)
 
 
@@ -540,7 +542,8 @@ def __enter__(self):
         """
         self.conn = self.pool.getconn()
         self.conn.autocommit = False
-        return self.conn.cursor(*self.a, **self.kw)
+        self.cursor = self.conn.cursor(*self.a, **self.kw)
+        return self.cursor
 
     def __exit__(self, *exc_info):
         """Put our connection back in the pool.
@@ -549,6 +552,7 @@ def __exit__(self, *exc_info):
             self.conn.commit()
         else:
             self.conn.rollback()
+        self.cursor.close()
         self.conn.autocommit = True
         self.pool.putconn(self.conn)
 

tests.py

@@ -6,6 +6,7 @@
 
 from postgres import Postgres, TooFew, TooMany
 from psycopg2.extras import NamedTupleCursor
+from psycopg2 import InterfaceError
 
 
 DATABASE_URL = os.environ['DATABASE_URL']
@@ -233,6 +234,13 @@ def test_we_can_use_cursor_closed(self):
             actual = cursor.closed
         assert not actual
 
+    def test_we_close_the_cursor(self):
+        with self.db.get_cursor() as cursor:
+            cursor.execute("SELECT * FROM foo ORDER BY bar")
+        self.assertRaises( InterfaceError
+                         , cursor.fetchall
+                          )
+
 
 # db.get_transaction
 # ==================
@@ -272,6 +280,13 @@ class Heck(Exception): pass
         actual = self.db.all("SELECT * FROM foo ORDER BY bar")
         assert actual == [{"bar": "baz"}, {"bar": "buz"}]
 
+    def test_we_close_the_cursor(self):
+        with self.db.get_transaction() as txn:
+            txn.execute("SELECT * FROM foo ORDER BY bar")
+        self.assertRaises( InterfaceError
+                         , txn.fetchall
+                          )
+
 
 # db.get_connection
 # =================