mirror of
https://github.com/kennethreitz/django-postgrespool.git
synced 2026-06-05 23:00:19 +00:00
Kenneth Reitz
84 lines
3.4 KiB
Python
84 lines
3.4 KiB
Python
# -*- coding: utf-8 -*-
|
|
|
|
import logging
|
|
from functools import partial
|
|
|
|
from sqlalchemy import event
|
|
from sqlalchemy.pool import manage, QueuePool
|
|
|
|
from django.conf import settings
|
|
from django.db.backends.postgresql_psycopg2.base import *
|
|
from django.db.backends.postgresql_psycopg2.base import DatabaseWrapper as Psycopg2DatabaseWrapper
|
|
|
|
POOL_SETTINGS = 'DATABASE_POOL_ARGS'
|
|
|
|
# DATABASE_POOL_ARGS should be something like:
|
|
# {'max_overflow':10, 'pool_size':5, 'recycle':300}
|
|
pool_args = getattr(settings, POOL_SETTINGS, {})
|
|
db_pool = manage(Database, **pool_args)
|
|
|
|
log = logging.getLogger('z.pool')
|
|
|
|
def _log(message, *args):
|
|
log.debug('%s to %s' % (message, args[0].get_host_info()))
|
|
|
|
# Only hook up the listeners if we are in debug mode.
|
|
if settings.DEBUG:
|
|
event.listen(QueuePool, 'checkout', partial(_log, 'retrieved from pool'))
|
|
event.listen(QueuePool, 'checkin', partial(_log, 'returned to pool'))
|
|
event.listen(QueuePool, 'connect', partial(_log, 'new connection'))
|
|
|
|
|
|
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 = db_pool.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)
|