diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..8138303 --- /dev/null +++ b/.gitignore @@ -0,0 +1,6 @@ +.idea +*.pyc +build +dist +*.egg-info +atlassian-ide-plugin.xml diff --git a/LICENSE b/LICENSE new file mode 100644 index 0000000..ce400a1 --- /dev/null +++ b/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. diff --git a/README.md b/README.md new file mode 100644 index 0000000..163648c --- /dev/null +++ b/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'] diff --git a/pgsql_wrapper.py b/pgsql_wrapper.py new file mode 100644 index 0000000..9db9d8f --- /dev/null +++ b/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) diff --git a/setup.py b/setup.py new file mode 100644 index 0000000..a97f930 --- /dev/null +++ b/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)