diff --git a/django_postgrespool.py b/django_postgrespool.py new file mode 100644 index 0000000..c9caa49 --- /dev/null +++ b/django_postgrespool.py @@ -0,0 +1,79 @@ +# -*- coding: utf-8 -*- + +""" + + + +""" + +from django.db.backends import * +from django.db.backends.signals import connection_created + +from django.utils.encoding import force_str +try: + # import psycopg2 as Database + import psycopg2.extensions +except ImportError as e: + from django.core.exceptions import ImproperlyConfigured + raise ImproperlyConfigured("Error loading psycopg2 module: %s" % e) + +from django.db.backends.postgresql_psycopg2.base import Database +from django.db.backends.postgresql_psycopg2.base import DatabaseWrapper as DatabaseWrapper_ + + +import sqlalchemy.pool as pool +Database_ = pool.manage(Database) + + +class DatabaseWrapper(DatabaseWrapper_): + """SQLAlchemy FTW.""" + + def __init__(self, *args, **kwargs): + super(DatabaseWrapper, self).__init__(*args, **kwargs) + + def _cursor(self): + settings_dict = self.settings_dict + if self.connection is None: + if not settings_dict['NAME']: + from django.core.exceptions import ImproperlyConfigured + raise ImproperlyConfigured( + "settings.DATABASES is improperly configured. " + "Please supply the NAME value.") + conn_params = { + 'database': settings_dict['NAME'], + } + conn_params.update(settings_dict['OPTIONS']) + if 'autocommit' in conn_params: + del conn_params['autocommit'] + if settings_dict['USER']: + conn_params['user'] = settings_dict['USER'] + if settings_dict['PASSWORD']: + conn_params['password'] = force_str(settings_dict['PASSWORD']) + if settings_dict['HOST']: + conn_params['host'] = settings_dict['HOST'] + if settings_dict['PORT']: + conn_params['port'] = settings_dict['PORT'] + self.connection = Database_.connect(**conn_params) + self.connection.set_client_encoding('UTF8') + tz = 'UTC' if settings.USE_TZ else settings_dict.get('TIME_ZONE') + if tz: + try: + get_parameter_status = self.connection.get_parameter_status + except AttributeError: + # psycopg2 < 2.0.12 doesn't have get_parameter_status + conn_tz = None + else: + conn_tz = get_parameter_status('TimeZone') + + if conn_tz != tz: + # Set the time zone in autocommit mode (see #17062) + self.connection.set_isolation_level( + psycopg2.extensions.ISOLATION_LEVEL_AUTOCOMMIT) + self.connection.cursor().execute( + self.ops.set_time_zone_sql(), [tz]) + self.connection.set_isolation_level(self.isolation_level) + self._get_pg_version() + connection_created.send(sender=self.__class__, connection=self) + cursor = self.connection.cursor() + cursor.tzinfo_factory = utc_tzinfo_factory if settings.USE_TZ else None + return CursorWrapper(cursor) diff --git a/setup.py b/setup.py index f4f367e..08d0b27 100644 --- a/setup.py +++ b/setup.py @@ -25,9 +25,9 @@ required = [ setup( name='django-postgrespool', - version='0.0.1', + version='0.1.0', description='Postgres Connection Pooling for Django.', - # long_description=open('README.rst').read(), + long_description=open('README.rst').read(), author='Kenneth Reitz', author_email='me@kennethreitz.com', url='https://github.com/kennethreitz/django-postgrespool',