Passed
Push — master ( ed1559...680383 )
by Swen
01:54
created

SchemaEditor._alter_field()   A

Complexity

Conditions 1

Size

Total Lines 9

Duplication

Lines 0
Ratio 0 %

Importance

Changes 1
Bugs 0 Features 0
Metric Value
cc 1
c 1
b 0
f 0
dl 0
loc 9
rs 9.6666
1
import importlib
2
3
from django.conf import settings
4
from django.core.exceptions import ImproperlyConfigured
5
from django.db.backends.postgresql.base import \
6
    DatabaseWrapper as Psycopg2DatabaseWrapper
7
8
9
def _get_backend_base():
10
    """Gets the base class for the custom database back-end.
11
12
    This should be the Django PostgreSQL back-end. However,
13
    some people are already using a custom back-end from
14
    another package. We are nice people and expose an option
15
    that allows them to configure the back-end we base upon.
16
17
    As long as the specified base eventually also has
18
    the PostgreSQL back-end as a base, then everything should
19
    work as intended.
20
    """
21
    base_class_name = getattr(
22
        settings,
23
        'LOCALIZED_FIELDS_DB_BACKEND_BASE',
24
        'django.db.backends.postgresql'
25
    )
26
27
    base_class_module = importlib.import_module(base_class_name + '.base')
28
    base_class = getattr(base_class_module, 'DatabaseWrapper', None)
29
30
    if not base_class:
31
        raise ImproperlyConfigured((
32
            '\'%s\' is not a valid database back-end.'
33
            ' The module does not define a DatabaseWrapper class.'
34
            ' Check the value of LOCALIZED_FIELDS_DB_BACKEND_BASE.'
35
        ))
36
37
    if isinstance(base_class, Psycopg2DatabaseWrapper):
38
        raise ImproperlyConfigured((
39
            '\'%s\' is not a valid database back-end.'
40
            ' It does inherit from the PostgreSQL back-end.'
41
            ' Check the value of LOCALIZED_FIELDS_DB_BACKEND_BASE.'
42
        ))
43
44
    return base_class
45
46
47
def _get_schema_editor_base():
48
    """Gets the base class for the schema editor.
49
50
    We have to use the configured base back-end's
51
    schema editor for this."""
52
    return _get_backend_base().SchemaEditorClass
53
54
55
class SchemaEditor(_get_schema_editor_base()):
56
    """Custom schema editor for hstore indexes.
57
58
    This allows us to put UNIQUE constraints for
59
    localized fields."""
60
61
    sql_hstore_unique_create = "CREATE UNIQUE INDEX {name} ON {table}{using} ({columns}){extra}"
62
    sql_hstore_unique_drop = "DROP INDEX IF EXISTS {name}"
63
64
    @staticmethod
65
    def _hstore_unique_name(model, field, keys):
66
        """Gets the name for a UNIQUE INDEX that applies
67
        to one or more keys in a hstore field.
68
69
        Arguments:
70
            model:
71
                The model the field is a part of.
72
73
            field:
74
                The hstore field to create a
75
                UNIQUE INDEX for.
76
77
            key:
78
                The name of the hstore key
79
                to create the name for.
80
81
                This can also be a tuple
82
                of multiple names.
83
84
        Returns:
85
            The name for the UNIQUE index.
86
        """
87
        postfix = '_'.join(keys)
88
        return '{table_name}_{field_name}_unique_{postfix}'.format(
89
            table_name=model._meta.db_table,
90
            field_name=field.column,
91
            postfix=postfix
92
        )
93
94
    def _drop_hstore_unique(self, model, field, keys):
95
        """Drops a UNIQUE constraint for the specified hstore keys."""
96
97
        name = self._hstore_unique_name(model, field, keys)
98
        sql = self.sql_hstore_unique_drop.format(name=name)
99
        self.execute(sql)
100
101
    def _create_hstore_unique(self, model, field, keys):
102
        """Creates a UNIQUE constraint for the specified hstore keys."""
103
104
        name = self._hstore_unique_name(model, field, keys)
105
        columns = [
106
            '(%s->\'%s\')' % (field.column, key)
107
            for key in keys
108
        ]
109
        sql = self.sql_hstore_unique_create.format(
110
            name=name,
111
            table=model._meta.db_table,
112
            using='',
113
            columns=','.join(columns),
114
            extra=''
115
        )
116
        self.execute(sql)
117
118
    def _update_hstore_constraints(self, model, old_field, new_field):
119
        """Updates the UNIQUE constraints for the specified field."""
120
121
        old_uniqueness = getattr(old_field, 'uniqueness', None)
122
        new_uniqueness = getattr(new_field, 'uniqueness', None)
123
124
        def _compose_keys(constraint):
125
            if isinstance(constraint, str):
126
                return [constraint]
127
128
            return constraint
129
130
        # drop any old uniqueness constraints
131
        if old_uniqueness:
132
            for keys in old_uniqueness:
133
                self._drop_hstore_unique(
134
                    model,
135
                    old_field,
136
                    _compose_keys(keys)
137
                )
138
139
        # (re-)create uniqueness constraints
140
        if new_uniqueness:
141
            for keys in new_uniqueness:
142
                self._create_hstore_unique(
143
                    model,
144
                    old_field,
145
                    _compose_keys(keys)
146
                )
147
148
    def _alter_field(self, model, old_field, new_field, *args, **kwargs):
149
        """Ran when the configuration on a field changed."""
150
151
        super()._alter_field(
152
            model, old_field, new_field,
153
            *args, **kwargs
154
        )
155
156
        self._update_hstore_constraints(model, old_field, new_field)
157
158
159
class DatabaseWrapper(_get_backend_base()):
160
    """Wraps the standard PostgreSQL database back-end.
161
162
    Overrides the schema editor with our custom
163
    schema editor and makes sure the `hstore`
164
    extension is enabled."""
165
166
    SchemaEditorClass = SchemaEditor
167
168
    def prepare_database(self):
169
        """Ran to prepare the configured database.
170
171
        This is where we enable the `hstore` extension
172
        if it wasn't enabled yet."""
173
174
        super().prepare_database()
175
        with self.cursor() as cursor:
176
            cursor.execute('CREATE EXTENSION IF NOT EXISTS hstore')
177