Ticket #30511: base.py

File base.py, 11.0 KB (added by Michael Kany, 4 years ago)
Line 
1"""
2PostgreSQL database backend for Django.
3
4Requires psycopg 2: http://initd.org/projects/psycopg2
5"""
6
7import threading
8import warnings
9
10from django.conf import settings
11from django.core.exceptions import ImproperlyConfigured
12from django.db import connections
13from django.db.backends.base.base import BaseDatabaseWrapper
14from django.db.utils import DatabaseError as WrappedDatabaseError
15from django.utils.functional import cached_property
16from django.utils.safestring import SafeText
17from django.utils.version import get_version_tuple
18
19try:
20    import psycopg2 as Database
21    import psycopg2.extensions
22    import psycopg2.extras
23except ImportError as e:
24    raise ImproperlyConfigured("Error loading psycopg2 module: %s" % e)
25
26
27def psycopg2_version():
28    version = psycopg2.__version__.split(' ', 1)[0]
29    return get_version_tuple(version)
30
31
32PSYCOPG2_VERSION = psycopg2_version()
33
34if PSYCOPG2_VERSION < (2, 5, 4):
35    raise ImproperlyConfigured("psycopg2_version 2.5.4 or newer is required; you have %s" % psycopg2.__version__)
36
37
38# Some of these import psycopg2, so import them after checking if it's installed.
39from .client import DatabaseClient                          # NOQA isort:skip
40from .creation import DatabaseCreation                      # NOQA isort:skip
41from .features import DatabaseFeatures                      # NOQA isort:skip
42from .introspection import DatabaseIntrospection            # NOQA isort:skip
43from .operations import DatabaseOperations                  # NOQA isort:skip
44from .schema import DatabaseSchemaEditor                    # NOQA isort:skip
45from .utils import utc_tzinfo_factory                       # NOQA isort:skip
46
47psycopg2.extensions.register_adapter(SafeText, psycopg2.extensions.QuotedString)
48psycopg2.extras.register_uuid()
49
50# Register support for inet[] manually so we don't have to handle the Inet()
51# object on load all the time.
52INETARRAY_OID = 1041
53INETARRAY = psycopg2.extensions.new_array_type(
54    (INETARRAY_OID,),
55    'INETARRAY',
56    psycopg2.extensions.UNICODE,
57)
58psycopg2.extensions.register_type(INETARRAY)
59
60
61class DatabaseWrapper(BaseDatabaseWrapper):
62    vendor = 'postgresql'
63    display_name = 'PostgreSQL'
64    # This dictionary maps Field objects to their associated PostgreSQL column
65    # types, as strings. Column-type strings can contain format strings; they'll
66    # be interpolated against the values of Field.__dict__ before being output.
67    # If a column type is set to None, it won't be included in the output.
68
69    data_types = {
70        'IdentityAutoField': 'integer',
71        'IdentityBigAutoField': 'bigint',
72        'AutoField': 'serial',
73        'BigAutoField': 'bigserial',
74        'BinaryField': 'bytea',
75        'BooleanField': 'boolean',
76        'CharField': 'varchar(%(max_length)s)',
77        'DateField': 'date',
78        'DateTimeField': 'timestamp with time zone',
79        'DecimalField': 'numeric(%(max_digits)s, %(decimal_places)s)',
80        'DurationField': 'interval',
81        'FileField': 'varchar(%(max_length)s)',
82        'FilePathField': 'varchar(%(max_length)s)',
83        'FloatField': 'double precision',
84        'IntegerField': 'integer',
85        'BigIntegerField': 'bigint',
86        'IPAddressField': 'inet',
87        'GenericIPAddressField': 'inet',
88        'NullBooleanField': 'boolean',
89        'OneToOneField': 'integer',
90        'PositiveIntegerField': 'integer',
91        'PositiveSmallIntegerField': 'smallint',
92        'SlugField': 'varchar(%(max_length)s)',
93        'SmallIntegerField': 'smallint',
94        'TextField': 'text',
95        'TimeField': 'time',
96        'UUIDField': 'uuid',
97    }
98    data_type_check_constraints = {
99        'PositiveIntegerField': '"%(column)s" >= 0',
100        'PositiveSmallIntegerField': '"%(column)s" >= 0',
101    }
102    operators = {
103        'exact': '= %s',
104        'iexact': '= UPPER(%s)',
105        'contains': 'LIKE %s',
106        'icontains': 'LIKE UPPER(%s)',
107        'regex': '~ %s',
108        'iregex': '~* %s',
109        'gt': '> %s',
110        'gte': '>= %s',
111        'lt': '< %s',
112        'lte': '<= %s',
113        'startswith': 'LIKE %s',
114        'endswith': 'LIKE %s',
115        'istartswith': 'LIKE UPPER(%s)',
116        'iendswith': 'LIKE UPPER(%s)',
117    }
118
119    # The patterns below are used to generate SQL pattern lookup clauses when
120    # the right-hand side of the lookup isn't a raw string (it might be an expression
121    # or the result of a bilateral transformation).
122    # In those cases, special characters for LIKE operators (e.g. \, *, _) should be
123    # escaped on database side.
124    #
125    # Note: we use str.format() here for readability as '%' is used as a wildcard for
126    # the LIKE operator.
127    pattern_esc = r"REPLACE(REPLACE(REPLACE({}, '\', '\\'), '%%', '\%%'), '_', '\_')"
128    pattern_ops = {
129        'contains': "LIKE '%%' || {} || '%%'",
130        'icontains': "LIKE '%%' || UPPER({}) || '%%'",
131        'startswith': "LIKE {} || '%%'",
132        'istartswith': "LIKE UPPER({}) || '%%'",
133        'endswith': "LIKE '%%' || {}",
134        'iendswith': "LIKE '%%' || UPPER({})",
135    }
136
137    Database = Database
138    SchemaEditorClass = DatabaseSchemaEditor
139    # Classes instantiated in __init__().
140    client_class = DatabaseClient
141    creation_class = DatabaseCreation
142    features_class = DatabaseFeatures
143    introspection_class = DatabaseIntrospection
144    ops_class = DatabaseOperations
145    # PostgreSQL backend-specific attributes.
146    _named_cursor_idx = 0
147
148    def get_connection_params(self):
149        settings_dict = self.settings_dict
150        # None may be used to connect to the default 'postgres' db
151        if settings_dict['NAME'] == '':
152            raise ImproperlyConfigured(
153                "settings.DATABASES is improperly configured. "
154                "Please supply the NAME value.")
155        if len(settings_dict['NAME'] or '') > self.ops.max_name_length():
156            raise ImproperlyConfigured(
157                "The database name '%s' (%d characters) is longer than "
158                "PostgreSQL's limit of %d characters. Supply a shorter NAME "
159                "in settings.DATABASES." % (
160                    settings_dict['NAME'],
161                    len(settings_dict['NAME']),
162                    self.ops.max_name_length(),
163                )
164            )
165        conn_params = {
166            'database': settings_dict['NAME'] or 'postgres',
167            **settings_dict['OPTIONS'],
168        }
169        conn_params.pop('isolation_level', None)
170        if settings_dict['USER']:
171            conn_params['user'] = settings_dict['USER']
172        if settings_dict['PASSWORD']:
173            conn_params['password'] = settings_dict['PASSWORD']
174        if settings_dict['HOST']:
175            conn_params['host'] = settings_dict['HOST']
176        if settings_dict['PORT']:
177            conn_params['port'] = settings_dict['PORT']
178        return conn_params
179
180    def get_new_connection(self, conn_params):
181        connection = Database.connect(**conn_params)
182
183        # self.isolation_level must be set:
184        # - after connecting to the database in order to obtain the database's
185        #   default when no value is explicitly specified in options.
186        # - before calling _set_autocommit() because if autocommit is on, that
187        #   will set connection.isolation_level to ISOLATION_LEVEL_AUTOCOMMIT.
188        options = self.settings_dict['OPTIONS']
189        try:
190            self.isolation_level = options['isolation_level']
191        except KeyError:
192            self.isolation_level = connection.isolation_level
193        else:
194            # Set the isolation level to the value from OPTIONS.
195            if self.isolation_level != connection.isolation_level:
196                connection.set_session(isolation_level=self.isolation_level)
197
198        return connection
199
200    def ensure_timezone(self):
201        self.ensure_connection()
202        conn_timezone_name = self.connection.get_parameter_status('TimeZone')
203        timezone_name = self.timezone_name
204        if timezone_name and conn_timezone_name != timezone_name:
205            with self.connection.cursor() as cursor:
206                cursor.execute(self.ops.set_time_zone_sql(), [timezone_name])
207            return True
208        return False
209
210    def init_connection_state(self):
211        self.connection.set_client_encoding('UTF8')
212
213        timezone_changed = self.ensure_timezone()
214        if timezone_changed:
215            # Commit after setting the time zone (see #17062)
216            if not self.get_autocommit():
217                self.connection.commit()
218
219    def create_cursor(self, name=None):
220        if name:
221            # In autocommit mode, the cursor will be used outside of a
222            # transaction, hence use a holdable cursor.
223            cursor = self.connection.cursor(name, scrollable=False, withhold=self.connection.autocommit)
224        else:
225            cursor = self.connection.cursor()
226        cursor.tzinfo_factory = utc_tzinfo_factory if settings.USE_TZ else None
227        return cursor
228
229    def chunked_cursor(self):
230        self._named_cursor_idx += 1
231        return self._cursor(
232            name='_django_curs_%d_%d' % (
233                # Avoid reusing name in other threads
234                threading.current_thread().ident,
235                self._named_cursor_idx,
236            )
237        )
238
239    def _set_autocommit(self, autocommit):
240        with self.wrap_database_errors:
241            self.connection.autocommit = autocommit
242
243    def check_constraints(self, table_names=None):
244        """
245        Check constraints by setting them to immediate. Return them to deferred
246        afterward.
247        """
248        self.cursor().execute('SET CONSTRAINTS ALL IMMEDIATE')
249        self.cursor().execute('SET CONSTRAINTS ALL DEFERRED')
250
251    def is_usable(self):
252        try:
253            # Use a psycopg cursor directly, bypassing Django's utilities.
254            self.connection.cursor().execute("SELECT 1")
255        except Database.Error:
256            return False
257        else:
258            return True
259
260    @property
261    def _nodb_connection(self):
262        nodb_connection = super()._nodb_connection
263        try:
264            nodb_connection.ensure_connection()
265        except (Database.DatabaseError, WrappedDatabaseError):
266            warnings.warn(
267                "Normally Django will use a connection to the 'postgres' database "
268                "to avoid running initialization queries against the production "
269                "database when it's not needed (for example, when running tests). "
270                "Django was unable to create a connection to the 'postgres' database "
271                "and will use the first PostgreSQL database instead.",
272                RuntimeWarning
273            )
274            for connection in connections.all():
275                if connection.vendor == 'postgresql' and connection.settings_dict['NAME'] != 'postgres':
276                    return self.__class__(
277                        {**self.settings_dict, 'NAME': connection.settings_dict['NAME']},
278                        alias=self.alias,
279                        allow_thread_sharing=False,
280                    )
281        return nodb_connection
282
283    @cached_property
284    def pg_version(self):
285        with self.temporary_connection():
286            return self.connection.server_version
Back to Top