This is a modification of http://djangosnippets.org/snippets/1707/ that handles the database going down or PG Bouncer killing the connection. This also works in things like Twisted to make sure the connection is alive before doing a real query. Thanks @mike_tk for the original post!
EDIT: Updated the wrapper to handle multi-db. Before it was using the first connection it made, now it creates an attribute name for the connection based on the name of the database.
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 | # Custom DB backend postgresql_psycopg2 based
# implements persistent database connection using thread local storage
from threading import local
from django.db.backends.postgresql_psycopg2.base import DatabaseError, \
DatabaseWrapper as BaseDatabaseWrapper, IntegrityError
from psycopg2 import OperationalError
threadlocal = local()
class DatabaseWrapper(BaseDatabaseWrapper):
def is_alive(self, connection):
try: # Check if connection is alive
connection.cursor().execute('SELECT 1')
except (DatabaseError, OperationalError):
return False
else:
return True
@property
def local_key(self):
return '{}_{}'.format(self.settings_dict.get('NAME', 'db'), 'connection')
@property
def local_connection(self):
return getattr(threadlocal, self.local_key, None)
@local_connection.setter
def local_connection(self, connection):
setattr(threadlocal, self.local_key, connection)
def _cursor(self, *args, **kwargs):
if self.local_connection is not None and self.connection is None:
if self.is_alive(self.local_connection):
self.connection = self.local_connection
else:
self.local_connection = None
elif self.connection is not None:
if self.is_alive(self.connection):
self.local_connection = self.connection
else:
self.connection.close()
self.connection = None
self.local_connection = None
cursor = super(DatabaseWrapper, self)._cursor(*args, **kwargs)
if self.local_connection is None and self.connection is not None:
self.local_connection = self.connection
return cursor
def close(self):
if self.connection is not None:
self.connection.commit()
self.connection = None
|
More like this
- Template tag - list punctuation for a list of items by shapiromatron 10 months, 3 weeks ago
- JSONRequestMiddleware adds a .json() method to your HttpRequests by cdcarter 11 months ago
- Serializer factory with Django Rest Framework by julio 1 year, 5 months ago
- Image compression before saving the new model / work with JPG, PNG by Schleidens 1 year, 6 months ago
- Help text hyperlinks by sa2812 1 year, 7 months ago
Comments
Please login first before commenting.