| 1 | # Custom DB backend postgresql_psycopg2 based
|
|---|
| 2 | # implements persistent database connection using global variable
|
|---|
| 3 |
|
|---|
| 4 | from django.db.backends.postgresql_psycopg2.base import DatabaseError, DatabaseWrapper as BaseDatabaseWrapper, \
|
|---|
| 5 | IntegrityError
|
|---|
| 6 | from psycopg2 import OperationalError
|
|---|
| 7 |
|
|---|
| 8 | connection = None
|
|---|
| 9 |
|
|---|
| 10 | class DatabaseWrapper(BaseDatabaseWrapper):
|
|---|
| 11 | def _cursor(self, *args, **kwargs):
|
|---|
| 12 | global connection
|
|---|
| 13 | if connection is not None and self.connection is None:
|
|---|
| 14 | try: # Check if connection is alive
|
|---|
| 15 | connection.cursor().execute('SELECT 1')
|
|---|
| 16 | except OperationalError: # The connection is not working, need reconnect
|
|---|
| 17 | connection = None
|
|---|
| 18 | else:
|
|---|
| 19 | self.connection = connection
|
|---|
| 20 | cursor = super(DatabaseWrapper, self)._cursor(*args, **kwargs)
|
|---|
| 21 | if connection is None and self.connection is not None:
|
|---|
| 22 | connection = self.connection
|
|---|
| 23 | return cursor
|
|---|
| 24 |
|
|---|
| 25 | def close(self):
|
|---|
| 26 | if self.connection is not None:
|
|---|
| 27 | self.connection.commit()
|
|---|
| 28 | self.connection = None
|
|---|