mirror of
https://github.com/kennethreitz/django-postgrespool.git
synced 2026-06-05 23:00:19 +00:00
Kenneth Reitz
63 lines
2.6 KiB
Python
63 lines
2.6 KiB
Python
# -*- coding: utf-8 -*-
|
|
|
|
import sqlalchemy.pool as pool
|
|
|
|
from django.db.backends.postgresql_psycopg2.base import *
|
|
from django.db.backends.postgresql_psycopg2.base import DatabaseWrapper as Psycopg2DatabaseWrapper
|
|
|
|
Database_ = pool.manage(Database)
|
|
|
|
|
|
class DatabaseWrapper(Psycopg2DatabaseWrapper):
|
|
"""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'] = 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)
|