python/lib/django_1_4/django/db/backends/postgresql_psycopg2/base.py

   1 """
   2 PostgreSQL database backend for Django.
   3
   4 Requires psycopg 2: http://initd.org/projects/psycopg2
   5 """
   6 import sys
   7
   8 from django.db import utils
   9 from django.db.backends import *
  10 from django.db.backends.signals import connection_created
  11 from django.db.backends.postgresql_psycopg2.operations import DatabaseOperations
  12 from django.db.backends.postgresql_psycopg2.client import DatabaseClient
  13 from django.db.backends.postgresql_psycopg2.creation import DatabaseCreation
  14 from django.db.backends.postgresql_psycopg2.version import get_version
  15 from django.db.backends.postgresql_psycopg2.introspection import DatabaseIntrospection
  16 from django.utils.log import getLogger
  17 from django.utils.safestring import SafeUnicode, SafeString
  18 from django.utils.timezone import utc
  19
  20 try:
  21     import psycopg2 as Database
  22     import psycopg2.extensions
  23 except ImportError, e:
  24     from django.core.exceptions import ImproperlyConfigured
  25     raise ImproperlyConfigured("Error loading psycopg2 module: %s" % e)
  26
  27 DatabaseError = Database.DatabaseError
  28 IntegrityError = Database.IntegrityError
  29
  30 psycopg2.extensions.register_type(psycopg2.extensions.UNICODE)
  31 psycopg2.extensions.register_adapter(SafeString, psycopg2.extensions.QuotedString)
  32 psycopg2.extensions.register_adapter(SafeUnicode, psycopg2.extensions.QuotedString)
  33
  34 logger = getLogger('django.db.backends')
  35
  36 def utc_tzinfo_factory(offset):
  37     if offset != 0:
  38         raise AssertionError("database connection isn't set to UTC")
  39     return utc
  40
  41 class CursorWrapper(object):
  42     """
  43     A thin wrapper around psycopg2's normal cursor class so that we can catch
  44     particular exception instances and reraise them with the right types.
  45     """
  46
  47     def __init__(self, cursor):
  48         self.cursor = cursor
  49
  50     def execute(self, query, args=None):
  51         try:
  52             return self.cursor.execute(query, args)
  53         except Database.IntegrityError, e:
  54             raise utils.IntegrityError, utils.IntegrityError(*tuple(e)), sys.exc_info()[2]
  55         except Database.DatabaseError, e:
  56             raise utils.DatabaseError, utils.DatabaseError(*tuple(e)), sys.exc_info()[2]
  57
  58     def executemany(self, query, args):
  59         try:
  60             return self.cursor.executemany(query, args)
  61         except Database.IntegrityError, e:
  62             raise utils.IntegrityError, utils.IntegrityError(*tuple(e)), sys.exc_info()[2]
  63         except Database.DatabaseError, e:
  64             raise utils.DatabaseError, utils.DatabaseError(*tuple(e)), sys.exc_info()[2]
  65
  66     def __getattr__(self, attr):
  67         if attr in self.__dict__:
  68             return self.__dict__[attr]
  69         else:
  70             return getattr(self.cursor, attr)
  71
  72     def __iter__(self):
  73         return iter(self.cursor)
  74
  75 class DatabaseFeatures(BaseDatabaseFeatures):
  76     needs_datetime_string_cast = False
  77     can_return_id_from_insert = True
  78     requires_rollback_on_dirty_transaction = True
  79     has_real_datatype = True
  80     can_defer_constraint_checks = True
  81     has_select_for_update = True
  82     has_select_for_update_nowait = True
  83     has_bulk_insert = True
  84     supports_tablespaces = True
  85     can_distinct_on_fields = True
  86
  87 class DatabaseWrapper(BaseDatabaseWrapper):
  88     vendor = 'postgresql'
  89     operators = {
  90         'exact': '= %s',
  91         'iexact': '= UPPER(%s)',
  92         'contains': 'LIKE %s',
  93         'icontains': 'LIKE UPPER(%s)',
  94         'regex': '~ %s',
  95         'iregex': '~* %s',
  96         'gt': '> %s',
  97         'gte': '>= %s',
  98         'lt': '< %s',
  99         'lte': '<= %s',
 100         'startswith': 'LIKE %s',
 101         'endswith': 'LIKE %s',
 102         'istartswith': 'LIKE UPPER(%s)',
 103         'iendswith': 'LIKE UPPER(%s)',
 104     }
 105
 106     def __init__(self, *args, **kwargs):
 107         super(DatabaseWrapper, self).__init__(*args, **kwargs)
 108
 109         self.features = DatabaseFeatures(self)
 110         autocommit = self.settings_dict["OPTIONS"].get('autocommit', False)
 111         self.features.uses_autocommit = autocommit
 112         if autocommit:
 113             level = psycopg2.extensions.ISOLATION_LEVEL_AUTOCOMMIT
 114         else:
 115             level = psycopg2.extensions.ISOLATION_LEVEL_READ_COMMITTED
 116         self._set_isolation_level(level)
 117         self.ops = DatabaseOperations(self)
 118         self.client = DatabaseClient(self)
 119         self.creation = DatabaseCreation(self)
 120         self.introspection = DatabaseIntrospection(self)
 121         self.validation = BaseDatabaseValidation(self)
 122         self._pg_version = None
 123
 124     def check_constraints(self, table_names=None):
 125         """
 126         To check constraints, we set constraints to immediate. Then, when, we're done we must ensure they
 127         are returned to deferred.
 128         """
 129         self.cursor().execute('SET CONSTRAINTS ALL IMMEDIATE')
 130         self.cursor().execute('SET CONSTRAINTS ALL DEFERRED')
 131
 132     def close(self):
 133         self.validate_thread_sharing()
 134         if self.connection is None:
 135             return
 136
 137         try:
 138             self.connection.close()
 139             self.connection = None
 140         except Database.Error:
 141             # In some cases (database restart, network connection lost etc...)
 142             # the connection to the database is lost without giving Django a
 143             # notification. If we don't set self.connection to None, the error
 144             # will occur a every request.
 145             self.connection = None
 146             logger.warning('psycopg2 error while closing the connection.',
 147                 exc_info=sys.exc_info()
 148             )
 149             raise
 150
 151     def _get_pg_version(self):
 152         if self._pg_version is None:
 153             self._pg_version = get_version(self.connection)
 154         return self._pg_version
 155     pg_version = property(_get_pg_version)
 156
 157     def _cursor(self):
 158         settings_dict = self.settings_dict
 159         if self.connection is None:
 160             if settings_dict['NAME'] == '':
 161                 from django.core.exceptions import ImproperlyConfigured
 162                 raise ImproperlyConfigured("You need to specify NAME in your Django settings file.")
 163             conn_params = {
 164                 'database': settings_dict['NAME'],
 165             }
 166             conn_params.update(settings_dict['OPTIONS'])
 167             if 'autocommit' in conn_params:
 168                 del conn_params['autocommit']
 169             if settings_dict['USER']:
 170                 conn_params['user'] = settings_dict['USER']
 171             if settings_dict['PASSWORD']:
 172                 conn_params['password'] = settings_dict['PASSWORD']
 173             if settings_dict['HOST']:
 174                 conn_params['host'] = settings_dict['HOST']
 175             if settings_dict['PORT']:
 176                 conn_params['port'] = settings_dict['PORT']
 177             self.connection = Database.connect(**conn_params)
 178             self.connection.set_client_encoding('UTF8')
 179             tz = 'UTC' if settings.USE_TZ else settings_dict.get('TIME_ZONE')
 180             if tz:
 181                 try:
 182                     get_parameter_status = self.connection.get_parameter_status
 183                 except AttributeError:
 184                     # psycopg2 < 2.0.12 doesn't have get_parameter_status
 185                     conn_tz = None
 186                 else:
 187                     conn_tz = get_parameter_status('TimeZone')
 188
 189                 if conn_tz != tz:
 190                     # Set the time zone in autocommit mode (see #17062)
 191                     self.connection.set_isolation_level(
 192                             psycopg2.extensions.ISOLATION_LEVEL_AUTOCOMMIT)
 193                     self.connection.cursor().execute(
 194                             self.ops.set_time_zone_sql(), [tz])
 195             self.connection.set_isolation_level(self.isolation_level)
 196             self._get_pg_version()
 197             connection_created.send(sender=self.__class__, connection=self)
 198         cursor = self.connection.cursor()
 199         cursor.tzinfo_factory = utc_tzinfo_factory if settings.USE_TZ else None
 200         return CursorWrapper(cursor)
 201
 202     def _enter_transaction_management(self, managed):
 203         """
 204         Switch the isolation level when needing transaction support, so that
 205         the same transaction is visible across all the queries.
 206         """
 207         if self.features.uses_autocommit and managed and not self.isolation_level:
 208             self._set_isolation_level(psycopg2.extensions.ISOLATION_LEVEL_READ_COMMITTED)
 209
 210     def _leave_transaction_management(self, managed):
 211         """
 212         If the normal operating mode is "autocommit", switch back to that when
 213         leaving transaction management.
 214         """
 215         if self.features.uses_autocommit and not managed and self.isolation_level:
 216             self._set_isolation_level(psycopg2.extensions.ISOLATION_LEVEL_AUTOCOMMIT)
 217
 218     def _set_isolation_level(self, level):
 219         """
 220         Do all the related feature configurations for changing isolation
 221         levels. This doesn't touch the uses_autocommit feature, since that
 222         controls the movement *between* isolation levels.
 223         """
 224         assert level in range(5)
 225         try:
 226             if self.connection is not None:
 227                 self.connection.set_isolation_level(level)
 228         finally:
 229             self.isolation_level = level
 230             self.features.uses_savepoints = bool(level)
 231
 232     def _commit(self):
 233         if self.connection is not None:
 234             try:
 235                 return self.connection.commit()
 236             except Database.IntegrityError, e:
 237                 raise utils.IntegrityError, utils.IntegrityError(*tuple(e)), sys.exc_info()[2]