Merge pull request #29 from SectorLabs/make-conditional-indexes-usable

Make ConditionalUniqueIndexes actually usable.

Author: Photonios

psqlextra/compiler.py

@@ -103,15 +103,26 @@ def _rewrite_insert_update(self, sql, params, returning):
         # for conflicts
         conflict_target = self._build_conflict_target()
 
+        index_predicate = self.query.index_predicate
+
+        sql_template = (
+            '{insert} ON CONFLICT {conflict_target} DO UPDATE '
+            'SET {update_columns} RETURNING {returning}'
+        )
+
+        if index_predicate:
+            sql_template = (
+                '{insert} ON CONFLICT {conflict_target} WHERE {index_predicate} DO UPDATE '
+                'SET {update_columns} RETURNING {returning}'
+            )
+
         return (
-            (
-                '{insert} ON CONFLICT {conflict_target} DO UPDATE'
-                ' SET {update_columns} RETURNING {returning}'
-            ).format(
+            sql_template.format(
                 insert=sql,
                 conflict_target=conflict_target,
                 update_columns=update_columns,
-                returning=returning
+                returning=returning,
+                index_predicate=index_predicate,
             ),
             params
         )

psqlextra/manager/manager.py

@@ -26,6 +26,7 @@ def __init__(self, model=None, query=None, using=None, hints=None):
 
         self.conflict_target = None
         self.conflict_action = None
+        self.index_predicate = None
 
     def annotate(self, **annotations):
         """Custom version of the standard annotate function
@@ -113,7 +114,7 @@ def update(self, **fields):
         # affected, let's do the same
         return len(rows)
 
-    def on_conflict(self, fields: List[Union[str, Tuple[str]]], action):
+    def on_conflict(self, fields: List[Union[str, Tuple[str]]], action, index_predicate: str=None):
         """Sets the action to take when conflicts arise when attempting
         to insert/create a new row.
 
@@ -123,10 +124,16 @@ def on_conflict(self, fields: List[Union[str, Tuple[str]]], action):
 
             action:
                 The action to take when the conflict occurs.
+
+            index_predicate:
+                The index predicate to satisfy an arbiter partial index (i.e. what partial index to use for checking
+                conflicts)
         """
 
         self.conflict_target = fields
         self.conflict_action = action
+        self.index_predicate = index_predicate
+
         return self
 
     def bulk_insert(self, rows):
@@ -216,7 +223,7 @@ def insert_and_get(self, **fields):
 
         return self.model(**model_init_fields)
 
-    def upsert(self, conflict_target: List, fields: Dict) -> int:
+    def upsert(self, conflict_target: List, fields: Dict, index_predicate: str=None) -> int:
         """Creates a new record or updates the existing one
         with the specified data.
 
@@ -227,11 +234,15 @@ def upsert(self, conflict_target: List, fields: Dict) -> int:
             fields:
                 Fields to insert/update.
 
+            index_predicate:
+                The index predicate to satisfy an arbiter partial index (i.e. what partial index to use for checking
+                conflicts)
+
         Returns:
             The primary key of the row that was created/updated.
         """
 
-        self.on_conflict(conflict_target, ConflictAction.UPDATE)
+        self.on_conflict(conflict_target, ConflictAction.UPDATE, index_predicate)
         return self.insert(**fields)
 
     def upsert_and_get(self, conflict_target: List, fields: Dict):
@@ -307,6 +318,7 @@ def _build_insert_compiler(self, rows: List[Dict]):
         query = PostgresInsertQuery(self.model)
         query.conflict_action = self.conflict_action
         query.conflict_target = self.conflict_target
+        query.index_predicate = self.index_predicate
         query.values(objs, insert_fields, update_fields)
 
         # use the postgresql insert query compiler to transform the insert
@@ -466,7 +478,7 @@ def on_conflict(self, fields: List[Union[str, Tuple[str]]], action):
         """
         return self.get_queryset().on_conflict(fields, action)
 
-    def upsert(self, conflict_target: List, fields: Dict) -> int:
+    def upsert(self, conflict_target: List, fields: Dict, index_predicate: str=None) -> int:
         """Creates a new record or updates the existing one
         with the specified data.
 
@@ -477,11 +489,14 @@ def upsert(self, conflict_target: List, fields: Dict) -> int:
             fields:
                 Fields to insert/update.
 
+            index_predicate:
+                The index predicate to satisfy an arbiter partial index.
+
         Returns:
             The primary key of the row that was created/updated.
         """
 
-        return self.get_queryset().upsert(conflict_target, fields)
+        return self.get_queryset().upsert(conflict_target, fields, index_predicate)
 
     def upsert_and_get(self, conflict_target: List, fields: Dict):
         """Creates a new record or updates the existing one

tests/test_conditional_unique_index.py

@@ -6,6 +6,8 @@
 from django.db import models, IntegrityError, transaction
 from django.db.migrations import AddIndex, CreateModel
 
+from .util import get_fake_model
+
 
 def test_deconstruct():
     """Tests whether the :see:HStoreField's deconstruct()
@@ -73,3 +75,82 @@ def test_migrations():
         Model.objects.create(id=1, name=None, other_name="other_name")
         with pytest.raises(IntegrityError):
             Model.objects.create(id=2, name=None, other_name="other_name")
+
+
+def test_upserting():
+    """Tests upserting respects the :see:ConditionalUniqueIndex rules"""
+    model = get_fake_model(
+        fields={
+            'a': models.IntegerField(),
+            'b': models.IntegerField(null=True),
+            'c': models.IntegerField(),
+        },
+        meta_options={
+            'indexes': [
+                ConditionalUniqueIndex(
+                    fields=['a', 'b'],
+                    condition='"b" IS NOT NULL'
+                ),
+                ConditionalUniqueIndex(
+                    fields=['a'],
+                    condition='"b" IS NULL'
+                )
+            ]
+        }
+    )
+
+    model.objects.upsert(conflict_target=['a'], index_predicate='"b" IS NULL', fields=dict(a=1, c=1))
+    assert model.objects.all().count() == 1
+    assert model.objects.filter(a=1, c=1).count() == 1
+
+    model.objects.upsert(conflict_target=['a'], index_predicate='"b" IS NULL', fields=dict(a=1, c=2))
+    assert model.objects.all().count() == 1
+    assert model.objects.filter(a=1, c=1).count() == 0
+    assert model.objects.filter(a=1, c=2).count() == 1
+
+    model.objects.upsert(conflict_target=['a', 'b'], index_predicate='"b" IS NOT NULL', fields=dict(a=1, b=1, c=1))
+    assert model.objects.all().count() == 2
+    assert model.objects.filter(a=1, c=2).count() == 1
+    assert model.objects.filter(a=1, b=1, c=1).count() == 1
+
+    model.objects.upsert(conflict_target=['a', 'b'], index_predicate='"b" IS NOT NULL', fields=dict(a=1, b=1, c=2))
+    assert model.objects.all().count() == 2
+    assert model.objects.filter(a=1, c=1).count() == 0
+    assert model.objects.filter(a=1, b=1, c=2).count() == 1
+
+
+def test_inserting():
+    """Tests inserting respects the :see:ConditionalUniqueIndex rules"""
+
+    model = get_fake_model(
+        fields={
+            'a': models.IntegerField(),
+            'b': models.IntegerField(null=True),
+            'c': models.IntegerField(),
+        },
+        meta_options={
+            'indexes': [
+                ConditionalUniqueIndex(
+                    fields=['a', 'b'],
+                    condition='"b" IS NOT NULL'
+                ),
+                ConditionalUniqueIndex(
+                    fields=['a'],
+                    condition='"b" IS NULL'
+                )
+            ]
+        }
+    )
+
+    model.objects.create(a=1, c=1)
+    with transaction.atomic():
+        with pytest.raises(IntegrityError):
+            model.objects.create(a=1, c=2)
+    model.objects.create(a=2, c=1)
+
+    model.objects.create(a=1, b=1, c=1)
+    with transaction.atomic():
+        with pytest.raises(IntegrityError):
+            model.objects.create(a=1, b=1, c=2)
+
+    model.objects.create(a=1, b=2, c=1)