|
1
|
1 |
|
from typing import Dict, List, Union, Tuple |
|
|
|
|
|
|
2
|
|
|
|
|
3
|
1 |
|
import django |
|
|
|
|
|
|
4
|
1 |
|
from django.conf import settings |
|
|
|
|
|
|
5
|
1 |
|
from django.db import models, transaction |
|
|
|
|
|
|
6
|
1 |
|
from django.db.models.sql import UpdateQuery |
|
|
|
|
|
|
7
|
1 |
|
from django.db.models.sql.constants import CURSOR |
|
|
|
|
|
|
8
|
1 |
|
from django.db.models.fields import NOT_PROVIDED |
|
|
|
|
|
|
9
|
1 |
|
from django.core.exceptions import ImproperlyConfigured, SuspiciousOperation |
|
|
|
|
|
|
10
|
|
|
|
|
11
|
1 |
|
from psqlextra import signals |
|
12
|
1 |
|
from psqlextra.compiler import (PostgresReturningUpdateCompiler, |
|
13
|
|
|
PostgresInsertCompiler) |
|
14
|
1 |
|
from psqlextra.query import PostgresQuery, PostgresInsertQuery, ConflictAction |
|
15
|
|
|
|
|
16
|
|
|
|
|
17
|
1 |
|
class PostgresQuerySet(models.QuerySet): |
|
18
|
|
|
"""Adds support for PostgreSQL specifics.""" |
|
19
|
|
|
|
|
20
|
1 |
|
def __init__(self, model=None, query=None, using=None, hints=None): |
|
21
|
|
|
"""Initializes a new instance of :see:PostgresQuerySet.""" |
|
22
|
|
|
|
|
23
|
1 |
|
super().__init__(model, query, using, hints) |
|
24
|
|
|
|
|
25
|
1 |
|
self.query = query or PostgresQuery(self.model) |
|
|
|
|
|
|
26
|
|
|
|
|
27
|
1 |
|
self.conflict_target = None |
|
28
|
1 |
|
self.conflict_action = None |
|
29
|
1 |
|
self.index_predicate = None |
|
30
|
|
|
|
|
31
|
1 |
|
def annotate(self, **annotations): |
|
32
|
|
|
"""Custom version of the standard annotate function |
|
33
|
|
|
that allows using field names as annotated fields. |
|
34
|
|
|
|
|
35
|
|
|
Normally, the annotate function doesn't allow you |
|
36
|
|
|
to use the name of an existing field on the model |
|
37
|
|
|
as the alias name. This version of the function does |
|
38
|
|
|
allow that. |
|
39
|
|
|
""" |
|
40
|
|
|
|
|
41
|
1 |
|
fields = { |
|
42
|
|
|
field.name: field |
|
43
|
|
|
for field in self.model._meta.get_fields() |
|
|
|
|
|
|
44
|
|
|
} |
|
45
|
|
|
|
|
46
|
|
|
# temporarily rename the fields that have the same |
|
47
|
|
|
# name as a field name, we'll rename them back after |
|
48
|
|
|
# the function in the base class ran |
|
49
|
1 |
|
new_annotations = {} |
|
50
|
1 |
|
renames = {} |
|
51
|
1 |
|
for name, value in annotations.items(): |
|
52
|
1 |
|
if name in fields: |
|
53
|
|
|
new_name = '%s_new' % name |
|
54
|
|
|
new_annotations[new_name] = value |
|
55
|
|
|
renames[new_name] = name |
|
56
|
|
|
else: |
|
57
|
1 |
|
new_annotations[name] = value |
|
58
|
|
|
|
|
59
|
|
|
# run the base class's annotate function |
|
60
|
1 |
|
result = super().annotate(**new_annotations) |
|
|
|
|
|
|
61
|
|
|
|
|
62
|
|
|
# rename the annotations back to as specified |
|
63
|
1 |
|
result.rename_annotations(**renames) |
|
|
|
|
|
|
64
|
1 |
|
return result |
|
65
|
|
|
|
|
66
|
1 |
|
def rename_annotations(self, **annotations): |
|
67
|
|
|
"""Renames the aliases for the specified annotations: |
|
68
|
|
|
|
|
69
|
|
|
.annotate(myfield=F('somestuf__myfield')) |
|
70
|
|
|
.rename_annotations(myfield='field') |
|
71
|
|
|
|
|
72
|
|
|
Arguments: |
|
73
|
|
|
annotations: |
|
74
|
|
|
The annotations to rename. Mapping the |
|
75
|
|
|
old name to the new name. |
|
76
|
|
|
""" |
|
77
|
|
|
|
|
78
|
1 |
|
self.query.rename_annotations(annotations) |
|
79
|
1 |
|
return self |
|
80
|
|
|
|
|
81
|
1 |
|
def join(self, **conditions): |
|
82
|
|
|
"""Adds extra conditions to existing joins. |
|
83
|
|
|
|
|
84
|
|
|
WARNING: This is an extremely experimental feature. |
|
85
|
|
|
DO NOT USE unless you know what you're doing. |
|
86
|
|
|
""" |
|
87
|
|
|
|
|
88
|
|
|
self.query.add_join_conditions(conditions) |
|
89
|
|
|
return self |
|
90
|
|
|
|
|
91
|
1 |
|
def update(self, **fields): |
|
92
|
|
|
"""Updates all rows that match the filter.""" |
|
93
|
|
|
|
|
94
|
|
|
# build up the query to execute |
|
95
|
1 |
|
self._for_write = True |
|
|
|
|
|
|
96
|
1 |
|
if django.VERSION >= (2, 0): |
|
97
|
1 |
|
query = self.query.chain(UpdateQuery) |
|
98
|
|
|
else: |
|
99
|
|
|
query = self.query.clone(UpdateQuery) |
|
100
|
1 |
|
query._annotations = None |
|
|
|
|
|
|
101
|
1 |
|
query.add_update_values(fields) |
|
102
|
|
|
|
|
103
|
|
|
# build the compiler for for the query |
|
104
|
1 |
|
connection = django.db.connections[self.db] |
|
|
|
|
|
|
105
|
1 |
|
compiler = PostgresReturningUpdateCompiler(query, connection, self.db) |
|
|
|
|
|
|
106
|
|
|
|
|
107
|
|
|
# execute the query |
|
108
|
1 |
|
with transaction.atomic(using=self.db, savepoint=False): |
|
|
|
|
|
|
109
|
1 |
|
rows = compiler.execute_sql(CURSOR) |
|
110
|
1 |
|
self._result_cache = None |
|
|
|
|
|
|
111
|
|
|
|
|
112
|
|
|
# send out a signal for each row |
|
113
|
1 |
|
for row in rows: |
|
114
|
1 |
|
signals.update.send(self.model, pk=row[0]) |
|
|
|
|
|
|
115
|
|
|
|
|
116
|
|
|
# the original update(..) returns the amount of rows |
|
117
|
|
|
# affected, let's do the same |
|
118
|
1 |
|
return len(rows) |
|
119
|
|
|
|
|
120
|
1 |
|
def on_conflict(self, fields: List[Union[str, Tuple[str]]], action, index_predicate: str=None): |
|
121
|
|
|
"""Sets the action to take when conflicts arise when attempting |
|
122
|
|
|
to insert/create a new row. |
|
123
|
|
|
|
|
124
|
|
|
Arguments: |
|
125
|
|
|
fields: |
|
126
|
|
|
The fields the conflicts can occur in. |
|
127
|
|
|
|
|
128
|
|
|
action: |
|
129
|
|
|
The action to take when the conflict occurs. |
|
130
|
|
|
|
|
131
|
|
|
index_predicate: |
|
132
|
|
|
The index predicate to satisfy an arbiter partial index (i.e. what partial index to use for checking |
|
133
|
|
|
conflicts) |
|
134
|
|
|
""" |
|
135
|
|
|
|
|
136
|
1 |
|
self.conflict_target = fields |
|
137
|
1 |
|
self.conflict_action = action |
|
138
|
1 |
|
self.index_predicate = index_predicate |
|
139
|
|
|
|
|
140
|
1 |
|
return self |
|
141
|
|
|
|
|
142
|
1 |
|
def bulk_insert(self, rows, return_model=False): |
|
143
|
|
|
"""Creates multiple new records in the database. |
|
144
|
|
|
|
|
145
|
|
|
This allows specifying custom conflict behavior using .on_conflict(). |
|
146
|
|
|
If no special behavior was specified, this uses the normal Django create(..) |
|
147
|
|
|
|
|
148
|
|
|
Arguments: |
|
149
|
|
|
rows: |
|
150
|
|
|
An array of dictionaries, where each dictionary |
|
151
|
|
|
describes the fields to insert. |
|
152
|
|
|
|
|
153
|
|
|
return_model (default: False): |
|
154
|
|
|
If model instances should be returned rather than |
|
155
|
|
|
just dicts. |
|
156
|
|
|
|
|
157
|
|
|
Returns: |
|
158
|
|
|
A list of either the dicts of the rows inserted, including the pk or |
|
159
|
|
|
the models of the rows inserted with defaults for any fields not specified |
|
160
|
|
|
""" |
|
161
|
|
|
|
|
162
|
1 |
|
if self.conflict_target or self.conflict_action: |
|
163
|
1 |
|
compiler = self._build_insert_compiler(rows) |
|
164
|
1 |
|
objs = compiler.execute_sql(return_id=True) |
|
165
|
1 |
|
if return_model: |
|
166
|
|
|
return [self.model(**dict(r, **k)) for r, k in zip(rows, objs)] |
|
|
|
|
|
|
167
|
|
|
else: |
|
168
|
1 |
|
return [dict(r, **k) for r, k in zip(rows, objs)] |
|
|
|
|
|
|
169
|
|
|
|
|
170
|
|
|
# no special action required, use the standard Django bulk_create(..) |
|
171
|
|
|
return super().bulk_create([self.model(**fields) for fields in rows]) |
|
|
|
|
|
|
172
|
|
|
|
|
173
|
1 |
|
def insert(self, **fields): |
|
174
|
|
|
"""Creates a new record in the database. |
|
175
|
|
|
|
|
176
|
|
|
This allows specifying custom conflict behavior using .on_conflict(). |
|
177
|
|
|
If no special behavior was specified, this uses the normal Django create(..) |
|
178
|
|
|
|
|
179
|
|
|
Arguments: |
|
180
|
|
|
fields: |
|
181
|
|
|
The fields of the row to create. |
|
182
|
|
|
|
|
183
|
|
|
Returns: |
|
184
|
|
|
The primary key of the record that was created. |
|
185
|
|
|
""" |
|
186
|
|
|
|
|
187
|
1 |
|
if self.conflict_target or self.conflict_action: |
|
188
|
1 |
|
compiler = self._build_insert_compiler([fields]) |
|
189
|
1 |
|
rows = compiler.execute_sql(return_id=True) |
|
190
|
1 |
|
if 'id' in rows[0]: |
|
191
|
1 |
|
return rows[0]['id'] |
|
192
|
|
|
return None |
|
193
|
|
|
|
|
194
|
|
|
# no special action required, use the standard Django create(..) |
|
195
|
|
|
return super().create(**fields).id |
|
196
|
|
|
|
|
197
|
1 |
|
def insert_and_get(self, **fields): |
|
198
|
|
|
"""Creates a new record in the database and then gets |
|
199
|
|
|
the entire row. |
|
200
|
|
|
|
|
201
|
|
|
This allows specifying custom conflict behavior using .on_conflict(). |
|
202
|
|
|
If no special behavior was specified, this uses the normal Django create(..) |
|
203
|
|
|
|
|
204
|
|
|
Arguments: |
|
205
|
|
|
fields: |
|
206
|
|
|
The fields of the row to create. |
|
207
|
|
|
|
|
208
|
|
|
Returns: |
|
209
|
|
|
The model instance representing the row that was created. |
|
210
|
|
|
""" |
|
211
|
|
|
|
|
212
|
1 |
|
if not self.conflict_target and not self.conflict_action: |
|
213
|
|
|
# no special action required, use the standard Django create(..) |
|
214
|
|
|
return super().create(**fields) |
|
215
|
|
|
|
|
216
|
1 |
|
compiler = self._build_insert_compiler([fields]) |
|
217
|
1 |
|
rows = compiler.execute_sql(return_id=False) |
|
218
|
|
|
|
|
219
|
1 |
|
columns = rows[0] |
|
220
|
|
|
|
|
221
|
|
|
# get a list of columns that are officially part of the model |
|
222
|
1 |
|
model_columns = [ |
|
223
|
|
|
field.column |
|
224
|
|
|
for field in self.model._meta.local_concrete_fields |
|
|
|
|
|
|
225
|
|
|
] |
|
226
|
|
|
|
|
227
|
|
|
# strip out any columns/fields returned by the db that |
|
228
|
|
|
# are not present in the model |
|
229
|
1 |
|
model_init_fields = {} |
|
230
|
1 |
|
for column_name, column_value in columns.items(): |
|
231
|
1 |
|
if column_name not in model_columns: |
|
232
|
1 |
|
continue |
|
233
|
|
|
|
|
234
|
1 |
|
model_init_fields[column_name] = column_value |
|
235
|
|
|
|
|
236
|
1 |
|
return self.model(**model_init_fields) |
|
|
|
|
|
|
237
|
|
|
|
|
238
|
1 |
|
def upsert(self, conflict_target: List, fields: Dict, index_predicate: str=None) -> int: |
|
239
|
|
|
"""Creates a new record or updates the existing one |
|
240
|
|
|
with the specified data. |
|
241
|
|
|
|
|
242
|
|
|
Arguments: |
|
243
|
|
|
conflict_target: |
|
244
|
|
|
Fields to pass into the ON CONFLICT clause. |
|
245
|
|
|
|
|
246
|
|
|
fields: |
|
247
|
|
|
Fields to insert/update. |
|
248
|
|
|
|
|
249
|
|
|
index_predicate: |
|
250
|
|
|
The index predicate to satisfy an arbiter partial index (i.e. what partial index to use for checking |
|
251
|
|
|
conflicts) |
|
252
|
|
|
|
|
253
|
|
|
Returns: |
|
254
|
|
|
The primary key of the row that was created/updated. |
|
255
|
|
|
""" |
|
256
|
|
|
|
|
257
|
1 |
|
self.on_conflict(conflict_target, ConflictAction.UPDATE, index_predicate) |
|
258
|
1 |
|
return self.insert(**fields) |
|
|
|
|
|
|
259
|
|
|
|
|
260
|
1 |
|
def upsert_and_get(self, conflict_target: List, fields: Dict, index_predicate: str=None): |
|
261
|
|
|
"""Creates a new record or updates the existing one |
|
262
|
|
|
with the specified data and then gets the row. |
|
263
|
|
|
|
|
264
|
|
|
Arguments: |
|
265
|
|
|
conflict_target: |
|
266
|
|
|
Fields to pass into the ON CONFLICT clause. |
|
267
|
|
|
|
|
268
|
|
|
fields: |
|
269
|
|
|
Fields to insert/update. |
|
270
|
|
|
|
|
271
|
|
|
index_predicate: |
|
272
|
|
|
The index predicate to satisfy an arbiter partial index (i.e. what partial index to use for checking |
|
273
|
|
|
conflicts) |
|
274
|
|
|
|
|
275
|
|
|
Returns: |
|
276
|
|
|
The model instance representing the row |
|
277
|
|
|
that was created/updated. |
|
278
|
|
|
""" |
|
279
|
|
|
|
|
280
|
1 |
|
self.on_conflict(conflict_target, ConflictAction.UPDATE, index_predicate) |
|
281
|
1 |
|
return self.insert_and_get(**fields) |
|
|
|
|
|
|
282
|
|
|
|
|
283
|
1 |
|
def bulk_upsert(self, conflict_target: List, rows: List[Dict], index_predicate: str=None): |
|
284
|
|
|
"""Creates a set of new records or updates the existing |
|
285
|
|
|
ones with the specified data. |
|
286
|
|
|
|
|
287
|
|
|
Arguments: |
|
288
|
|
|
conflict_target: |
|
289
|
|
|
Fields to pass into the ON CONFLICT clause. |
|
290
|
|
|
|
|
291
|
|
|
rows: |
|
292
|
|
|
Rows to upsert. |
|
293
|
|
|
|
|
294
|
|
|
index_predicate: |
|
295
|
|
|
The index predicate to satisfy an arbiter partial index (i.e. what partial index to use for checking |
|
296
|
|
|
conflicts) |
|
297
|
|
|
""" |
|
298
|
|
|
|
|
299
|
1 |
|
if not rows or len(rows) <= 0: |
|
300
|
1 |
|
return |
|
301
|
|
|
|
|
302
|
1 |
|
self.on_conflict(conflict_target, ConflictAction.UPDATE, index_predicate) |
|
303
|
1 |
|
return self.bulk_insert(rows) |
|
304
|
|
|
|
|
305
|
1 |
|
def _build_insert_compiler(self, rows: List[Dict]): |
|
306
|
|
|
"""Builds the SQL compiler for a insert query. |
|
307
|
|
|
|
|
308
|
|
|
Arguments: |
|
309
|
|
|
rows: |
|
310
|
|
|
A list of dictionaries, where each entry |
|
311
|
|
|
describes a record to insert. |
|
312
|
|
|
|
|
313
|
|
|
Returns: |
|
314
|
|
|
The SQL compiler for the insert. |
|
315
|
|
|
""" |
|
316
|
|
|
|
|
317
|
|
|
# create model objects, we also have to detect cases |
|
318
|
|
|
# such as: |
|
319
|
|
|
# [dict(first_name='swen'), dict(fist_name='swen', last_name='kooij')] |
|
320
|
|
|
# we need to be certain that each row specifies the exact same |
|
321
|
|
|
# amount of fields/columns |
|
322
|
1 |
|
objs = [] |
|
323
|
1 |
|
field_count = len(rows[0]) |
|
324
|
1 |
|
for index, row in enumerate(rows): |
|
325
|
1 |
|
if field_count != len(row): |
|
326
|
|
|
raise SuspiciousOperation(( |
|
327
|
|
|
'In bulk upserts, you cannot have rows with different field ' |
|
328
|
|
|
'configurations. Row {0} has a different field config than ' |
|
329
|
|
|
'the first row.' |
|
330
|
|
|
).format(index)) |
|
331
|
|
|
|
|
332
|
1 |
|
objs.append(self.model(**row)) |
|
|
|
|
|
|
333
|
|
|
|
|
334
|
|
|
# indicate this query is going to perform write |
|
335
|
1 |
|
self._for_write = True |
|
|
|
|
|
|
336
|
|
|
|
|
337
|
|
|
# get the fields to be used during update/insert |
|
338
|
1 |
|
insert_fields, update_fields = self._get_upsert_fields(rows[0]) |
|
339
|
|
|
|
|
340
|
|
|
# build a normal insert query |
|
341
|
1 |
|
query = PostgresInsertQuery(self.model) |
|
|
|
|
|
|
342
|
1 |
|
query.conflict_action = self.conflict_action |
|
343
|
1 |
|
query.conflict_target = self.conflict_target |
|
344
|
1 |
|
query.index_predicate = self.index_predicate |
|
345
|
1 |
|
query.values(objs, insert_fields, update_fields) |
|
346
|
|
|
|
|
347
|
|
|
# use the postgresql insert query compiler to transform the insert |
|
348
|
|
|
# into an special postgresql insert |
|
349
|
1 |
|
connection = django.db.connections[self.db] |
|
|
|
|
|
|
350
|
1 |
|
compiler = PostgresInsertCompiler(query, connection, self.db) |
|
|
|
|
|
|
351
|
|
|
|
|
352
|
1 |
|
return compiler |
|
353
|
|
|
|
|
354
|
1 |
|
def _is_magical_field(self, model_instance, field, is_insert: bool): |
|
|
|
|
|
|
355
|
|
|
"""Verifies whether this field is gonna modify something |
|
356
|
|
|
on its own. |
|
357
|
|
|
|
|
358
|
|
|
"Magical" means that a field modifies the field value |
|
359
|
|
|
during the pre_save. |
|
360
|
|
|
|
|
361
|
|
|
Arguments: |
|
362
|
|
|
model_instance: |
|
363
|
|
|
The model instance the field is defined on. |
|
364
|
|
|
|
|
365
|
|
|
field: |
|
366
|
|
|
The field to get of whether the field is |
|
367
|
|
|
magical. |
|
368
|
|
|
|
|
369
|
|
|
is_insert: |
|
370
|
|
|
Pretend whether this is an insert? |
|
371
|
|
|
|
|
372
|
|
|
Returns: |
|
373
|
|
|
True when this field modifies something. |
|
374
|
|
|
""" |
|
375
|
|
|
|
|
376
|
|
|
# does this field modify someting upon insert? |
|
377
|
1 |
|
old_value = getattr(model_instance, field.name, None) |
|
378
|
1 |
|
field.pre_save(model_instance, is_insert) |
|
379
|
1 |
|
new_value = getattr(model_instance, field.name, None) |
|
380
|
|
|
|
|
381
|
1 |
|
return old_value != new_value |
|
382
|
|
|
|
|
383
|
1 |
|
def _get_upsert_fields(self, kwargs): |
|
384
|
|
|
"""Gets the fields to use in an upsert. |
|
385
|
|
|
|
|
386
|
|
|
This some nice magic. We'll split the fields into |
|
387
|
|
|
a group of "insert fields" and "update fields": |
|
388
|
|
|
|
|
389
|
|
|
INSERT INTO bla ("val1", "val2") ON CONFLICT DO UPDATE SET val1 = EXCLUDED.val1 |
|
390
|
|
|
|
|
391
|
|
|
^^^^^^^^^^^^^^ ^^^^^^^^^^^^^^^^^^^^ |
|
392
|
|
|
insert_fields update_fields |
|
393
|
|
|
|
|
394
|
|
|
Often, fields appear in both lists. But, for example, |
|
395
|
|
|
a :see:DateTime field with `auto_now_add=True` set, will |
|
396
|
|
|
only appear in "insert_fields", since it won't be set |
|
397
|
|
|
on existing rows. |
|
398
|
|
|
|
|
399
|
|
|
Other than that, the user specificies a list of fields |
|
400
|
|
|
in the upsert() call. That migt not be all fields. The |
|
401
|
|
|
user could decide to leave out optional fields. If we |
|
402
|
|
|
end up doing an update, we don't want to overwrite |
|
403
|
|
|
those non-specified fields. |
|
404
|
|
|
|
|
405
|
|
|
We cannot just take the list of fields the user |
|
406
|
|
|
specifies, because as mentioned, some fields |
|
407
|
|
|
make modifications to the model on their own. |
|
408
|
|
|
|
|
409
|
|
|
We'll have to detect which fields make modifications |
|
410
|
|
|
and include them in the list of insert/update fields. |
|
411
|
|
|
""" |
|
412
|
|
|
|
|
413
|
1 |
|
model_instance = self.model(**kwargs) |
|
|
|
|
|
|
414
|
1 |
|
insert_fields = [] |
|
415
|
1 |
|
update_fields = [] |
|
416
|
|
|
|
|
417
|
1 |
|
for field in model_instance._meta.local_concrete_fields: |
|
|
|
|
|
|
418
|
1 |
|
has_default = field.default != NOT_PROVIDED |
|
419
|
1 |
|
if (field.name in kwargs or field.column in kwargs): |
|
|
|
|
|
|
420
|
1 |
|
insert_fields.append(field) |
|
421
|
1 |
|
update_fields.append(field) |
|
422
|
1 |
|
continue |
|
423
|
1 |
|
elif has_default: |
|
424
|
1 |
|
insert_fields.append(field) |
|
425
|
1 |
|
continue |
|
426
|
|
|
|
|
427
|
|
|
# special handling for 'pk' which always refers to |
|
428
|
|
|
# the primary key, so if we the user specifies `pk` |
|
429
|
|
|
# instead of a concrete field, we have to handle that |
|
430
|
1 |
|
if field.primary_key is True and 'pk' in kwargs: |
|
431
|
1 |
|
insert_fields.append(field) |
|
432
|
1 |
|
update_fields.append(field) |
|
433
|
1 |
|
continue |
|
434
|
|
|
|
|
435
|
1 |
|
if self._is_magical_field(model_instance, field, is_insert=True): |
|
436
|
1 |
|
insert_fields.append(field) |
|
437
|
|
|
|
|
438
|
1 |
|
if self._is_magical_field(model_instance, field, is_insert=False): |
|
439
|
1 |
|
update_fields.append(field) |
|
440
|
|
|
|
|
441
|
1 |
|
return insert_fields, update_fields |
|
442
|
|
|
|
|
443
|
|
|
|
|
444
|
1 |
|
class PostgresManager(models.Manager): |
|
445
|
|
|
"""Adds support for PostgreSQL specifics.""" |
|
446
|
|
|
|
|
447
|
1 |
|
use_in_migrations = True |
|
448
|
|
|
|
|
449
|
1 |
|
def __init__(self, *args, **kwargs): |
|
450
|
|
|
"""Initializes a new instance of :see:PostgresManager.""" |
|
451
|
|
|
|
|
452
|
1 |
|
super(PostgresManager, self).__init__(*args, **kwargs) |
|
453
|
|
|
|
|
454
|
|
|
# make sure our back-end is set and refuse to proceed |
|
455
|
|
|
# if it's not set |
|
456
|
1 |
|
db_backend = settings.DATABASES['default']['ENGINE'] |
|
457
|
1 |
|
if 'psqlextra' not in db_backend: |
|
458
|
|
|
raise ImproperlyConfigured(( |
|
459
|
|
|
'\'%s\' is not the \'psqlextra.backend\'. ' |
|
460
|
|
|
'django-postgres-extra cannot function without ' |
|
461
|
|
|
'the \'psqlextra.backend\'. Set DATABASES.ENGINE.' |
|
462
|
|
|
) % db_backend) |
|
463
|
|
|
|
|
464
|
|
|
# hook into django signals to then trigger our own |
|
465
|
|
|
|
|
466
|
1 |
|
django.db.models.signals.post_save.connect( |
|
467
|
|
|
self._on_model_save, sender=self.model, weak=False) |
|
468
|
|
|
|
|
469
|
1 |
|
django.db.models.signals.pre_delete.connect( |
|
470
|
|
|
self._on_model_delete, sender=self.model, weak=False) |
|
471
|
|
|
|
|
472
|
1 |
|
self._signals_connected = True |
|
473
|
|
|
|
|
474
|
1 |
|
def __del__(self): |
|
475
|
|
|
"""Disconnects signals.""" |
|
476
|
|
|
|
|
477
|
1 |
|
if self._signals_connected is False: |
|
478
|
|
|
return |
|
479
|
|
|
|
|
480
|
|
|
# django.db.models.signals.post_save.disconnect( |
|
481
|
|
|
# self._on_model_save, sender=self.model) |
|
482
|
|
|
|
|
483
|
|
|
# django.db.models.signals.pre_delete.disconnect( |
|
484
|
|
|
# self._on_model_delete, sender=self.model) |
|
485
|
|
|
|
|
486
|
1 |
|
def get_queryset(self): |
|
487
|
|
|
"""Gets the query set to be used on this manager.""" |
|
488
|
|
|
|
|
489
|
1 |
|
return PostgresQuerySet(self.model, using=self._db) |
|
|
|
|
|
|
490
|
|
|
|
|
491
|
1 |
|
def on_conflict(self, fields: List[Union[str, Tuple[str]]], action, index_predicate: str=None): |
|
492
|
|
|
"""Sets the action to take when conflicts arise when attempting |
|
493
|
|
|
to insert/create a new row. |
|
494
|
|
|
|
|
495
|
|
|
Arguments: |
|
496
|
|
|
fields: |
|
497
|
|
|
The fields the conflicts can occur in. |
|
498
|
|
|
|
|
499
|
|
|
action: |
|
500
|
|
|
The action to take when the conflict occurs. |
|
501
|
|
|
|
|
502
|
|
|
index_predicate: |
|
503
|
|
|
The index predicate to satisfy an arbiter partial index. |
|
504
|
|
|
""" |
|
505
|
1 |
|
return self.get_queryset().on_conflict(fields, action, index_predicate) |
|
506
|
|
|
|
|
507
|
1 |
|
def upsert(self, conflict_target: List, fields: Dict, index_predicate: str=None) -> int: |
|
508
|
|
|
"""Creates a new record or updates the existing one |
|
509
|
|
|
with the specified data. |
|
510
|
|
|
|
|
511
|
|
|
Arguments: |
|
512
|
|
|
conflict_target: |
|
513
|
|
|
Fields to pass into the ON CONFLICT clause. |
|
514
|
|
|
|
|
515
|
|
|
fields: |
|
516
|
|
|
Fields to insert/update. |
|
517
|
|
|
|
|
518
|
|
|
index_predicate: |
|
519
|
|
|
The index predicate to satisfy an arbiter partial index. |
|
520
|
|
|
|
|
521
|
|
|
Returns: |
|
522
|
|
|
The primary key of the row that was created/updated. |
|
523
|
|
|
""" |
|
524
|
|
|
|
|
525
|
1 |
|
return self.get_queryset().upsert(conflict_target, fields, index_predicate) |
|
526
|
|
|
|
|
527
|
1 |
|
def upsert_and_get(self, conflict_target: List, fields: Dict, index_predicate: str=None): |
|
528
|
|
|
"""Creates a new record or updates the existing one |
|
529
|
|
|
with the specified data and then gets the row. |
|
530
|
|
|
|
|
531
|
|
|
Arguments: |
|
532
|
|
|
conflict_target: |
|
533
|
|
|
Fields to pass into the ON CONFLICT clause. |
|
534
|
|
|
|
|
535
|
|
|
fields: |
|
536
|
|
|
Fields to insert/update. |
|
537
|
|
|
|
|
538
|
|
|
index_predicate: |
|
539
|
|
|
The index predicate to satisfy an arbiter partial index. |
|
540
|
|
|
|
|
541
|
|
|
Returns: |
|
542
|
|
|
The model instance representing the row |
|
543
|
|
|
that was created/updated. |
|
544
|
|
|
""" |
|
545
|
|
|
|
|
546
|
1 |
|
return self.get_queryset().upsert_and_get(conflict_target, fields, index_predicate) |
|
547
|
|
|
|
|
548
|
1 |
|
def bulk_upsert(self, conflict_target: List, rows: List[Dict], index_predicate: str=None): |
|
549
|
|
|
"""Creates a set of new records or updates the existing |
|
550
|
|
|
ones with the specified data. |
|
551
|
|
|
|
|
552
|
|
|
Arguments: |
|
553
|
|
|
conflict_target: |
|
554
|
|
|
Fields to pass into the ON CONFLICT clause. |
|
555
|
|
|
|
|
556
|
|
|
index_predicate: |
|
557
|
|
|
The index predicate to satisfy an arbiter partial index. |
|
558
|
|
|
|
|
559
|
|
|
rows: |
|
560
|
|
|
Rows to upsert. |
|
561
|
|
|
""" |
|
562
|
|
|
|
|
563
|
1 |
|
return self.get_queryset().bulk_upsert(conflict_target, rows, index_predicate) |
|
564
|
|
|
|
|
565
|
1 |
|
@staticmethod |
|
566
|
|
|
def _on_model_save(sender, **kwargs): |
|
567
|
|
|
"""When a model gets created or updated.""" |
|
568
|
|
|
|
|
569
|
1 |
|
created, instance = kwargs['created'], kwargs['instance'] |
|
570
|
|
|
|
|
571
|
1 |
|
if created: |
|
572
|
1 |
|
signals.create.send(sender, pk=instance.pk) |
|
573
|
|
|
else: |
|
574
|
1 |
|
signals.update.send(sender, pk=instance.pk) |
|
575
|
|
|
|
|
576
|
1 |
|
@staticmethod |
|
577
|
|
|
def _on_model_delete(sender, **kwargs): |
|
578
|
|
|
"""When a model gets deleted.""" |
|
579
|
|
|
|
|
580
|
1 |
|
instance = kwargs['instance'] |
|
581
|
|
|
signals.delete.send(sender, pk=instance.pk) |
|
582
|
|
|
|
This can be caused by one of the following:
1. Missing Dependencies
This error could indicate a configuration issue of Pylint. Make sure that your libraries are available by adding the necessary commands.
2. Missing __init__.py files
This error could also result from missing
__init__.pyfiles in your module folders. Make sure that you place one file in each sub-folder.