Add support for conflict_target

Author: Photonios

docs/features.md

@@ -46,16 +46,29 @@ In order to combat this, PostgreSQL added native upserts. Also known as [`ON CON
 
 
 ## upsert
-The `upsert` method attempts to insert a row with the specified data or updates (and overwrites) the duplicate row, and then returns the primary key of the row that was created/updated:
+The `upsert` method attempts to insert a row with the specified data or updates (and overwrites) the duplicate row, and then returns the primary key of the row that was created/updated.
+
+Upserts work by catching conflicts. PostgreSQL requires to know which conflicts to react to. You have to specify the name of the column which's constraint you want to react to. This is specified in the `conflict_target` field. If the constraint you're trying to react to consists of multiple columns, specify multiple columns.
 
     from django.db import models
     from psqlextra.models import PostgresModel
 
     class MyModel(PostgresModel):
         myfield = models.CharField(max_length=255, unique=True)
 
-    id1 = MyModel.objects.upsert(myfield='beer')
-    id2 = MyModel.objects.upsert(myfield='beer')
+    id1 = MyModel.objects.upsert(
+        conflict_target=['myfield'],
+        fields=dict(
+            myfield='beer'
+        )
+    )
+
+    id2 = MyModel.objects.upsert(
+        conflict_target=['myfield'],
+        fields=dict(
+            myfield='beer'
+        )
+    )
 
     assert id1 == id2
 
@@ -73,4 +86,18 @@ Note that a single call to `upsert` results in a single `INSERT INTO ... ON CONF
     obj1 = MyModel.objects.upsert_and_get(myfield='beer')
     obj2 = MyModel.objects.upsert_and_get(myfield='beer')
 
+    obj1 = MyModel.objects.upsert_and_get(
+        conflict_target=['myfield'],
+        fields=dict(
+            myfield='beer'
+        )
+    )
+
+    obj2 = MyModel.objects.upsert_and_get(
+        conflict_target=['myfield'],
+        fields=dict(
+            myfield='beer'
+        )
+    )
+
     assert obj1.id == obj2.id

psqlextra/compiler.py

@@ -1,7 +1,5 @@
-from django.db.models.sql.compiler import SQLInsertCompiler
-
-from .fields import HStoreField
 from django.core.exceptions import SuspiciousOperation
+from django.db.models.sql.compiler import SQLInsertCompiler
 
 
 class PostgresSQLUpsertCompiler(SQLInsertCompiler):
@@ -31,10 +29,13 @@ def execute_sql(self, return_id=False):
         # return the primary key, which is stored in
         # the first column that is returned
         if return_id:
-            return rows[0][0]
+            return dict(id=rows[0][0])
 
-        # return the entire row instead
-        return rows[0]
+        # create a mapping between column names and column value
+        return {
+            column.name: rows[0][column_index]
+            for column_index, column in enumerate(cursor.description)
+        }
 
     def _rewrite_insert(self, sql, returning='id'):
         """Rewrites a formed SQL INSERT query to include
@@ -54,63 +55,58 @@ def _rewrite_insert(self, sql, returning='id'):
         """
         qn = self.connection.ops.quote_name
 
-        # ON CONFLICT requires a list of columns to operate on, form
-        # a list of columns to pass in
-        unique_columns = ', '.join(self._get_unique_columns())
-        if len(unique_columns) == 0:
-            raise SuspiciousOperation((
-                'You\'re trying to do a upsert on a table that '
-                'doesn\'t have any unique columns.'
-            ))
-
         # construct a list of columns to update when there's a conflict
         update_columns = ', '.join([
             '{0} = EXCLUDED.{0}'.format(qn(field.column))
             for field in self.query.update_fields
         ])
 
+        # build the conflict target, the columns to watch
+        # for conflict basically
+        conflict_target = self._build_conflict_target()
+
         # form the new sql query that does the insert
         new_sql = (
-            '{insert} ON CONFLICT ({unique_columns}) '
+            '{insert} ON CONFLICT ({conflict_target}) '
             'DO UPDATE SET {update_columns} RETURNING {returning}'
         ).format(
             insert=sql,
-            unique_columns=unique_columns,
+            conflict_target=conflict_target,
             update_columns=update_columns,
             returning=returning
         )
 
         return new_sql
 
-    def _get_unique_columns(self):
-        """Gets a list of columns that are marked as 'UNIQUE'.
-
-        This is used in the ON CONFLICT clause. This also
-        works for :see:HStoreField."""
+    def _build_conflict_target(self):
+        """Builds the `conflict_target` for the ON CONFLICT
+        clause."""
 
         qn = self.connection.ops.quote_name
-        unique_columns = []
+        conflict_target = []
 
-        for field in self.query.fields:
-            if field.unique is True:
-                unique_columns.append(qn(field.column))
+        if not isinstance(self.query.conflict_target, list):
+            raise SuspiciousOperation((
+                '%s is not a valid conflict target, specify '
+                'a list of column names, or tuples with column '
+                'names and hstore key.'
+            ) % str(self.query.conflict_target))
+
+        for field in self.query.conflict_target:
+            if isinstance(field, str):
+                conflict_target.append(qn(field))
                 continue
 
-            # we must also go into possible tuples since those
-            # are used to indicate "unique together"
-            if isinstance(field, HStoreField):
-                uniqueness = getattr(field, 'uniqueness', None)
-                if not uniqueness:
-                    continue
-                for key in uniqueness:
-                    if isinstance(key, tuple):
-                        for sub_key in key:
-                            unique_columns.append(
-                                '(%s->\'%s\')' % (qn(field.column), sub_key))
-                    else:
-                        unique_columns.append(
-                            '(%s->\'%s\')' % (qn(field.column), key))
-
+            if isinstance(field, tuple):
+                field, key = field
+                conflict_target.append(
+                    '(%s -> \'%s\')' % (qn(field), key))
                 continue
 
-        return unique_columns
+            raise SuspiciousOperation((
+                '%s is not a valid conflict target, specify '
+                'a list of column names, or tuples with column '
+                'names and hstore key.'
+            ) % str(field))
+
+        return ','.join(conflict_target)

psqlextra/manager.py

@@ -1,13 +1,12 @@
-import copy
+from typing import List, Dict
 
-from django.db import models
+import django
 from django.conf import settings
-from django.db.models.sql import InsertQuery
 from django.core.exceptions import ImproperlyConfigured
-import django
+from django.db import models
 
-from .query import PostgresUpsertQuery
 from .compiler import PostgresSQLUpsertCompiler
+from .query import PostgresUpsertQuery
 
 
 class PostgresManager(models.Manager):
@@ -28,43 +27,51 @@ def __init__(self, *args, **kwargs):
                 'the \'psqlextra.backend\'. Set DATABASES.ENGINE.'
             ) % db_backend)
 
-    def upsert(self, **kwargs) -> int:
+    def upsert(self, conflict_target: List, fields: Dict) -> int:
         """Creates a new record or updates the existing one
         with the specified data.
 
         Arguments:
-            kwargs:
+            conflict_target:
+                Fields to pass into the ON CONFLICT clause.
+
+            fields:
                 Fields to insert/update.
 
         Returns:
             The primary key of the row that was created/updated.
         """
 
-        compiler = self._build_upsert_compiler(kwargs)
-        return compiler.execute_sql(return_id=True)
+        compiler = self._build_upsert_compiler(conflict_target, fields)
+        return compiler.execute_sql(return_id=True)['id']
 
-    def upsert_and_get(self, **kwargs):
+    def upsert_and_get(self, conflict_target: List, fields: Dict):
         """Creates a new record or updates the existing one
         with the specified data and then gets the row.
 
         Arguments:
-            kwargs:
+            conflict_target:
+                Fields to pass into the ON CONFLICT clause.
+
+            fields:
                 Fields to insert/update.
 
         Returns:
             The model instance representing the row
             that was created/updated.
         """
 
-        compiler = self._build_upsert_compiler(kwargs)
-        row = compiler.execute_sql(return_id=False)
-        field_names = [f.name for f in self.model._meta.concrete_fields]
-        return self.model.from_db(self.db, field_names, row)
+        compiler = self._build_upsert_compiler(conflict_target, fields)
+        column_data = compiler.execute_sql(return_id=False)
+        return self.model(**column_data)
 
-    def _build_upsert_compiler(self, kwargs):
+    def _build_upsert_compiler(self, conflict_target: List, kwargs):
         """Builds the SQL compiler for a insert/update query.
 
         Arguments:
+            conflict_target:
+                Fields to pass into the ON CONFLICT clause.
+
             kwargs:
                 Field values.
 
@@ -87,6 +94,7 @@ def _build_upsert_compiler(self, kwargs):
 
         # build a normal insert query
         query = PostgresUpsertQuery(self.model)
+        query.conflict_target = conflict_target
         query.values([obj], insert_fields, update_fields)
 
         # use the upsert query compiler to transform the insert

psqlextra/query.py

@@ -1,3 +1,5 @@
+from typing import List
+
 from django.db.models.sql import InsertQuery
 
 
@@ -10,8 +12,9 @@ def __init__(self, *args, **kwargs):
         super(PostgresUpsertQuery, self).__init__(*args, **kwargs)
 
         self.update_fields = []
+        self.conflict_target = []
 
-    def values(self, objs, insert_fields, update_fields):
+    def values(self, objs: List, insert_fields: List, update_fields: List):
         """Sets the values to be used in this query.
 
         Insert fields are fields that are definitely

tests/benchmarks/test_upsert.py

@@ -23,7 +23,10 @@ def _native_upsert(model, random_value):
     """Performs a concurrency safe upsert
     using the native PostgreSQL upsert."""
 
-    return model.objects.upsert_and_get(field=random_value)
+    return model.objects.upsert_and_get(
+        conflict_target=['field'],
+        fields=dict(field=random_value)
+    )
 
 
 @pytest.mark.django_db()