Add tests for bulk upserts

Author: Photonios

psqlextra/manager.py

@@ -129,6 +129,23 @@ def on_conflict(self, fields: List[Union[str, Tuple[str]]], action):
         self.conflict_action = action
         return self
 
+    def bulk_insert(self, rows):
+        """Creates multiple new records in the database.
+
+        This allows specifying custom conflict behavior using .on_conflict().
+        If no special behavior was specified, this uses the normal Django create(..)
+
+        Arguments:
+            rows:
+                An array of dictionaries, where each dictionary
+                describes the fields to insert.
+
+        Returns:
+        """
+
+        for row in rows:
+            self.insert(**row)
+
     def insert(self, **fields):
         """Creates a new record in the database.
 

tests/benchmarks/test_insert_nothing.py

@@ -10,49 +10,45 @@
 from ..fake_model import get_fake_model
 
 
-@pytest.mark.django_db()
-class TestInsertNothing(TestCase):
-
-    @pytest.mark.benchmark()
-    @staticmethod
-    def test_insert_nothing_traditional(benchmark):
-        model = get_fake_model({
-            'field': models.CharField(max_length=255, unique=True)
-        })
-
-        random_value = str(uuid.uuid4())[:8]
-        model.objects.create(field=random_value)
-
-        def _traditional_insert(model, random_value):
-            """Performs a concurrency safe insert the
-            traditional way."""
-
-            try:
-                with transaction.atomic():
-                    return model.objects.create(field=random_value)
-            except IntegrityError:
-                return model.objects.filter(field=random_value).first()
-
-        benchmark(_traditional_insert, model, random_value)
-
-    @pytest.mark.benchmark()
-    @staticmethod
-    def test_insert_nothing_native(benchmark):
-        model = get_fake_model({
-            'field': models.CharField(max_length=255, unique=True)
-        })
-
-        random_value = str(uuid.uuid4())[:8]
-        model.objects.create(field=random_value)
-
-        def _native_insert(model, random_value):
-            """Performs a concurrency safeinsert
-            using the native PostgreSQL conflict resolution."""
-
-            return (
-                model.objects
-                .on_conflict(['field'], ConflictAction.NOTHING)
-                .insert_and_get(field=random_value)
-            )
-
-        benchmark(_native_insert, model, random_value)
+@pytest.mark.benchmark()
+def test_insert_nothing_traditional(benchmark):
+    model = get_fake_model({
+        'field': models.CharField(max_length=255, unique=True)
+    })
+
+    random_value = str(uuid.uuid4())[:8]
+    model.objects.create(field=random_value)
+
+    def _traditional_insert(model, random_value):
+        """Performs a concurrency safe insert the
+        traditional way."""
+
+        try:
+            with transaction.atomic():
+                return model.objects.create(field=random_value)
+        except IntegrityError:
+            return model.objects.filter(field=random_value).first()
+
+    benchmark(_traditional_insert, model, random_value)
+
+
+@pytest.mark.benchmark()
+def test_insert_nothing_native(benchmark):
+    model = get_fake_model({
+        'field': models.CharField(max_length=255, unique=True)
+    })
+
+    random_value = str(uuid.uuid4())[:8]
+    model.objects.create(field=random_value)
+
+    def _native_insert(model, random_value):
+        """Performs a concurrency safeinsert
+        using the native PostgreSQL conflict resolution."""
+
+        return (
+            model.objects
+            .on_conflict(['field'], ConflictAction.NOTHING)
+            .insert_and_get(field=random_value)
+        )
+
+    benchmark(_native_insert, model, random_value)

tests/benchmarks/test_upsert.py

@@ -8,51 +8,46 @@
 from ..fake_model import get_fake_model
 
 
-@pytest.mark.django_db()
 @pytest.mark.benchmark()
-class TestUpsert(TestCase):
+def test_upsert_traditional(benchmark):
+    model = get_fake_model({
+        'field': models.CharField(max_length=255, unique=True)
+    })
 
-    @pytest.mark.benchmark()
-    @staticmethod
-    def test_upsert_traditional(benchmark):
-        model = get_fake_model({
-            'field': models.CharField(max_length=255, unique=True)
-        })
+    random_value = str(uuid.uuid4())[:8]
+    model.objects.create(field=random_value)
 
-        random_value = str(uuid.uuid4())[:8]
-        model.objects.create(field=random_value)
+    def _traditional_upsert(model, random_value):
+        """Performs a concurrency safe upsert
+        the traditional way."""
 
-        def _traditional_upsert(model, random_value):
-            """Performs a concurrency safe upsert
-            the traditional way."""
+        try:
 
-            try:
+            with transaction.atomic():
+                return model.objects.create(field=random_value)
+        except IntegrityError:
+            model.objects.update(field=random_value)
+            return model.objects.get(field=random_value)
 
-                with transaction.atomic():
-                    return model.objects.create(field=random_value)
-            except IntegrityError:
-                model.objects.update(field=random_value)
-                return model.objects.get(field=random_value)
+    benchmark(_traditional_upsert, model, random_value)
 
-        benchmark(_traditional_upsert, model, random_value)
 
-    @pytest.mark.benchmark()
-    @staticmethod
-    def test_upsert_native(benchmark):
-        model = get_fake_model({
-            'field': models.CharField(max_length=255, unique=True)
-        })
+@pytest.mark.benchmark()
+def test_upsert_native(benchmark):
+    model = get_fake_model({
+        'field': models.CharField(max_length=255, unique=True)
+    })
 
-        random_value = str(uuid.uuid4())[:8]
-        model.objects.create(field=random_value)
+    random_value = str(uuid.uuid4())[:8]
+    model.objects.create(field=random_value)
 
-        def _native_upsert(model, random_value):
-            """Performs a concurrency safe upsert
-            using the native PostgreSQL upsert."""
+    def _native_upsert(model, random_value):
+        """Performs a concurrency safe upsert
+        using the native PostgreSQL upsert."""
 
-            return model.objects.upsert_and_get(
-                conflict_target=['field'],
-                fields=dict(field=random_value)
-            )
+        return model.objects.upsert_and_get(
+            conflict_target=['field'],
+            fields=dict(field=random_value)
+        )
 
-        benchmark(_native_upsert, model, random_value)
+    benchmark(_native_upsert, model, random_value)

tests/benchmarks/test_upsert_bulk.py

@@ -0,0 +1,72 @@
+import uuid
+
+import pytest
+
+from django.test import TestCase
+from django.db import models, transaction
+from django.db.utils import IntegrityError
+
+from psqlextra.query import ConflictAction
+
+from ..fake_model import get_fake_model
+
+ROW_COUNT = 10000
+
+
+@pytest.mark.benchmark()
+def test_upsert_bulk_naive(benchmark):
+    model = get_fake_model({
+        'field': models.CharField(max_length=255, unique=True)
+    })
+
+    rows = []
+    random_values = []
+    for i in range(0, ROW_COUNT):
+        random_value = str(uuid.uuid4())
+        random_values.append(random_value)
+        rows.append(model(field=random_value))
+
+    model.objects.bulk_create(rows)
+
+    def _native_upsert(model, random_values):
+        """Performs a concurrency safe upsert
+        using the native PostgreSQL upsert."""
+
+        rows = [
+            dict(field=random_value)
+            for random_value in random_values
+        ]
+
+        for row in rows:
+            model.objects.on_conflict(['field'], ConflictAction.UPDATE).insert(**row)
+
+    benchmark(_native_upsert, model, random_values)
+
+
+@pytest.mark.benchmark()
+def test_upsert_bulk_native(benchmark):
+    model = get_fake_model({
+        'field': models.CharField(max_length=255, unique=True)
+    })
+
+    rows = []
+    random_values = []
+    for i in range(0, ROW_COUNT):
+        random_value = str(uuid.uuid4())
+        random_values.append(random_value)
+        rows.append(model(field=random_value))
+
+    model.objects.bulk_create(rows)
+
+    def _native_upsert(model, random_values):
+        """Performs a concurrency safe upsert
+        using the native PostgreSQL upsert."""
+
+        rows = [
+            dict(field=random_value)
+            for random_value in random_values
+        ]
+
+        model.objects.on_conflict(['field'], ConflictAction.UPDATE).bulk_insert(rows)
+
+    benchmark(_native_upsert, model, random_values)

tests/test_on_conflict.py

@@ -387,3 +387,29 @@ def test_on_conflict_default_value_no_overwrite():
 
     assert obj1.id == obj2.id
     assert obj2.title == 'mytitle'
+
+
+def test_on_conflict_bulk():
+    """Tests whether using `on_conflict` with `insert_bulk`
+    properly works."""
+
+    model = get_fake_model({
+        'title': models.CharField(max_length=255, unique=True)
+    })
+
+    rows = [
+        dict(title='this is my title'),
+        dict(title='this is another title'),
+        dict(title='and another one')
+    ]
+
+    (
+        model.objects
+        .on_conflict(['title'], ConflictAction.UPDATE)
+        .bulk_insert(rows)
+    )
+
+    assert model.objects.all().count() == len(rows)
+
+    for index, obj in enumerate(list(model.objects.all())):
+        assert obj.title == rows[index]['title']