Skip to content

Commit 4376299

Browse files
Address palewire#91.
1 parent 9159b12 commit 4376299

File tree

1 file changed

+45
-9
lines changed

1 file changed

+45
-9
lines changed

postgres_copy/copy_from.py

Lines changed: 45 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@
99
import logging
1010
from collections import OrderedDict
1111
from io import TextIOWrapper
12+
import warnings
1213
from django.db import NotSupportedError
1314
from django.db import connections, router
1415
from django.core.exceptions import FieldDoesNotExist
@@ -33,6 +34,7 @@ def __init__(
3334
force_null=None,
3435
encoding=None,
3536
ignore_conflicts=False,
37+
on_conflict={},
3638
static_mapping=None,
3739
temp_table_name=None
3840
):
@@ -57,8 +59,9 @@ def __init__(
5759
self.force_not_null = force_not_null
5860
self.force_null = force_null
5961
self.encoding = encoding
60-
self.supports_ignore_conflicts = True
62+
self.supports_on_conflict = True
6163
self.ignore_conflicts = ignore_conflicts
64+
self.on_conflict = on_conflict
6265
if static_mapping is not None:
6366
self.static_mapping = OrderedDict(static_mapping)
6467
else:
@@ -76,10 +79,18 @@ def __init__(
7679
if self.conn.vendor != 'postgresql':
7780
raise TypeError("Only PostgreSQL backends supported")
7881

79-
# Check if it is PSQL 9.5 or greater, which determines if ignore_conflicts is supported
80-
self.supports_ignore_conflicts = self.is_postgresql_9_5()
81-
if self.ignore_conflicts and not self.supports_ignore_conflicts:
82-
raise NotSupportedError('This database backend does not support ignoring conflicts.')
82+
# Check if it is PSQL 9.5 or greater, which determines if on_conflict is supported
83+
self.supports_on_conflict = self.is_postgresql_9_5()
84+
if self.ignore_conflicts:
85+
self.on_conflict = {
86+
'action': 'ignore',
87+
}
88+
warnings.warn(
89+
"The `ignore_conflicts` kwarg has been replaced with "
90+
"on_conflict={'action': 'ignore'}."
91+
)
92+
if self.on_conflict and not self.supports_on_conflict:
93+
raise NotSupportedError('This database backend does not support conflict logic.')
8394

8495
# Pull the CSV headers
8596
self.headers = self.get_headers()
@@ -317,10 +328,35 @@ def insert_suffix(self):
317328
"""
318329
Preps the suffix to the insert query.
319330
"""
320-
if self.ignore_conflicts:
321-
return """
322-
ON CONFLICT DO NOTHING;
323-
"""
331+
if self.on_conflict:
332+
try:
333+
action = self.on_conflict['action']
334+
except KeyError:
335+
raise ValueError("Must specify an `action` when passing `on_conflict`.")
336+
if action is None:
337+
target, action = "", "DO NOTHING"
338+
elif action == 'update':
339+
try:
340+
target = self.on_conflict['target']
341+
except KeyError:
342+
raise ValueError("Must specify `target` when action == 'update'.")
343+
if target in [f.name for f in self.model._meta.fields]:
344+
target = "({0})".format(target)
345+
elif target in [c.name for c in self.model._meta.constraints]:
346+
target = "ON CONSTRAINT {0}".format(target)
347+
else:
348+
raise ValueError("`target` must be a field name or constraint name.")
349+
350+
if 'columns' in self.on_conflict:
351+
columns = ', '.join([
352+
"{0} = excluded.{0}".format(col)
353+
for col in self.on_conflict['columns']
354+
])
355+
else:
356+
raise ValueError("Must specify `columns` when action == 'update'.")
357+
358+
action = "DO UPDATE SET {0}".format(columns)
359+
return "ON CONFLICT {0} {1};".format(target, action)
324360
else:
325361
return ";"
326362

0 commit comments

Comments
 (0)