Add ability to set update values, fixes #56

Author: Photonios

docs/source/conflict_handling.rst

@@ -197,6 +197,42 @@ Alternatively, with Django 3.1 or newer, :class:`~django:django.db.models.Q` obj
     Q(name__gt=ExcludedCol('priority'))
 
 
+Update values
+"""""""""""""
+
+Optionally, the fields to update can be overriden. The default is to update the same fields that were specified in the rows to insert.
+
+Refer to the insert values using the :class:`psqlextra.expressions.ExcludedCol` expression which translates to PostgreSQL's ``EXCLUDED.<column>`` expression. All expressions and features that can be used with Django's :meth:`~django:django.db.models.query.QuerySet.update` can be used here.
+
+.. warning::
+
+   Specifying an empty ``update_values`` (``{}``) will transform the query into :attr:`~psqlextra.types.ConflictAction.NOTHING`. Only ``None`` makes the default behaviour kick in of updating all fields that were specified.
+
+.. code-block:: python
+
+    from django.db.models import F
+
+    from psqlextra.expressions import ExcludedCol
+
+    (
+        MyModel
+        .objects
+        .on_conflict(
+            ['name'],
+            ConflictAction.UPDATE,
+            update_values=dict(
+                name=ExcludedCol('name'),
+                count=F('count') + 1,
+            ),
+        )
+        .insert(
+            name='henk',
+            count=0,
+        )
+    )
+
+
+
 ConflictAction.NOTHING
 **********************
 

psqlextra/compiler.py

@@ -233,13 +233,6 @@ def _rewrite_insert_on_conflict(
         """Rewrites a normal SQL INSERT query to add the 'ON CONFLICT'
         clause."""
 
-        update_columns = ", ".join(
-            [
-                "{0} = EXCLUDED.{0}".format(self.qn(field.column))
-                for field in self.query.update_fields
-            ]
-        )
-
         # build the conflict target, the columns to watch
         # for conflicts
         conflict_target = self._build_conflict_target()
@@ -253,10 +246,21 @@ def _rewrite_insert_on_conflict(
             rewritten_sql += f" WHERE {expr_sql}"
             params += tuple(expr_params)
 
+        # Fallback in case the user didn't specify any update values. We can still
+        # make the query work if we switch to ConflictAction.NOTHING
+        if (
+            conflict_action == ConflictAction.UPDATE.value
+            and not self.query.update_values
+        ):
+            conflict_action = ConflictAction.NOTHING
+
         rewritten_sql += f" DO {conflict_action}"
 
-        if conflict_action == "UPDATE":
-            rewritten_sql += f" SET {update_columns}"
+        if conflict_action == ConflictAction.UPDATE.value:
+            set_sql, sql_params = self._build_set_statement()
+
+            rewritten_sql += f" SET {set_sql}"
+            params += sql_params
 
             if update_condition:
                 expr_sql, expr_params = self._compile_expression(
@@ -269,6 +273,23 @@ def _rewrite_insert_on_conflict(
 
         return (rewritten_sql, params)
 
+    def _build_set_statement(self) -> Tuple[str, tuple]:
+        """Builds the SET statement for the ON CONFLICT DO UPDATE clause.
+
+        This uses the update compiler to provide full compatibility with
+        the standard Django's `update(...)`.
+        """
+
+        # Local import to work around the circular dependency between
+        # the compiler and the queries.
+        from .sql import PostgresUpdateQuery
+
+        query = self.query.chain(PostgresUpdateQuery)
+        query.add_update_values(self.query.update_values)
+
+        sql, params = query.get_compiler(self.connection.alias).as_sql()
+        return sql.split("SET")[1].split(" WHERE")[0], tuple(params)
+
     def _build_conflict_target(self):
         """Builds the `conflict_target` for the ON CONFLICT clause."""
 

psqlextra/query.py

@@ -1,12 +1,13 @@
 from collections import OrderedDict
 from itertools import chain
-from typing import Dict, Iterable, List, Optional, Tuple, Union
+from typing import Any, Dict, Iterable, List, Optional, Tuple, Union
 
 from django.core.exceptions import SuspiciousOperation
 from django.db import connections, models, router
 from django.db.models import Expression, Q
 from django.db.models.fields import NOT_PROVIDED
 
+from .expressions import ExcludedCol
 from .sql import PostgresInsertQuery, PostgresQuery
 from .types import ConflictAction
 
@@ -27,6 +28,7 @@ def __init__(self, model=None, query=None, using=None, hints=None):
         self.conflict_action = None
         self.conflict_update_condition = None
         self.index_predicate = None
+        self.update_values = None
 
     def annotate(self, **annotations):
         """Custom version of the standard annotate function that allows using
@@ -84,6 +86,7 @@ def on_conflict(
         action: ConflictAction,
         index_predicate: Optional[Union[Expression, Q, str]] = None,
         update_condition: Optional[Union[Expression, Q, str]] = None,
+        update_values: Optional[Dict[str, Union[Any, Expression]]] = None,
     ):
         """Sets the action to take when conflicts arise when attempting to
         insert/create a new row.
@@ -101,12 +104,18 @@ def on_conflict(
 
             update_condition:
                 Only update if this SQL expression evaluates to true.
+
+            update_values:
+                Optionally, values/expressions to use when rows
+                conflict. If not specified, all columns specified
+                in the rows are updated with the values you specified.
         """
 
         self.conflict_target = fields
         self.conflict_action = action
         self.conflict_update_condition = update_condition
         self.index_predicate = index_predicate
+        self.update_values = update_values
 
         return self
 
@@ -260,6 +269,7 @@ def upsert(
         index_predicate: Optional[Union[Expression, Q, str]] = None,
         using: Optional[str] = None,
         update_condition: Optional[Union[Expression, Q, str]] = None,
+        update_values: Optional[Dict[str, Union[Any, Expression]]] = None,
     ) -> int:
         """Creates a new record or updates the existing one with the specified
         data.
@@ -282,6 +292,11 @@ def upsert(
             update_condition:
                 Only update if this SQL expression evaluates to true.
 
+            update_values:
+                Optionally, values/expressions to use when rows
+                conflict. If not specified, all columns specified
+                in the rows are updated with the values you specified.
+
         Returns:
             The primary key of the row that was created/updated.
         """
@@ -291,6 +306,7 @@ def upsert(
             ConflictAction.UPDATE,
             index_predicate=index_predicate,
             update_condition=update_condition,
+            update_values=update_values,
         )
         return self.insert(**fields, using=using)
 
@@ -301,6 +317,7 @@ def upsert_and_get(
         index_predicate: Optional[Union[Expression, Q, str]] = None,
         using: Optional[str] = None,
         update_condition: Optional[Union[Expression, Q, str]] = None,
+        update_values: Optional[Dict[str, Union[Any, Expression]]] = None,
     ):
         """Creates a new record or updates the existing one with the specified
         data and then gets the row.
@@ -323,6 +340,11 @@ def upsert_and_get(
             update_condition:
                 Only update if this SQL expression evaluates to true.
 
+            update_values:
+                Optionally, values/expressions to use when rows
+                conflict. If not specified, all columns specified
+                in the rows are updated with the values you specified.
+
         Returns:
             The model instance representing the row
             that was created/updated.
@@ -333,6 +355,7 @@ def upsert_and_get(
             ConflictAction.UPDATE,
             index_predicate=index_predicate,
             update_condition=update_condition,
+            update_values=update_values,
         )
         return self.insert_and_get(**fields, using=using)
 
@@ -344,6 +367,7 @@ def bulk_upsert(
         return_model: bool = False,
         using: Optional[str] = None,
         update_condition: Optional[Union[Expression, Q, str]] = None,
+        update_values: Optional[Dict[str, Union[Any, Expression]]] = None,
     ):
         """Creates a set of new records or updates the existing ones with the
         specified data.
@@ -370,6 +394,11 @@ def bulk_upsert(
             update_condition:
                 Only update if this SQL expression evaluates to true.
 
+            update_values:
+                Optionally, values/expressions to use when rows
+                conflict. If not specified, all columns specified
+                in the rows are updated with the values you specified.
+
         Returns:
             A list of either the dicts of the rows upserted, including the pk or
             the models of the rows upserted
@@ -386,7 +415,9 @@ def is_empty(r):
             ConflictAction.UPDATE,
             index_predicate=index_predicate,
             update_condition=update_condition,
+            update_values=update_values,
         )
+
         return self.bulk_insert(rows, return_model, using=using)
 
     def _create_model_instance(
@@ -474,15 +505,19 @@ def _build_insert_compiler(
             )
 
         # get the fields to be used during update/insert
-        insert_fields, update_fields = self._get_upsert_fields(first_row)
+        insert_fields, update_values = self._get_upsert_fields(first_row)
+
+        # allow the user to override what should happen on update
+        if self.update_values is not None:
+            update_values = self.update_values
 
         # build a normal insert query
         query = PostgresInsertQuery(self.model)
         query.conflict_action = self.conflict_action
         query.conflict_target = self.conflict_target
         query.conflict_update_condition = self.conflict_update_condition
         query.index_predicate = self.index_predicate
-        query.values(objs, insert_fields, update_fields)
+        query.values(objs, insert_fields, update_values)
 
         compiler = query.get_compiler(using)
         return compiler
@@ -547,13 +582,13 @@ def _get_upsert_fields(self, kwargs):
 
         model_instance = self.model(**kwargs)
         insert_fields = []
-        update_fields = []
+        update_values = {}
 
         for field in model_instance._meta.local_concrete_fields:
             has_default = field.default != NOT_PROVIDED
             if field.name in kwargs or field.column in kwargs:
                 insert_fields.append(field)
-                update_fields.append(field)
+                update_values[field.name] = ExcludedCol(field.column)
                 continue
             elif has_default:
                 insert_fields.append(field)
@@ -564,13 +599,13 @@ def _get_upsert_fields(self, kwargs):
             # instead of a concrete field, we have to handle that
             if field.primary_key is True and "pk" in kwargs:
                 insert_fields.append(field)
-                update_fields.append(field)
+                update_values[field.name] = ExcludedCol(field.column)
                 continue
 
             if self._is_magical_field(model_instance, field, is_insert=True):
                 insert_fields.append(field)
 
             if self._is_magical_field(model_instance, field, is_insert=False):
-                update_fields.append(field)
+                update_values[field.name] = ExcludedCol(field.column)
 
-        return insert_fields, update_fields
+        return insert_fields, update_values

psqlextra/sql.py

@@ -1,11 +1,11 @@
 from collections import OrderedDict
-from typing import List, Optional, Tuple
+from typing import Any, Dict, List, Optional, Tuple, Union
 
 import django
 
 from django.core.exceptions import SuspiciousOperation
 from django.db import connections, models
-from django.db.models import sql
+from django.db.models import Expression, sql
 from django.db.models.constants import LOOKUP_SEP
 
 from .compiler import PostgresInsertOnConflictCompiler
@@ -148,10 +148,14 @@ def __init__(self, *args, **kwargs):
         self.conflict_action = ConflictAction.UPDATE
         self.conflict_update_condition = None
         self.index_predicate = None
-
-        self.update_fields = []
-
-    def values(self, objs: List, insert_fields: List, update_fields: List = []):
+        self.update_values = {}
+
+    def values(
+        self,
+        objs: List,
+        insert_fields: List,
+        update_values: Dict[str, Union[Any, Expression]] = [],
+    ):
         """Sets the values to be used in this query.
 
         Insert fields are fields that are definitely
@@ -170,12 +174,13 @@ def values(self, objs: List, insert_fields: List, update_fields: List = []):
             insert_fields:
                 The fields to use in the INSERT statement
 
-            update_fields:
-                The fields to only use in the UPDATE statement.
+            update_values:
+                Expressions/values to use when a conflict
+                occurs and an UPDATE is performed.
         """
 
         self.insert_values(insert_fields, objs, raw=False)
-        self.update_fields = update_fields
+        self.update_values = update_values
 
     def get_compiler(self, using=None, connection=None):
         if using:

psqlextra/types.py

@@ -28,6 +28,9 @@ class ConflictAction(Enum):
     def all(cls) -> List["ConflictAction"]:
         return [choice for choice in cls]
 
+    def __str__(self) -> str:
+        return self.value
+
 
 class PostgresPartitioningMethod(StrEnum):
     """Methods of partitioning supported by PostgreSQL 11.x native support for