Add ability to set update values, fixes #56

Author: Photonios

docs/source/conflict_handling.rst

@@ -232,6 +232,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

@@ -104,8 +104,7 @@ def as_sql(self, *args, **kwargs):
         return append_caller_to_sql(sql), params
 
     def _prepare_query_values(self):
-        """Extra prep on query values by converting dictionaries into.
-
+        """Extra prep on query values by converting dictionaries into
         :see:HStoreValue expressions.
 
         This allows putting expressions in a dictionary. The
@@ -234,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  # type: ignore[attr-defined]
-            ]
-        )
-
         # build the conflict target, the columns to watch
         # for conflicts
         on_conflict_clause = self._build_on_conflict_clause()
@@ -254,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(
@@ -270,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_on_conflict_clause(self):
         if django.VERSION >= (2, 2):
             from django.db.models.constraints import BaseConstraint

psqlextra/query.py

@@ -2,6 +2,7 @@
 from itertools import chain
 from typing import (
     TYPE_CHECKING,
+    Any,
     Dict,
     Generic,
     Iterable,
@@ -17,6 +18,7 @@
 from django.db.models import Expression, Q, QuerySet
 from django.db.models.fields import NOT_PROVIDED
 
+from .expressions import ExcludedCol
 from .sql import PostgresInsertQuery, PostgresQuery
 from .types import ConflictAction
 
@@ -51,6 +53,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) -> "Self":  # type: ignore[valid-type, override]
         """Custom version of the standard annotate function that allows using
@@ -108,6 +111,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.
@@ -125,12 +129,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
 
@@ -293,6 +303,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.
@@ -315,6 +326,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.
         """
@@ -324,6 +340,7 @@ def upsert(
             ConflictAction.UPDATE,
             index_predicate=index_predicate,
             update_condition=update_condition,
+            update_values=update_values,
         )
 
         kwargs = {**fields, "using": using}
@@ -336,6 +353,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.
@@ -358,6 +376,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.
@@ -368,6 +391,7 @@ def upsert_and_get(
             ConflictAction.UPDATE,
             index_predicate=index_predicate,
             update_condition=update_condition,
+            update_values=update_values,
         )
 
         kwargs = {**fields, "using": using}
@@ -381,6 +405,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.
@@ -407,6 +432,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
@@ -417,7 +447,9 @@ def bulk_upsert(
             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(
@@ -505,15 +537,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.insert_on_conflict_values(objs, insert_fields, update_values)
 
         compiler = query.get_compiler(using)
         return compiler
@@ -578,13 +614,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)
@@ -595,13 +631,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,5 +1,5 @@
 from collections import OrderedDict
-from typing import Optional, Tuple
+from typing import Any, Dict, List, Optional, Tuple, Union
 
 import django
 
@@ -154,10 +154,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, insert_fields, update_fields=[]):
+        self.update_values = {}
+
+    def insert_on_conflict_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
@@ -176,12 +180,13 @@ def values(self, objs, insert_fields, update_fields=[]):
             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