Ticket #13161: base.py

File base.py, 8.1 KB (added by spoksss, 10 years ago)

Small usefull change.

Line 
1"""
2PostgreSQL database backend for Django.
3
4Requires psycopg 2: http://initd.org/projects/psycopg2
5"""
6
7import sys
8
9from django.db import utils
10from django.db.backends import *
11from django.db.backends.signals import connection_created
12from django.db.backends.postgresql.operations import DatabaseOperations as PostgresqlDatabaseOperations
13from django.db.backends.postgresql.client import DatabaseClient
14from django.db.backends.postgresql.creation import DatabaseCreation
15from django.db.backends.postgresql.version import get_version
16from django.db.backends.postgresql_psycopg2.introspection import DatabaseIntrospection
17from django.utils.safestring import SafeUnicode, SafeString
18
19try:
20    import psycopg2 as Database
21    import psycopg2.extensions
22except ImportError, e:
23    from django.core.exceptions import ImproperlyConfigured
24    raise ImproperlyConfigured("Error loading psycopg2 module: %s" % e)
25
26DatabaseError = Database.DatabaseError
27IntegrityError = Database.IntegrityError
28# Removed because it shouldn't be global settings,
29# some application (not django) could use SQLASCII not UNICODE
30# it is better if we set that extension only for connections which we are using in Django
31#psycopg2.extensions.register_type(psycopg2.extensions.UNICODE)
32psycopg2.extensions.register_adapter(SafeString, psycopg2.extensions.QuotedString)
33psycopg2.extensions.register_adapter(SafeUnicode, psycopg2.extensions.QuotedString)
34
35class CursorWrapper(object):
36    """
37    A thin wrapper around psycopg2's normal cursor class so that we can catch
38    particular exception instances and reraise them with the right types.
39    """
40
41    def __init__(self, cursor):
42        self.cursor = cursor
43
44    def execute(self, query, args=None):
45        try:
46            return self.cursor.execute(query, args)
47        except Database.IntegrityError, e:
48            raise utils.IntegrityError, utils.IntegrityError(*tuple(e)), sys.exc_info()[2]
49        except Database.DatabaseError, e:
50            raise utils.DatabaseError, utils.DatabaseError(*tuple(e)), sys.exc_info()[2]
51
52    def executemany(self, query, args):
53        try:
54            return self.cursor.executemany(query, args)
55        except Database.IntegrityError, e:
56            raise utils.IntegrityError, utils.IntegrityError(*tuple(e)), sys.exc_info()[2]
57        except Database.DatabaseError, e:
58            raise utils.DatabaseError, utils.DatabaseError(*tuple(e)), sys.exc_info()[2]
59
60    def __getattr__(self, attr):
61        if attr in self.__dict__:
62            return self.__dict__[attr]
63        else:
64            return getattr(self.cursor, attr)
65
66    def __iter__(self):
67        return iter(self.cursor)
68
69class DatabaseFeatures(BaseDatabaseFeatures):
70    needs_datetime_string_cast = False
71    can_return_id_from_insert = False
72
73class DatabaseOperations(PostgresqlDatabaseOperations):
74    def last_executed_query(self, cursor, sql, params):
75        # With psycopg2, cursor objects have a "query" attribute that is the
76        # exact query sent to the database. See docs here:
77        # http://www.initd.org/tracker/psycopg/wiki/psycopg2_documentation#postgresql-status-message-and-executed-query
78        return cursor.query
79
80    def return_insert_id(self):
81        return "RETURNING %s", ()
82
83class DatabaseWrapper(BaseDatabaseWrapper):
84    operators = {
85        'exact': '= %s',
86        'iexact': '= UPPER(%s)',
87        'contains': 'LIKE %s',
88        'icontains': 'LIKE UPPER(%s)',
89        'regex': '~ %s',
90        'iregex': '~* %s',
91        'gt': '> %s',
92        'gte': '>= %s',
93        'lt': '< %s',
94        'lte': '<= %s',
95        'startswith': 'LIKE %s',
96        'endswith': 'LIKE %s',
97        'istartswith': 'LIKE UPPER(%s)',
98        'iendswith': 'LIKE UPPER(%s)',
99    }
100
101    def __init__(self, *args, **kwargs):
102        super(DatabaseWrapper, self).__init__(*args, **kwargs)
103        self._set_extension()
104        self.features = DatabaseFeatures()
105        autocommit = self.settings_dict["OPTIONS"].get('autocommit', False)
106        self.features.uses_autocommit = autocommit
107        self._set_isolation_level(int(not autocommit))
108        self.ops = DatabaseOperations(self)
109        self.client = DatabaseClient(self)
110        self.creation = DatabaseCreation(self)
111        self.introspection = DatabaseIntrospection(self)
112        self.validation = BaseDatabaseValidation(self)
113   
114    def _set_extension(self):
115        if self.connection:
116            psycopg2.extensions.register_type(psycopg2.extensions.UNICODE,  self.connection)
117   
118    def _cursor(self):
119        new_connection = False
120        set_tz = False
121        settings_dict = self.settings_dict
122        if self.connection is None:
123            new_connection = True
124            set_tz = settings_dict.get('TIME_ZONE')
125            if settings_dict['NAME'] == '':
126                from django.core.exceptions import ImproperlyConfigured
127                raise ImproperlyConfigured("You need to specify NAME in your Django settings file.")
128            conn_params = {
129                'database': settings_dict['NAME'],
130            }
131            conn_params.update(settings_dict['OPTIONS'])
132            if 'autocommit' in conn_params:
133                del conn_params['autocommit']
134            if settings_dict['USER']:
135                conn_params['user'] = settings_dict['USER']
136            if settings_dict['PASSWORD']:
137                conn_params['password'] = settings_dict['PASSWORD']
138            if settings_dict['HOST']:
139                conn_params['host'] = settings_dict['HOST']
140            if settings_dict['PORT']:
141                conn_params['port'] = settings_dict['PORT']
142            self.connection = Database.connect(**conn_params)
143            self._set_extension()
144            self.connection.set_client_encoding('UTF8')
145            self.connection.set_isolation_level(self.isolation_level)
146            connection_created.send(sender=self.__class__)
147        cursor = self.connection.cursor()
148        cursor.tzinfo_factory = None
149        if new_connection:
150            if set_tz:
151                cursor.execute("SET TIME ZONE %s", [settings_dict['TIME_ZONE']])
152            if not hasattr(self, '_version'):
153                self.__class__._version = get_version(cursor)
154            if self._version[0:2] < (8, 0):
155                # No savepoint support for earlier version of PostgreSQL.
156                self.features.uses_savepoints = False
157            if self.features.uses_autocommit:
158                if self._version[0:2] < (8, 2):
159                    # FIXME: Needs extra code to do reliable model insert
160                    # handling, so we forbid it for now.
161                    from django.core.exceptions import ImproperlyConfigured
162                    raise ImproperlyConfigured("You cannot use autocommit=True with PostgreSQL prior to 8.2 at the moment.")
163                else:
164                    # FIXME: Eventually we're enable this by default for
165                    # versions that support it, but, right now, that's hard to
166                    # do without breaking other things (#10509).
167                    self.features.can_return_id_from_insert = True
168        return CursorWrapper(cursor)
169
170    def _enter_transaction_management(self, managed):
171        """
172        Switch the isolation level when needing transaction support, so that
173        the same transaction is visible across all the queries.
174        """
175        if self.features.uses_autocommit and managed and not self.isolation_level:
176            self._set_isolation_level(1)
177
178    def _leave_transaction_management(self, managed):
179        """
180        If the normal operating mode is "autocommit", switch back to that when
181        leaving transaction management.
182        """
183        if self.features.uses_autocommit and not managed and self.isolation_level:
184            self._set_isolation_level(0)
185
186    def _set_isolation_level(self, level):
187        """
188        Do all the related feature configurations for changing isolation
189        levels. This doesn't touch the uses_autocommit feature, since that
190        controls the movement *between* isolation levels.
191        """
192        assert level in (0, 1)
193        try:
194            if self.connection is not None:
195                self.connection.set_isolation_level(level)
196        finally:
197            self.isolation_level = level
198            self.features.uses_savepoints = bool(level)
Back to Top