Initial public release

.gitignore

@@ -0,0 +1,6 @@
+.idea
+*.pyc
+build
+dist
+*.egg-info
+atlassian-ide-plugin.xml

LICENSE

@@ -0,0 +1,25 @@
+Copyright (c) 2012 MeetMe, Inc.
+All rights reserved.
+
+Redistribution and use in source and binary forms, with or without modification,
+are permitted provided that the following conditions are met:
+
+ * Redistributions of source code must retain the above copyright notice, this
+   list of conditions and the following disclaimer.
+ * Redistributions in binary form must reproduce the above copyright notice,
+   this list of conditions and the following disclaimer in the documentation
+   and/or other materials provided with the distribution.
+ * Neither the name of the MeetMe, Inc. nor the names of its
+   contributors may be used to endorse or promote products derived from this
+   software without specific prior written permission.
+
+THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
+ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
+WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED.
+IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT,
+INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING,
+BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
+DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF
+LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE
+OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF
+ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

README.md

@@ -0,0 +1,35 @@
+pgsql_wrapper
+=============
+An opinionated wrapper for interfacing with PostgreSQL that offers caching of
+connections and support for PyPy via psycopg2ct. By default the PgSQL class
+sets the cursor type to extras.DictCursor, and turns on both Unicode and UUID
+support. In addition, isolation level is set to auto-commit.
+
+As a convenience tool, pgsql_wrapper reduces the steps required in connecting to
+and setting up the connections and cursors required to talk to PostgreSQL.
+
+Without requiring any additional code, the module level caching of connections
+allows for multiple modules in the same interpreter to use the same PostgreSQL
+connection.
+
+Requirements
+------------
+
+ - psycopg2
+ - psycopg2ct (for PyPy support)
+
+Example
+-------
+
+    import pgsql_wrapper
+
+    HOST = 'localhost'
+    PORT = 5432
+    DBNAME = 'production'
+    USER = 'www'
+
+
+    connection = pgsql_wrapper.PgSQL(HOST, PORT, DBNAME, USER)
+    connection.cursor.execute('SELECT 1 as value')
+    data = connection.cursor.fetchone()
+    print data['value']

pgsql_wrapper.py

@@ -0,0 +1,269 @@
+"""PostgreSQL Class module
+
+An opinionated wrapper for interfacing with PostgreSQL that offers caching of
+connections and support for PyPy via psycopg2ct. By default the PgSQL class
+sets the cursor type to extras.DictCursor, and turns on both Unicode and UUID
+support.
+
+"""
+__version__ = '1.1.1'
+
+import hashlib
+import logging
+import platform
+import time
+
+# Import PyPy compatibility
+target = platform.python_implementation()
+if target == 'PyPy':
+    from psycopg2ct import compat
+    compat.register()
+    PYPY = True
+else:
+    PYPY = False
+
+# Import psycopg2 and it's extensions and extras
+import psycopg2
+from psycopg2 import extensions
+from psycopg2 import extras
+
+# Expose exceptions so clients do not need to import psycopg2 too
+from psycopg2 import DataError
+from psycopg2 import DatabaseError
+from psycopg2 import IntegrityError
+from psycopg2 import InterfaceError
+from psycopg2 import InternalError
+from psycopg2 import NotSupportedError
+from psycopg2 import OperationalError
+from psycopg2 import ProgrammingError
+from psycopg2.extensions import QueryCanceledError
+from psycopg2.extensions import TransactionRollbackError
+
+LOGGER = logging.getLogger(__name__)
+
+# Empty connection TTL
+CACHE_TTL = 60
+
+CONNECTIONS = dict()
+
+
+def _add_cached_connection(hash_value, connection):
+    """Add the connection to our module level connection dictionary
+
+    :param str hash_value: Hash generated by _generate_connection
+    :param psycopg2._psycopg.connection connection: PostgreSQL connection
+
+    """
+    global CONNECTIONS
+    if hash_value not in CONNECTIONS:
+        CONNECTIONS[hash_value] = {'clients': 1,
+                                   'handle': connection,
+                                   'last_client': 0}
+        LOGGER.info('Appended connection %s to module pool', hash_value)
+    else:
+        LOGGER.critical('Connection is already assigned: %s', hash_value)
+
+def _check_for_unused_expired_connections():
+
+    global CONNECTIONS
+    for hash_value in CONNECTIONS.keys():
+        if (not CONNECTIONS[hash_value]['clients'] and
+            (time.time() > CONNECTIONS[hash_value]['last_client'] + CACHE_TTL)):
+            LOGGER.info('Removing expired connection: %s', hash_value)
+            del CONNECTIONS[hash_value]
+
+def _generate_connection_hash(dsn):
+    """Generates a connection hash for the given parameters.
+
+    :param str dsn: DSN for connection
+    :rtype: str
+
+    """
+    return str(hashlib.sha1(dsn).hexdigest())
+
+
+def _get_cached_connection(hash_value):
+    """Check our global connection stack to see if we already have a
+    connection with the same exact connection parameters and use it if so.
+
+    :param str hash_value: Hash generated by _generate_connection_hash
+    :rtype: psycopg2._psycopg.connection or None
+
+    """
+    _check_for_unused_expired_connections()
+    if hash_value in CONNECTIONS:
+        LOGGER.debug('Returning cached connection and incrementing counter')
+        CONNECTIONS[hash_value]['clients'] += 1
+        return CONNECTIONS[hash_value]['handle']
+    return None
+
+
+def _free_cached_connection(hash_value):
+    """Decrement our use counter for the hash and if it is the only one, delete
+    the cached connection.
+
+    :param str hash_value: Hash generated by _generate_connection_hash
+
+    """
+    global CONNECTIONS
+    if hash_value in CONNECTIONS:
+        LOGGER.debug('Decrementing client count for: %s', hash_value)
+        CONNECTIONS[hash_value]['clients'] -= 1
+        if not CONNECTIONS[hash_value]['clients']:
+            LOGGER.debug('Noting last client time for %s', hash_value)
+            CONNECTIONS[hash_value]['last_client'] = time.time()
+
+
+class PgSQL(object):
+    """PostgreSQL connection object.
+
+    Uses a module level cache of connections to reduce overhead.
+
+    """
+    def __init__(self,
+                 host='localhost',
+                 port=5432,
+                 dbname=None,
+                 user='www',
+                 password=None,
+                 cursor_factory=extras.DictCursor):
+        """Connect to a PostgreSQL server using the module wide connection and
+        set the isolation level.
+
+        :param str host: PostgreSQL Host
+        :param port: PostgreSQL port
+        :param str dbname: Dabase name
+        :param str user: PostgreSQL user name
+        :param str password: User's password
+        :param psycopg2.cursor: The cursor type to use
+
+        """
+        dsn = self._get_dsn(host, port, dbname, user, password)
+        self._connection = self._get_connection(dsn)
+        self._cursor = self._get_cursor(cursor_factory)
+        self._register_unicode_cursor(self._cursor)
+        self._register_uuid()
+
+    @property
+    def connection(self):
+        """Returns the psycopg2 postgresql connection instance"""
+        return self._connection
+
+    @property
+    def cursor(self):
+        """Returns the cursor instance
+
+        :rtype: psycopg2._psycopg.cursor
+
+        """
+        return self._cursor
+
+    def __del__(self):
+        """Remove our connection from the stack when we're done."""
+        if self._connection_hash:
+            _free_cached_connection(self._connection_hash)
+
+    def _add_connection_to_cache(self, connection):
+        """Add the connection to the module wide connection pool.
+
+        :param psycopg2.connection: The new connection
+
+        """
+        _add_cached_connection(self._connection_hash, connection)
+
+    def _connect(self, dsn):
+        """Connect to PostgreSQL using the DSN.
+
+        :param str dsn: The connection dsn
+        :return: psycopg2.connection
+
+        """
+        LOGGER.info('Connecting to %s', dsn)
+        return psycopg2.connect(dsn)
+
+    def _get_connection(self, dsn):
+        """Return a connection, cached or otherwise for the given DSN.
+
+        :param str dsn: The connection DSN
+        :rtype: psycopg2.connection
+
+        """
+        # Generate a connection hash for module level instance of connection
+        self._connection_hash = _generate_connection_hash(dsn)
+
+        LOGGER.debug('Connection hash: %s', self._connection_hash)
+
+        # Attempt to get a cached connection from our module level pool
+        connection = _get_cached_connection(self._connection_hash)
+
+        # If we got a result, just log our success in doing so
+        if connection:
+            LOGGER.debug("Reusing cached connection: %s", self._connection_hash)
+            return connection
+
+        # Create a new PostgreSQL connection and cache it
+        connection = self._connect(dsn)
+
+        # Cache the connection
+        self._add_connection_to_cache(connection)
+
+        # Added in because psycopg2ct connects and leaves the connection in
+        # a weird state: consts.STATUS_DATESTYLE, returning from
+        # Connection._setup without setting the state as const.STATUS_OK
+        if PYPY:
+            connection.reset()
+
+        self._set_isolation_level(connection)
+
+        # Return the connection
+        return connection
+
+    def _get_cursor(self, cursor_factory):
+        """Return a cursor for the given cursor_factory.
+
+        :param psycopg2.cursor: The cursor type to use
+        :rtype: psycopg2.extensions.cursor
+
+        """
+        return self._connection.cursor(cursor_factory=cursor_factory)
+
+    def _register_unicode_cursor(self, cursor):
+        """Register the cursor to be able to receive Unicode string.
+
+        :param psycopg2.cursor: The cursor to add unicode support to
+
+        """
+        psycopg2.extensions.register_type(psycopg2.extensions.UNICODE, cursor)
+
+    def _get_dsn(self, host, port, dbname, user, password):
+        """Create a DSN for the specified attributes.
+
+        :param str host: The postgresql host
+        :param int port: The postgresql port
+        :param str dbname: The database name
+        :param str user: The postgresql user
+        :param str password: The password to use, None for no password
+        :return str: The DSN to connect
+
+        """
+        dsn = "host='%s' port=%i dbname='%s' user='%s'" % (host,
+                                                           port,
+                                                           dbname,
+                                                           user)
+        # Add the password if specified
+        if password:
+            dsn += " password='%s'" % password
+
+        return dsn
+
+    def _register_uuid(self):
+        """Register the UUID extension from psycopg2"""
+        psycopg2.extras.register_uuid(self._connection)
+
+    def _set_isolation_level(self, connection):
+        """Set the isolation level automatically to commit after every query
+
+        :param psycopg2.connection connection: Connection to set level on
+
+        """
+        connection.set_isolation_level(extensions.ISOLATION_LEVEL_AUTOCOMMIT)

setup.py

@@ -0,0 +1,19 @@
+from setuptools import setup
+import platform
+
+# Make the install_requires
+target = platform.python_implementation()
+if target == 'PyPy':
+    install_requires = ['psycopg2ct']
+else:
+    install_requires = ['psycopg2']
+
+setup(name='pgsql_wrapper',
+      version='1.1.1',
+      description="PostgreSQL / psycopg2 caching wrapper class",
+      maintainer="Gavin M. Roy",
+      maintainer_email="gmr@meetme.com",
+      url="http://github.com/MeetMe/pgsql_wrapper",
+      install_requires=install_requires,
+      py_modules=['pgsql_wrapper'],
+      zip_safe=True)