Allow Q objects to be used for upsert update condition and index predicate

Author: Photonios

docs/source/api_reference.rst

@@ -29,8 +29,11 @@ API Reference
 
    .. autoclass:: DateTimeEpoch
 
+   .. autoclass:: ExcludedCol
+
 .. automodule:: psqlextra.indexes
 
+   .. autoclass:: UniqueIndex
    .. autoclass:: ConditionalUniqueIndex
    .. autoclass:: CaseInsensitiveUniqueIndex
 

docs/source/conflict_handling.rst

@@ -155,35 +155,21 @@ A row level lock is acquired before evaluating the condition and proceeding with
         An expression that returns a value of type boolean. Only rows for which this expression returns true will be updated, although all rows will be locked when the ON CONFLICT DO UPDATE action is taken. Note that condition is evaluated last, after a conflict has been identified as a candidate to update.
 
 
-.. warning::
-
-    Always parameterize the input to avoid SQL injections.
-
-    Do:
-
-        .. code-block:: python
-
-            my_name = 'henk'
-            RawSQL("name != %s", (my_name,))
-
-    Not:
-
-        .. code-block:: python
-
-            RawSQL("name != " + henk, tuple())
-
-
 .. code-block:: python
 
-    from django.db.models.expressions import RawSQL
+    from psqlextra.expressions import CombinedExpression, ExcludedCol
 
     pk = (
         MyModel
         .objects
         .on_conflict(
             ['name'],
             ConflictAction.UPDATE,
-            update_condition=RawSQL("priority >= EXCLUDED.priority"),
+            update_condition=CombinedExpression(
+                MyModel._meta.get_field('priority').get_col(MyModel._meta.db_table),
+                '>',
+                ExcludedCol('priority'),
+            )
         )
         .insert(
             name='henk',
@@ -197,30 +183,18 @@ A row level lock is acquired before evaluating the condition and proceeding with
         print('condition was false-ish and no changes were made')
 
 
-When writing expressions, refer to the data you're trying to upsert with ``EXCLUDED``. Refer to the existing row by prefixing the name of the table:
-
-.. code-block:: python
-
-    RawSQL(MyModel._meta.db_table + '.mycolumn = EXCLUDED.mycolumn')
-
-You can use :meth:`~django:django.db.models.expressions.CombinedExpression` to build simple comparion expressions:
+When writing expressions, refer to the data you're trying to upsert with the :class:`psqlextra.expressions.ExcludedCol` expression.
 
+Alternatively, with Django 3.1, :class:`~django:django.db.models.Q` objects can be used instead:
 
 .. code-block:: python
 
-    from django.db.models import CombinedExpression, Col, Value
-
-    CombinedExpression(
-        MyModel._meta.get_field('name').get_col(MyModel._meta.db_table)
-        '=',
-        Col('EXCLUDED', 'name'),
-    )
+    from django.db.models import Q
+    from psqlextra.expressions import ExcludedCol
 
-    CombinedExpression(
-        MyModel._meta.get_field('active').get_col(MyModel._meta.db_table)
-        '=',
-        Value(True),
-    )
+    Q(name=ExcludedCol('name'))
+    Q(name__isnull=True)
+    Q(name__gt=ExcludedCol('priority'))
 
 
 ConflictAction.NOTHING

docs/source/expressions.rst

@@ -89,3 +89,31 @@ Use the :class:`~psqlextra.expressions.IsNotNone` expression to perform somethin
       .values_list('name', flat=True)
       .first()
    )
+
+
+Excluded column
+---------------
+
+Use the :class:`~psqlextra.expressions.ExcludedCol` expression when performing an upsert using `ON CONFLICT`_ to refer to a column/field in the data is about to be upserted.
+
+PostgreSQL keeps that data to be upserted in a special table named `EXCLUDED`. This expression is used to refer to a column in that table.
+
+.. code-block:: python
+
+    from django.db.models import Q
+    from psqlextra.expressions import ExcludedCol
+
+    (
+        MyModel
+        .objects
+        .on_conflict(
+            ['name'],
+            ConflictAction.UPDATE,
+            # translates to `priority > EXCLUDED.priority`
+            update_condition=Q(priority__gt=ExcludedCol('priority')),
+        )
+        .insert(
+            name='henk',
+            priority=1,
+        )
+    )

psqlextra/compiler.py

@@ -1,7 +1,10 @@
 from collections.abc import Iterable
+from typing import Tuple, Union
+
+import django
 
 from django.core.exceptions import SuspiciousOperation
-from django.db.models import Expression, Model
+from django.db.models import Expression, Model, Q
 from django.db.models.fields.related import RelatedField
 from django.db.models.sql.compiler import SQLInsertCompiler, SQLUpdateCompiler
 from django.db.utils import ProgrammingError
@@ -155,20 +158,19 @@ def _rewrite_insert_on_conflict(
         rewritten_sql = f"{sql} ON CONFLICT {conflict_target}"
 
         if index_predicate:
-            if isinstance(index_predicate, Expression):
-                expr_sql, expr_params = self.compile(index_predicate)
-                rewritten_sql += f" WHERE {expr_sql}"
-                params += tuple(expr_params)
-            else:
-                rewritten_sql += f" WHERE {index_predicate}"
+            expr_sql, expr_params = self._compile_expression(index_predicate)
+            rewritten_sql += f" WHERE {expr_sql}"
+            params += tuple(expr_params)
 
         rewritten_sql += f" DO {conflict_action}"
 
         if conflict_action == "UPDATE":
             rewritten_sql += f" SET {update_columns}"
 
             if update_condition:
-                expr_sql, expr_params = self.compile(update_condition)
+                expr_sql, expr_params = self._compile_expression(
+                    update_condition
+                )
                 rewritten_sql += f" WHERE {expr_sql}"
                 params += tuple(expr_params)
 
@@ -319,6 +321,27 @@ def _format_field_value(self, field_name) -> str:
             value,
         )
 
+    def _compile_expression(
+        self, expression: Union[Expression, Q, str]
+    ) -> Tuple[str, tuple]:
+        """Compiles an expression, Q object or raw SQL string into SQL and
+        tuple of parameters."""
+
+        if isinstance(expression, Q):
+            if django.VERSION < (3, 1):
+                raise SuspiciousOperation(
+                    "Q objects in psqlextra can only be used with Django 3.1 and newer"
+                )
+
+            return self.query.build_where(expression).as_sql(
+                self, self.connection
+            )
+
+        elif isinstance(expression, Expression):
+            return self.compile(expression)
+
+        return expression, tuple()
+
     def _assert_valid_field(self, field_name: str):
         """Asserts that a field with the specified name exists on the model and
         raises :see:SuspiciousOperation if it does not."""

psqlextra/expressions.py

@@ -204,3 +204,18 @@ def IsNotNone(*fields, default=None):
         default=expressions.Value(default),
         output_field=CharField(),
     )
+
+
+class ExcludedCol(expressions.Expression):
+    """References a column in PostgreSQL's special EXCLUDED column, which is
+    used in upserts to refer to the data about to be inserted/updated.
+
+    See: https://www.postgresql.org/docs/9.5/sql-insert.html#SQL-ON-CONFLICT
+    """
+
+    def __init__(self, name: str):
+        self.name = name
+
+    def as_sql(self, compiler, connection):
+        quoted_name = connection.ops.quote_name(self.name)
+        return f"EXCLUDED.{quoted_name}", tuple()

psqlextra/query.py

@@ -4,7 +4,7 @@
 
 from django.core.exceptions import SuspiciousOperation
 from django.db import connections, models, router
-from django.db.models import Expression
+from django.db.models import Expression, Q
 from django.db.models.fields import NOT_PROVIDED
 
 from .sql import PostgresInsertQuery, PostgresQuery
@@ -82,8 +82,8 @@ def on_conflict(
         self,
         fields: ConflictTarget,
         action: ConflictAction,
-        index_predicate: Optional[Union[Expression, str]] = None,
-        update_condition: Optional[Expression] = None,
+        index_predicate: Optional[Union[Expression, Q, str]] = None,
+        update_condition: Optional[Union[Expression, Q, str]] = None,
     ):
         """Sets the action to take when conflicts arise when attempting to
         insert/create a new row.
@@ -257,9 +257,9 @@ def upsert(
         self,
         conflict_target: ConflictTarget,
         fields: dict,
-        index_predicate: Optional[Union[Expression, str]] = None,
+        index_predicate: Optional[Union[Expression, Q, str]] = None,
         using: Optional[str] = None,
-        update_condition: Optional[Expression] = None,
+        update_condition: Optional[Union[Expression, Q, str]] = None,
     ) -> int:
         """Creates a new record or updates the existing one with the specified
         data.
@@ -298,9 +298,9 @@ def upsert_and_get(
         self,
         conflict_target: ConflictTarget,
         fields: dict,
-        index_predicate: Optional[Union[Expression, str]] = None,
+        index_predicate: Optional[Union[Expression, Q, str]] = None,
         using: Optional[str] = None,
-        update_condition: Optional[Expression] = None,
+        update_condition: Optional[Union[Expression, Q, str]] = None,
     ):
         """Creates a new record or updates the existing one with the specified
         data and then gets the row.
@@ -340,10 +340,10 @@ def bulk_upsert(
         self,
         conflict_target: ConflictTarget,
         rows: Iterable[Dict],
-        index_predicate: Optional[Union[Expression, str]] = None,
+        index_predicate: Optional[Union[Expression, Q, str]] = None,
         return_model: bool = False,
         using: Optional[str] = None,
-        update_condition: Optional[Expression] = None,
+        update_condition: Optional[Union[Expression, Q, str]] = None,
     ):
         """Creates a set of new records or updates the existing ones with the
         specified data.

tests/test_query_values.py

@@ -51,15 +51,16 @@ def test_query_values_list_hstore_key(model, modelobj):
     assert result[1] == modelobj.title["ar"]
 
 
+@pytest.mark.skipif(
+    django.VERSION < (2, 1), reason="requires django 2.1 or newer"
+)
 def test_query_values_hstore_key_through_fk():
     """Tests whether selecting a single key from a :see:HStoreField using the
     query set's .values() works properly when there's a foreign key
     relationship involved."""
 
     # this starting working in django 2.1
     # see: https://github.com/django/django/commit/20bab2cf9d02a5c6477d8aac066a635986e0d3f3
-    if django.VERSION < (2, 1):
-        return
 
     fmodel = get_fake_model({"name": HStoreField()})
 

tests/test_upsert.py

@@ -1,6 +1,11 @@
+import django
+import pytest
+
 from django.db import models
+from django.db.models import Q
 from django.db.models.expressions import CombinedExpression, Value
 
+from psqlextra.expressions import ExcludedCol
 from psqlextra.fields import HStoreField
 
 from .fake_model import get_fake_model
@@ -78,7 +83,7 @@ def test_upsert_explicit_pk():
 
 
 def test_upsert_with_update_condition():
-    """Tests that a custom expression can be passed as an update condition."""
+    """Tests that an expression can be used as an upsert update condition."""
 
     model = get_fake_model(
         {
@@ -96,7 +101,7 @@ def test_upsert_with_update_condition():
         update_condition=CombinedExpression(
             model._meta.get_field("active").get_col(model._meta.db_table),
             "=",
-            Value(True),
+            ExcludedCol("active"),
         ),
         fields=dict(name="joe", priority=2, active=True),
     )
@@ -122,6 +127,46 @@ def test_upsert_with_update_condition():
     assert obj1.active
 
 
+@pytest.mark.skipif(
+    django.VERSION < (3, 1), reason="requires django 3.1 or newer"
+)
+def test_upsert_with_update_condition_with_q_object():
+    """Tests that :see:Q objects can be used as an upsert update condition."""
+
+    model = get_fake_model(
+        {
+            "name": models.TextField(unique=True),
+            "priority": models.IntegerField(),
+            "active": models.BooleanField(),
+        }
+    )
+
+    obj1 = model.objects.create(name="joe", priority=1, active=False)
+
+    # should not return anything because no rows were affected
+    assert not model.objects.upsert(
+        conflict_target=["name"],
+        update_condition=Q(active=ExcludedCol("active")),
+        fields=dict(name="joe", priority=2, active=True),
+    )
+
+    obj1.refresh_from_db()
+    assert obj1.priority == 1
+    assert not obj1.active
+
+    # should return something because one row was affected
+    obj1_pk = model.objects.upsert(
+        conflict_target=["name"],
+        update_condition=Q(active=Value(False)),
+        fields=dict(name="joe", priority=2, active=True),
+    )
+
+    obj1.refresh_from_db()
+    assert obj1.pk == obj1_pk
+    assert obj1.priority == 2
+    assert obj1.active
+
+
 def test_upsert_and_get_applies_converters():
     """Tests that converters are properly applied when using upsert_and_get."""