feat(river): persist progressions with topicleaner

Author: Jasopaum

django/river/api/serializers/serializers.py

@@ -3,6 +3,7 @@
 
 from rest_framework import serializers
 
+from pyrog import models as pyrog_models
 from river import models
 from river.adapters.progression_counter import RedisProgressionCounter
 
@@ -13,9 +14,24 @@ class Meta:
         fields = "__all__"
 
 
+class ResourceForProgressionSerializer(serializers.ModelSerializer):
+    class Meta:
+        model = pyrog_models.Resource
+        fields = ["id", "definition_id", "label"]
+
+
+class ProgressionSerializer(serializers.ModelSerializer):
+    resource = ResourceForProgressionSerializer(read_only=True)
+
+    class Meta:
+        model = models.Progression
+        fields = ["resource", "extracted", "loaded", "failed"]
+
+
 class BatchSerializer(serializers.ModelSerializer):
+    # FIXME errors is a property of BaseSerializer, we shouldn't override it
     errors = ErrorSerializer(many=True, read_only=True)
-    progressions = serializers.SerializerMethodField()
+    progressions = ProgressionSerializer(many=True, read_only=True)
 
     class Meta:
         model = models.Batch

django/river/api/views.py

@@ -45,6 +45,12 @@ def create(self, request, *args, **kwargs):
 
         batch_instance = serializer.save()
 
+        # Create Progressions
+        for resource in resources:
+            river_models.Progression.objects.create(
+                batch=batch_instance, resource=resource, extracted=None, loaded=None, failed=None
+            )
+
         topics_manager = KafkaTopicsManager()
         event_publisher = KafkaEventPublisher()
 

django/river/services.py

@@ -1,4 +1,5 @@
 import json
+import logging
 from typing import Any, List, Optional, Tuple
 
 from django.utils import timezone
@@ -16,6 +17,8 @@
 from river.transformer.transformer import Transformer
 from utils.json import CustomJSONEncoder
 
+logger = logging.getLogger(__name__)
+
 
 def batch(
     batch_id: str,
@@ -37,18 +40,20 @@ def abort(batch: models.Batch, topics_manager: TopicsManager, counter: Progressi
     for base_topic in ["batch", "extract", "transform", "load"]:
         topics_manager.delete(f"{base_topic}.{batch.id}")
 
-    # Persist progressions in DB
+    # Update Progressions in DB
     for resource in batch.resources.all():
-        resource_progression = counter.get(f"{batch.id}:{resource.id}")
-        if not resource_progression:
+        redis_progression = counter.get(f"{batch.id}:{resource.id}")
+        if not redis_progression:
             continue
-        models.Progression.objects.create(
-            batch=batch,
-            resource=resource,
-            extracted=resource_progression.extracted,
-            loaded=resource_progression.loaded,
-            failed=resource_progression.failed,
-        )
+        try:
+            resource_progression = models.Progression.objects.get(batch=batch, resource=resource)
+        except models.Progression.DoesNotExist:
+            logger.warning(f"Could not find progression of resource {resource} in batch {batch}")
+            continue
+        resource_progression.extracted = redis_progression.extracted
+        resource_progression.loaded = redis_progression.loaded
+        resource_progression.failed = redis_progression.failed
+        resource_progression.save()
 
     batch.canceled_at = timezone.now()
     batch.save()

django/river/topicleaner/service.py

@@ -16,7 +16,7 @@ def teardown_after_batch(batch: Batch, topics: TopicsManager):
         topics.delete(f"{base_topic}.{batch.id}")
 
 
-def clean(counter: ProgressionCounter, topics: TopicsManager):
+def task(counter: ProgressionCounter, topics: TopicsManager):
     current_batches = Batch.objects.filter(completed_at__isnull=True, canceled_at__isnull=True).prefetch_related(
         "resources"
     )
@@ -26,6 +26,21 @@ def clean(counter: ProgressionCounter, topics: TopicsManager):
             resource: counter.get(f"{batch.id}:{resource.id}") for resource in batch.resources.all()
         }
 
+        # Update Progressions in DB
+        for resource, redis_progression in resources_progressions.items():
+            if not redis_progression:
+                continue
+            try:
+                resource_progression = models.Progression.objects.get(batch=batch, resource=resource)
+            except models.Progression.DoesNotExist:
+                logger.warning(f"Could not find progression of resource {resource} in batch {batch}")
+                continue
+            resource_progression.extracted = redis_progression.extracted
+            resource_progression.loaded = redis_progression.loaded
+            resource_progression.failed = redis_progression.failed
+            resource_progression.save()
+
+        # Clear if needed
         if all(
             [
                 progression is not None
@@ -36,17 +51,6 @@ def clean(counter: ProgressionCounter, topics: TopicsManager):
         ):
             logger.info(f"Deleting batch {batch}.")
 
-            for resource, progression in resources_progressions.items():
-                if not progression:
-                    continue
-                models.Progression.objects.create(
-                    batch=batch,
-                    resource=resource,
-                    extracted=progression.extracted,
-                    loaded=progression.loaded,
-                    failed=progression.failed,
-                )
-
             teardown_after_batch(batch, topics)
             batch.completed_at = timezone.now()
             batch.save()
@@ -56,5 +60,5 @@ def clean(counter: ProgressionCounter, topics: TopicsManager):
 
 def run(counter: ProgressionCounter, topics: TopicsManager):
     while True:
-        clean(counter=counter, topics=topics)
+        task(counter=counter, topics=topics)
         sleep(10)

tests/river/conftest.py

@@ -21,6 +21,7 @@
 
 register(factories.BatchFactory)
 register(factories.ErrorFactory)
+register(factories.ProgressionFactory)
 register(ResourceFactory)
 register(SourceFactory)
 

tests/river/e2e/api/test_batch.py

@@ -58,11 +58,13 @@ def test_list_batch(api_client, batch_factory):
 
 
 @pytest.mark.redis
-def test_get_batch_progression(api_client, batch_factory, resource_factory, redis_client):
+def test_get_batch_progression(api_client, resource_factory, batch_factory, progression_factory, redis_client):
     url = reverse("batches-list")
     r1 = resource_factory.create(definition_id="Patient")
     r2 = resource_factory.create(definition_id="Practitioner")
     batch = batch_factory.create(resources=[r1, r2])
+    progression_factory.create(batch=batch, resource=r1)
+    progression_factory.create(batch=batch, resource=r2)
 
     set_counters(redis_client, batch, r1, r2)
 
@@ -72,8 +74,18 @@ def test_get_batch_progression(api_client, batch_factory, resource_factory, redi
     assert len(response.data) == 1
     batch_response = response.json()[0]
     assert batch_response["progressions"] == [
-        ["Patient", {"extracted": 10, "loaded": 5, "failed": None}],
-        ["Practitioner", {"extracted": 20, "loaded": 5, "failed": 3}],
+        {
+            "resource": {"id": r1.id, "definition_id": "Patient", "label": ""},
+            "extracted": 100,
+            "loaded": 50,
+            "failed": None,
+        },
+        {
+            "resource": {"id": r2.id, "definition_id": "Practitioner", "label": ""},
+            "extracted": 100,
+            "loaded": 50,
+            "failed": None,
+        },
     ]
 
 
@@ -106,10 +118,12 @@ def test_retrieve_batch(api_client, batch_factory, resource_factory):
 
 
 @pytest.mark.redis
-def test_delete_batch(api_client, redis_client, batch_factory, resource_factory, kafka_admin):
+def test_delete_batch(api_client, redis_client, resource_factory, batch_factory, progression_factory, kafka_admin):
     r1 = resource_factory.create(definition_id="Patient")
     r2 = resource_factory.create(definition_id="Practitioner")
     batch = batch_factory.create(resources=[r1, r2])
+    progression_factory.create(batch=batch, resource=r1)
+    progression_factory.create(batch=batch, resource=r2)
     url = reverse("batches-detail", kwargs={"pk": batch.id})
 
     set_counters(redis_client, batch, r1, r2)
@@ -120,11 +134,21 @@ def test_delete_batch(api_client, redis_client, batch_factory, resource_factory,
     clear_counters(redis_client, batch, r1, r2)
 
     response_get = api_client.get(url)
-    print(response_get.json())
     assert response_get.json()["canceled_at"] is not None
+    print(response_get.json()["progressions"])
     assert response_get.json()["progressions"] == [
-        ["Patient", {"extracted": 10, "loaded": 5, "failed": None}],
-        ["Practitioner", {"extracted": 20, "loaded": 5, "failed": 3}],
+        {
+            "resource": {"id": r1.id, "definition_id": "Patient", "label": ""},
+            "extracted": 10,
+            "loaded": 5,
+            "failed": None,
+        },
+        {
+            "resource": {"id": r2.id, "definition_id": "Practitioner", "label": ""},
+            "extracted": 20,
+            "loaded": 5,
+            "failed": 3,
+        },
     ]
 
     # Check that topics are deleted

tests/river/factories.py

@@ -37,3 +37,13 @@ class Meta:
 
     id = factory.Sequence(lambda n: f"batch_id_{n:04d}")
     batch = factory.SubFactory(BatchFactory)
+
+
+class ProgressionFactory(factory.django.DjangoModelFactory):
+    class Meta:
+        model = "river.Progression"
+
+    id = factory.Sequence(lambda n: f"progression_id_{n:04d}")
+    extracted = 100
+    loaded = 50
+    failed = None

tests/river/unit/test_services.py

@@ -25,10 +25,12 @@ def test_batch(batch_factory, resource_factory):
     ]
 
 
-def test_abort(batch_factory, resource_factory):
+def test_abort(resource_factory, batch_factory, progression_factory):
     r1 = resource_factory.create(definition_id="Patient")
     r2 = resource_factory.create(definition_id="Practitioner")
     batch = batch_factory.create(resources=[r1, r2])
+    progression_factory.create(batch=batch, resource=r1)
+    progression_factory.create(batch=batch, resource=r2)
 
     topics = InMemoryTopicsManager(
         topics=[f"{base_topic}.{batch.id}" for base_topic in ["batch", "extract", "transform", "load"]]
@@ -43,16 +45,14 @@ def test_abort(batch_factory, resource_factory):
 
     assert topics._topics == set()
     assert batch.canceled_at is not None
-    r1_progressions = models.Progression.objects.filter(batch=batch, resource=r1)
-    assert len(r1_progressions) == 1
-    assert r1_progressions[0].extracted == 100
-    assert r1_progressions[0].loaded == 20
-    assert r1_progressions[0].failed == 3
-    r2_progressions = models.Progression.objects.filter(batch=batch, resource=r2)
-    assert len(r2_progressions) == 1
-    assert r2_progressions[0].extracted == 200
-    assert r2_progressions[0].loaded == 10
-    assert r2_progressions[0].failed is None
+    r1_progressions = models.Progression.objects.get(batch=batch, resource=r1)
+    assert r1_progressions.extracted == 100
+    assert r1_progressions.loaded == 20
+    assert r1_progressions.failed == 3
+    r2_progressions = models.Progression.objects.get(batch=batch, resource=r2)
+    assert r2_progressions.extracted == 200
+    assert r2_progressions.loaded == 10
+    assert r2_progressions.failed is None
 
 
 @pytest.mark.skip(reason="feature not implemented yet")

tests/river/unit/topicleaner/test_service.py

@@ -3,22 +3,24 @@
 from river import models
 from river.adapters.progression_counter import InMemoryProgressionCounter
 from river.adapters.topics import InMemoryTopicsManager
-from river.topicleaner.service import clean
+from river.topicleaner.service import task as topicleaner_task
 
 pytestmark = pytest.mark.django_db
 
 
-def test_done_batch_is_cleaned(batch_factory, resource_factory):
+def test_done_batch_is_cleaned(resource_factory, batch_factory, progression_factory):
     r1, r2 = resource_factory.create_batch(2)
     batch = batch_factory.create(resources=[r1, r2])
+    progression_factory.create(batch=batch, resource=r1)
+    progression_factory.create(batch=batch, resource=r2)
     counters = InMemoryProgressionCounter(
         counts={f"{batch.id}:{resource.id}": {"extracted": 10, "loaded": 10} for resource in batch.resources.all()}
     )
     topics = InMemoryTopicsManager(
         topics=[f"{base_topic}.{batch.id}" for base_topic in ["batch", "extract", "transform", "load"]]
     )
 
-    clean(counters, topics)
+    topicleaner_task(counters, topics)
 
     assert topics._topics == set()
     batches = models.Batch.objects.all()
@@ -30,9 +32,11 @@ def test_done_batch_is_cleaned(batch_factory, resource_factory):
         assert progression.failed is None
 
 
-def test_done_batch_is_cleaned_with_failed(batch_factory, resource_factory):
+def test_done_batch_is_cleaned_with_failed(resource_factory, batch_factory, progression_factory):
     r1, r2 = resource_factory.create_batch(2)
     batch = batch_factory.create(resources=[r1, r2])
+    progression_factory.create(batch=batch, resource=r1)
+    progression_factory.create(batch=batch, resource=r2)
     counters = InMemoryProgressionCounter(
         counts={
             f"{batch.id}:{resource.id}": {"extracted": 10, "loaded": 6, "failed": 4}
@@ -42,7 +46,7 @@ def test_done_batch_is_cleaned_with_failed(batch_factory, resource_factory):
     topics = InMemoryTopicsManager(
         topics=[f"{base_topic}.{batch.id}" for base_topic in ["batch", "extract", "transform", "load"]]
     )
-    clean(counters, topics)
+    topicleaner_task(counters, topics)
 
     assert topics._topics == set()
     batches = models.Batch.objects.all()
@@ -64,7 +68,7 @@ def test_ongoing_batch_is_not_cleaned(batch_factory, resource_factory):
         topics=[f"{base_topic}.{batch.id}" for base_topic in ["batch", "extract", "transform", "load"]]
     )
 
-    clean(counters, topics)
+    topicleaner_task(counters, topics)
 
     assert topics._topics != set()
 
@@ -82,7 +86,7 @@ def test_ongoing_batch_is_not_cleaned_with_failed(batch_factory, resource_factor
         topics=[f"{base_topic}.{batch.id}" for base_topic in ["batch", "extract", "transform", "load"]]
     )
 
-    clean(counters, topics)
+    topicleaner_task(counters, topics)
 
     assert topics._topics != set()
 
@@ -97,7 +101,7 @@ def test_none_counter_prevents_cleaning(batch_factory, resource_factory):
         topics=[f"{base_topic}.{batch.id}" for base_topic in ["batch", "extract", "transform", "load"]]
     )
 
-    clean(counters, topics)
+    topicleaner_task(counters, topics)
 
     assert topics._topics != set()
 
@@ -112,6 +116,6 @@ def test_missing_counter_prevents_cleaning(batch_factory, resource_factory):
         topics=[f"{base_topic}.{batch.id}" for base_topic in ["batch", "extract", "transform", "load"]]
     )
 
-    clean(counters, topics)
+    topicleaner_task(counters, topics)
 
     assert topics._topics != set()