From 9b2d75fc8d6bddeb38af3b7ca5fcd8eae9f4cf9b Mon Sep 17 00:00:00 2001
From: Chad Whitacre <chad@zetaweb.com>
Date: Wed, 7 Aug 2013 12:02:44 -0400
Subject: [PATCH] Remove Postgres prefix from classes; #9

---
 postgres.py | 79 +++++++++++++++++++++++++++++------------------------
 1 file changed, 44 insertions(+), 35 deletions(-)

diff --git a/postgres.py b/postgres.py
index 2c9e2e5..ca740ca 100644
--- a/postgres.py
+++ b/postgres.py
@@ -116,6 +116,9 @@
 from psycopg2.pool import ThreadedConnectionPool as ConnectionPool
 
 
+# A Helper
+# ========
+
 # Teach urlparse about postgres:// URLs.
 if 'postgres' not in urlparse.uses_netloc:
     urlparse.uses_netloc.append('postgres')
@@ -137,6 +140,9 @@ def url_to_dsn(url):
     return dsn
 
 
+# The Main Event
+# ==============
+
 class Postgres(object):
     """Interact with a `PostgreSQL <http://www.postgresql.org/>`_ datastore.
 
@@ -164,7 +170,7 @@ def __init__(self, url, minconn=1, maxconn=10):
         self.pool = ConnectionPool( minconn=minconn
                                   , maxconn=maxconn
                                   , dsn=dsn
-                                  , connection_factory=PostgresConnection
+                                  , connection_factory=Connection
                                    )
 
     def execute(self, *a, **kw):
@@ -194,7 +200,7 @@ def get_cursor(self, *a, **kw):
         upon both successful and exceptional executions against the cursor.
 
         """
-        return PostgresCursorContextManager(self.pool, *a, **kw)
+        return CursorContextManager(self.pool, *a, **kw)
 
     def get_transaction(self, *a, **kw):
         """Return a context manager wrapping a transactional cursor.
@@ -206,10 +212,10 @@ def get_transaction(self, *a, **kw):
         fine-grained  control over the transaction.
 
         """
-        return PostgresTransactionContextManager(self.pool, *a, **kw)
+        return TransactionContextManager(self.pool, *a, **kw)
 
     def get_connection(self):
-        """Return a context manager wrapping a PostgresConnection.
+        """Return a context manager wrapping a :py:class:`postgres.Connection`.
 
         This manager turns autocommit off, and back on when you're done with
         the connection. The connection is rolled back on exit, so be sure to
@@ -217,10 +223,10 @@ def get_connection(self):
         full fine-grained transaction control.
 
         """
-        return PostgresConnectionContextManager(self.pool)
+        return ConnectionContextManager(self.pool)
 
 
-class PostgresConnection(psycopg2.extensions.connection):
+class Connection(psycopg2.extensions.connection):
     """This is a subclass of psycopg2.extensions.connection.
 
     Changes:
@@ -246,81 +252,84 @@ def cursor(self, *a, **kw):
         return psycopg2.extensions.connection.cursor(self, *a, **kw)
 
 
-class PostgresTransactionContextManager(object):
-    """Instantiated once per db.get_transaction call.
+# Context Managers
+# ================
+
+class CursorContextManager(object):
+    """Instantiated once per cursor-level db access.
     """
 
     def __init__(self, pool, *a, **kw):
         self.pool = pool
+        self.a = a
+        self.kw = kw
         self.conn = None
 
-    def __enter__(self, *a, **kw):
+    def __enter__(self):
         """Get a connection from the pool.
         """
         self.conn = self.pool.getconn()
-        self.conn.autocommit = False
-        return self.conn.cursor(*a, **kw)
+        cursor = self.conn.cursor()
+        try:
+            cursor.execute(*self.a, **self.kw)
+        except:
+            # If we get an exception from execute (like, the query fails:
+            # pretty common), then the __exit__ clause is not triggered. We
+            # trigger it ourselves to avoid draining the pool.
+            self.__exit__()
+            raise
+        return cursor
 
     def __exit__(self, *exc_info):
         """Put our connection back in the pool.
         """
-        if exc_info == (None, None, None):
-            self.conn.commit()
-        else:
-            self.conn.rollback()
-        self.conn.autocommit = True
         self.pool.putconn(self.conn)
 
 
-class PostgresConnectionContextManager(object):
-    """Instantiated once per db.get_connection call.
+class TransactionContextManager(object):
+    """Instantiated once per db.get_transaction call.
     """
 
     def __init__(self, pool, *a, **kw):
         self.pool = pool
         self.conn = None
 
-    def __enter__(self):
+    def __enter__(self, *a, **kw):
         """Get a connection from the pool.
         """
         self.conn = self.pool.getconn()
         self.conn.autocommit = False
-        return self.conn
+        return self.conn.cursor(*a, **kw)
 
     def __exit__(self, *exc_info):
         """Put our connection back in the pool.
         """
-        self.conn.rollback()
+        if exc_info == (None, None, None):
+            self.conn.commit()
+        else:
+            self.conn.rollback()
         self.conn.autocommit = True
         self.pool.putconn(self.conn)
 
 
-class PostgresCursorContextManager(object):
-    """Instantiated once per cursor-level db access.
+class ConnectionContextManager(object):
+    """Instantiated once per db.get_connection call.
     """
 
     def __init__(self, pool, *a, **kw):
         self.pool = pool
-        self.a = a
-        self.kw = kw
         self.conn = None
 
     def __enter__(self):
         """Get a connection from the pool.
         """
         self.conn = self.pool.getconn()
-        cursor = self.conn.cursor()
-        try:
-            cursor.execute(*self.a, **self.kw)
-        except:
-            # If we get an exception from execute (like, the query fails:
-            # pretty common), then the __exit__ clause is not triggered. We
-            # trigger it ourselves to avoid draining the pool.
-            self.__exit__()
-            raise
-        return cursor
+        self.conn.autocommit = False
+        return self.conn
 
     def __exit__(self, *exc_info):
         """Put our connection back in the pool.
         """
+        self.conn.rollback()
+        self.conn.autocommit = True
         self.pool.putconn(self.conn)