From 0fd7bd0e4a4c72b34d43455452a075cd5a828df4 Mon Sep 17 00:00:00 2001
From: Kenneth Reitz <me@kennethreitz.com>
Date: Thu, 27 Sep 2012 20:02:34 -0400
Subject: [PATCH] DATABASE_POOL_ARGS

Closes #1
---
 README.rst                  | 11 +++++++++++
 django_postgrespool/base.py | 26 +++++++++++++++++++++++---
 2 files changed, 34 insertions(+), 3 deletions(-)

diff --git a/README.rst b/README.rst
index 9e1f50b..ee096aa 100644
--- a/README.rst
+++ b/README.rst
@@ -36,3 +36,14 @@ Installation
 Installing Django-PostgresPool is simple, with pip::
 
     $ pip install django-postgrespool
+
+Configuration
+-------------
+
+Optionally, you can provide additional options to pass to SQLAlchemy's pool creation::
+
+    DATABASE_POOL_ARGS = {
+        'max_overflow': 10,
+        'pool_size': 5,
+        'recycle': 300
+    }
\ No newline at end of file
diff --git a/django_postgrespool/base.py b/django_postgrespool/base.py
index 592ef58..6f11606 100644
--- a/django_postgrespool/base.py
+++ b/django_postgrespool/base.py
@@ -1,11 +1,31 @@
 # -*- coding: utf-8 -*-
 
-import sqlalchemy.pool as pool
+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
 
-Database_ = pool.manage(Database)
+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):
@@ -36,7 +56,7 @@ class DatabaseWrapper(Psycopg2DatabaseWrapper):
                 conn_params['host'] = settings_dict['HOST']
             if settings_dict['PORT']:
                 conn_params['port'] = settings_dict['PORT']
-            self.connection = Database_.connect(**conn_params)
+            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: