Improve data import dry_run functionality

Author: jolpica

jolpica/formula_one/importer/importer.py

@@ -19,18 +19,6 @@ def deserialise(self, data: dict) -> DeserialisationResult:
         deserialiser = self.factory.get_deserialiser(data["object_type"])
         return deserialiser.deserialise(data)
 
-    def get_object_priority(self, object_dict: dict) -> int:
-        object_type = object_dict.get("object_type")
-
-        if object_type in self.OBJECT_TYPE_PRIORITY:
-            return self.OBJECT_TYPE_PRIORITY.index(object_type)
-        else:
-            return len(self.OBJECT_TYPE_PRIORITY)
-
-    def get_model_import_priority(self, import_key: ModelImport) -> int:
-        object_type = import_key.model_class.__name__
-        return self.get_object_priority({"object_type": object_type})
-
     def deserialise_all(self, data: list[dict]) -> DeserialisationResult:
         prioritised_data = sorted(enumerate(data), key=lambda x: self.get_object_priority(x[1]))
         indexed_results = [(i, self.deserialise(item)) for i, item in prioritised_data]
@@ -51,11 +39,52 @@ def deserialise_all(self, data: list[dict]) -> DeserialisationResult:
             errors=errors,
         )
 
-    def save_deserialisation_result_to_db(self, result: DeserialisationResult):
-        prioritised_items = sorted(result.instances.items(), key=lambda x: self.get_model_import_priority(x[0]))
+    @classmethod
+    def get_object_priority(cls, object_dict: dict) -> int:
+        object_type = object_dict.get("object_type")
+
+        if object_type in cls.OBJECT_TYPE_PRIORITY:
+            return cls.OBJECT_TYPE_PRIORITY.index(object_type)
+        else:
+            return len(cls.OBJECT_TYPE_PRIORITY)
+
+    @classmethod
+    def get_model_import_priority(cls, import_key: ModelImport) -> int:
+        object_type = import_key.model_class.__name__
+        return cls.get_object_priority({"object_type": object_type})
+
+    @classmethod
+    def save_deserialisation_result_to_db(cls, result: DeserialisationResult) -> dict:
+        prioritised_items = sorted(result.instances.items(), key=lambda x: cls.get_model_import_priority(x[0]))
+
+        import_stats = {"updated_count": 0, "created_count": 0, "models": {}}
+
         for model_import, instances in prioritised_items:
+            model_name = model_import.model_class.__name__
+            if model_name not in import_stats["models"]:
+                import_stats["models"][model_name] = {
+                    "updated_count": 0,
+                    "created_count": 0,
+                    "updated": [],
+                    "created": [],
+                }
+
             for ins in instances:
-                ins.id = model_import.model_class.objects.update_or_create(  # type: ignore[attr-defined]
+                updated_ins, is_created = model_import.model_class.objects.update_or_create(  # type: ignore[attr-defined]
                     **{field: getattr(ins, field) for field in model_import.unique_fields},
                     defaults={field: getattr(ins, field) for field in model_import.update_fields},
-                )[0].id
+                )
+                ins.pk = updated_ins.pk
+
+                if is_created:
+                    import_stats["created_count"] += 1
+                    import_stats["models"][model_name]["created_count"] += 1
+                    import_stats["models"][model_name]["created"].append(updated_ins.pk)
+                else:
+                    import_stats["updated_count"] += 1
+                    import_stats["models"][model_name]["updated_count"] += 1
+                    import_stats["models"][model_name]["updated"].append(updated_ins.pk)
+
+        import_stats["total_count"] = import_stats["created_count"] + import_stats["updated_count"]
+
+        return import_stats

jolpica_api/data_import/admin.py

@@ -5,7 +5,10 @@
 
 class DataImportLogAdmin(admin.ModelAdmin):
     def __init__(self, model, admin_site):
-        self.list_display = [field.name for field in model._meta.fields if field.name != "updated_records"]
+        self.list_display = [
+            field.name for field in model._meta.fields if field.name not in {"import_result", "errors"}
+        ]
+        self.list_filter = ["dry_run", "error_type", "is_success"]
         super().__init__(model, admin_site)
 
 

jolpica_api/data_import/migrations/0002_dataimportlog_description.py jolpica_api/data_import/migrations/0002_rename_updated_records_dataimportlog_import_result_and_more.py

@@ -1,4 +1,4 @@
-# Generated by Django 5.1.6 on 2025-03-03 21:29
+# Generated by Django 5.1.6 on 2025-03-03 23:08
 
 from django.db import migrations, models
 
@@ -10,10 +10,20 @@ class Migration(migrations.Migration):
     ]
 
     operations = [
+        migrations.RenameField(
+            model_name="dataimportlog",
+            old_name="updated_records",
+            new_name="import_result",
+        ),
         migrations.AddField(
             model_name="dataimportlog",
             name="description",
             field=models.CharField(default="", max_length=255),
             preserve_default=False,
         ),
+        migrations.AddField(
+            model_name="dataimportlog",
+            name="dry_run",
+            field=models.BooleanField(default=True),
+        ),
     ]

jolpica_api/data_import/models.py

@@ -8,9 +8,10 @@ class DataImportLog(models.Model):
     id = models.BigAutoField(primary_key=True)
     description = models.CharField(max_length=255)
     user = models.ForeignKey(User, on_delete=models.SET_NULL, null=True)
+    dry_run = models.BooleanField(default=True)
     is_success = models.BooleanField(default=False)
     completed_at = models.DateTimeField(auto_now_add=True)
     total_records = models.PositiveIntegerField(null=True)
-    updated_records = models.JSONField(null=True)
+    import_result = models.JSONField(null=True)
     error_type = models.CharField(max_length=255, null=True)
     errors = models.JSONField(null=True)

jolpica_api/data_import/tests/test_views.py

@@ -216,8 +216,12 @@ def test_successful_import(client):
     assert DataImportLog.objects.count() == 1
     log = DataImportLog.objects.first()
     assert log.completed_at is not None
+    assert not log.dry_run
     assert log.error_type is None
-    assert log.updated_records == {"Driver": [831]}
+    assert log.import_result["created_count"] == 0
+    assert log.import_result["updated_count"] == 1
+    assert list(log.import_result["models"].keys()) == ["Driver"]
+    assert log.import_result["models"]["Driver"]["updated"] == [831]
     assert log.is_success
     assert log.error_type is None
     assert log.errors is None
@@ -242,6 +246,7 @@ def test_validation_error_has_logs(client):
     assert DataImportLog.objects.count() == 1
     log = DataImportLog.objects.first()
     assert not log.is_success
+    assert not log.dry_run
     assert log.error_type == "VALIDATION"
     assert log.errors[0]["type"] == "missing"
 
@@ -265,6 +270,7 @@ def test_deserialisation_error_has_log(client):
     assert DataImportLog.objects.count() == 1
     log = DataImportLog.objects.first()
     assert not log.is_success
+    assert not log.dry_run
     assert log.error_type == "DESERIALISATION"
     assert response.json()["errors"][0] == {
         "index": 0,
@@ -276,27 +282,31 @@ def test_deserialisation_error_has_log(client):
 @pytest.mark.django_db
 def test_dry_run(client):
     """Test dry run."""
+    assert f1.Driver.objects.get(reference="max_verstappen").forename == "Max"
     data = {
         "dry_run": True,
         "data": [
             {
                 "object_type": "Driver",
                 "foreign_keys": {},
-                "objects": [{"reference": "max_verstappen", "forename": "Max"}],
+                "objects": [{"reference": "max_verstappen", "forename": "Maxxx"}],
             }
         ],
     }
     response = client.put("/data/import/", data, format="json")
 
     assert response.status_code == status.HTTP_200_OK
-    assert DataImportLog.objects.count() == 0  # No log for dry run
+    assert f1.Driver.objects.get(reference="max_verstappen").forename == "Max"
+    assert DataImportLog.objects.count() == 1
+    assert DataImportLog.objects.first().dry_run
 
 
+@pytest.mark.parametrize(["dry_run"], [(True,), (False,)])
 @pytest.mark.django_db
-def test_db_error(client):
+def test_db_error(client, dry_run):
     """Test database error during import."""
     data = {
-        "dry_run": False,
+        "dry_run": dry_run,
         "legacy_import": False,
         "data": [
             {
@@ -317,5 +327,6 @@ def test_db_error(client):
     assert response.status_code == status.HTTP_400_BAD_REQUEST
     assert DataImportLog.objects.count() == 1
     log = DataImportLog.objects.first()
+    assert log.dry_run == dry_run
     assert not log.is_success
     assert log.error_type == "IMPORT"

jolpica_api/data_import/views.py

@@ -1,6 +1,5 @@
-from collections import defaultdict
-
 from django.contrib.auth.models import User
+from django.db import transaction
 from pydantic import BaseModel, Field, ValidationError
 from rest_framework.permissions import IsAdminUser, IsAuthenticated
 from rest_framework.request import Request
@@ -14,6 +13,10 @@
 from .models import DataImportLog
 
 
+class DryRunError(Exception):
+    pass
+
+
 class ImportDataRequestData(BaseModel):
     # If True, the data will be validated but not saved to the database
     dry_run: bool = True
@@ -38,51 +41,74 @@ def put(self, request: Request) -> Response:
             request_data = ImportDataRequestData.model_validate(request.data)
         except ValidationError as ex:
             errors = ex.errors(include_url=False)
-            DataImportLog(user=request.user, is_success=False, error_type="VALIDATION", errors=errors).save()
+            dry_run = request.data.get("dry_run", True)
+            log_data_import_result(
+                request.user,
+                dry_run=dry_run if isinstance(dry_run, bool) else True,
+                error_type="VALIDATION",
+                errors=errors,
+            )
             return Response({"errors": errors}, status=400)
 
         model_importer = JSONModelImporter(legacy_import=request_data.legacy_import)
         result = model_importer.deserialise_all(request.data["data"])
 
         if not result.success:
-            DataImportLog(
+            log_data_import_result(
+                request.user,
+                dry_run=request_data.dry_run,
                 description=request_data.description,
-                user=request.user,
-                is_success=False,
                 error_type="DESERIALISATION",
                 errors=result.errors,
-            ).save()
+            )
             return Response({"errors": result.errors}, status=400)
 
-        if not request_data.dry_run:
-            try:
-                model_importer.save_deserialisation_result_to_db(result)
-            except Exception as ex:
-                DataImportLog(
-                    description=request_data.description,
-                    user=request.user,
-                    is_success=False,
-                    error_type="IMPORT",
-                    errors=[repr(ex)],
-                ).save()
-                return Response({"errors": [{"type": "import_error", "message": repr(ex)}]}, status=400)
-
-            save_successful_import_to_db(request_data.description, request.user, result)
-        return Response({})
-
-
-def save_successful_import_to_db(description: str, user: User | None, result: DeserialisationResult) -> None:
-    updated_record_count = 0
-    updated_records = defaultdict(list)
-    for model_import, instances in result.instances.items():
-        instance_pks = [ins.pk for ins in instances]
-        updated_record_count += len(instance_pks)
-        updated_records[model_import.model_class.__name__].extend(instance_pks)
+        try:
+            import_stats = save_deserialisation_result_to_db(result, request_data.dry_run)
+        except Exception as ex:
+            errors = [{"type": "import_error", "message": repr(ex)}]
+            log_data_import_result(
+                request.user,
+                dry_run=request_data.dry_run,
+                description=request_data.description,
+                error_type="IMPORT",
+                errors=errors,
+            )
+            return Response({"errors": errors}, status=400)
+
+        log_data_import_result(
+            request.user, dry_run=request_data.dry_run, description=request_data.description, import_stats=import_stats
+        )
+        return Response(import_stats)
 
+
+def save_deserialisation_result_to_db(result: DeserialisationResult, dry_run: bool) -> dict:
+    try:
+        with transaction.atomic():
+            import_stats = JSONModelImporter.save_deserialisation_result_to_db(result)
+
+            if dry_run:
+                raise DryRunError("Transaction should be rolled back as this is a dry run")  # noqa: TRY301
+    except DryRunError:
+        pass  # Rollback the transaction, but keep import_stats
+    return import_stats
+
+
+def log_data_import_result(
+    user: User,
+    dry_run: bool,
+    description: str = "",
+    import_stats: dict | None = None,
+    error_type: str | None = None,
+    errors: list | None = None,
+):
     DataImportLog(
-        description=description,
+        is_success=False if error_type else True,
+        dry_run=dry_run,
         user=user,
-        is_success=True,
-        total_records=updated_record_count,
-        updated_records=updated_records,
+        description=description,
+        total_records=import_stats.get("total_count") if import_stats else None,
+        import_result=import_stats,
+        error_type=error_type,
+        errors=errors,
     ).save()