feat(cohorts): expose sync state and allow updating the managed segment (#8386)

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: flagsmith-engineering[bot] <flagsmith-engineering[bot]@users.noreply.github.com>

Author: 3 people

api/cohorts/migrations/0005_cohort_last_synced_at.py

@@ -0,0 +1,15 @@
+from django.db import migrations, models
+
+
+class Migration(migrations.Migration):
+    dependencies = [
+        ("cohorts", "0004_mixpanel_source"),
+    ]
+
+    operations = [
+        migrations.AddField(
+            model_name="cohort",
+            name="last_synced_at",
+            field=models.DateTimeField(blank=True, null=True),
+        ),
+    ]

api/cohorts/models.py

@@ -36,6 +36,7 @@ class Cohort(SoftDeleteExportableModel):
     external_id = models.CharField(max_length=255, null=True, blank=True)
     version = models.PositiveIntegerField(default=0)
     created_at = models.DateTimeField(auto_now_add=True)
+    last_synced_at = models.DateTimeField(null=True, blank=True)
     # Deletion drains memberships from the identity store first; the cohort is
     # only soft-deleted once drained. This marks it as awaiting that final step.
     deletion_requested_at = models.DateTimeField(null=True, blank=True)

api/cohorts/serializers.py

@@ -1,11 +1,13 @@
 import typing
 
 from django.core.files.uploadedfile import UploadedFile
+from django.db.models import Count, Q
+from drf_spectacular.utils import extend_schema_field
 from rest_framework import serializers
 
 from cohorts.constants import COHORT_CSV_MAX_FILE_SIZE_BYTES
 from cohorts.exceptions import CsvFileTooLargeError
-from cohorts.models import Cohort, CohortSyncKey
+from cohorts.models import Cohort, CohortMembershipState, CohortSyncKey
 from cohorts.services import create_cohort
 from environments.models import Environment
 from metadata.serializers import MetadataSerializer, MetadataSerializerMixin
@@ -19,12 +21,19 @@ class Meta:
         model = Segment
 
 
+class CohortMembershipCountsSerializer(serializers.Serializer):  # type: ignore[type-arg]
+    applied = serializers.IntegerField(min_value=0)
+    pending_add = serializers.IntegerField(min_value=0)
+    pending_remove = serializers.IntegerField(min_value=0)
+
+
 class CohortSerializer(serializers.ModelSerializer[Cohort]):
     name = serializers.CharField(max_length=2000, source="segment.name")
     description = serializers.CharField(
         source="segment.description", required=False, allow_null=True
     )
     metadata = MetadataSerializer(required=False, many=True, write_only=True)
+    membership_counts = serializers.SerializerMethodField()
 
     class Meta:
         model = Cohort
@@ -38,11 +47,41 @@ class Meta:
             "source_type",
             "version",
             "created_at",
+            "last_synced_at",
+            "membership_counts",
+        )
+        read_only_fields = (
+            "segment",
+            "source_type",
+            "version",
+            "created_at",
+            "last_synced_at",
+        )
+
+    @extend_schema_field(CohortMembershipCountsSerializer)
+    def get_membership_counts(self, cohort: Cohort) -> dict[str, int]:
+        # Clients derive sync status and progress from these. The viewset
+        # annotates the counts; a freshly created cohort isn't annotated.
+        if (applied := getattr(cohort, "applied_count", None)) is not None:
+            return {
+                "applied": applied,
+                "pending_add": getattr(cohort, "pending_add_count", 0),
+                "pending_remove": getattr(cohort, "pending_remove_count", 0),
+            }
+        return cohort.memberships.aggregate(
+            applied=Count("id", filter=Q(state=CohortMembershipState.APPLIED)),
+            pending_add=Count("id", filter=Q(state=CohortMembershipState.PENDING_ADD)),
+            pending_remove=Count(
+                "id", filter=Q(state=CohortMembershipState.PENDING_REMOVE)
+            ),
         )
-        read_only_fields = ("segment", "source_type", "version", "created_at")
 
     def validate(self, attrs: dict[str, typing.Any]) -> dict[str, typing.Any]:
         attrs = super().validate(attrs)
+        if self.instance is not None and "metadata" not in attrs:
+            # A partial update without metadata must not fail the
+            # required-metadata check.
+            return attrs
         environment = Environment.objects.get(
             api_key=self.context["view"].kwargs["environment_api_key"]
         )
@@ -64,6 +103,19 @@ def create(self, validated_data: dict[str, typing.Any]) -> Cohort:
             _SegmentMetadataHandler()._update_metadata(cohort.segment, metadata_data)
         return cohort
 
+    def update(self, instance: Cohort, validated_data: dict[str, typing.Any]) -> Cohort:
+        # Only the managed segment's fields are updatable.
+        metadata_data = validated_data.pop("metadata", None)
+        segment_data = validated_data.pop("segment", {})
+        if segment_data:
+            segment = instance.segment
+            for field, value in segment_data.items():
+                setattr(segment, field, value)
+            segment.save(update_fields=list(segment_data))
+        if metadata_data is not None:
+            _SegmentMetadataHandler()._update_metadata(instance.segment, metadata_data)
+        return instance
+
 
 class CohortSyncKeySerializer(serializers.ModelSerializer[CohortSyncKey]):
     key = serializers.SerializerMethodField()

api/cohorts/services.py

@@ -404,8 +404,10 @@ def sync_cohort_memberships_from_csv(
             state=CohortMembershipState.PENDING_REMOVE, updated_at=timezone.now()
         )
 
-    Cohort.objects.filter(id=cohort.id).update(version=F("version") + 1)
-    cohort.refresh_from_db(fields=["version"])
+    Cohort.objects.filter(id=cohort.id).update(
+        version=F("version") + 1, last_synced_at=timezone.now()
+    )
+    cohort.refresh_from_db(fields=["version", "last_synced_at"])
     apply_cohort_membership_deltas.delay(kwargs={"cohort_id": cohort.id})
 
     flagsmith_cohorts_csv_syncs_total.inc()

api/cohorts/views.py

@@ -1,4 +1,4 @@
-from django.db.models import QuerySet
+from django.db.models import Count, Q, QuerySet
 from drf_spectacular.utils import extend_schema, extend_schema_view
 from rest_framework import mixins, status, viewsets
 from rest_framework.decorators import action
@@ -10,7 +10,7 @@
 
 from api.serializers import ErrorSerializer
 from cohorts import services
-from cohorts.models import Cohort, CohortSyncKey
+from cohorts.models import Cohort, CohortMembershipState, CohortSyncKey
 from cohorts.permissions import CohortPermission, CohortPlanPermission
 from cohorts.serializers import (
     CohortCsvSyncResultSerializer,
@@ -29,6 +29,11 @@
         description="Create a cohort and the managed segment that targets it."
     ),
     retrieve=extend_schema(description="Retrieve a cohort."),
+    partial_update=extend_schema(
+        description=(
+            "Update the cohort's managed segment name, description and metadata."
+        )
+    ),
     destroy=extend_schema(
         description=(
             "Request cohort deletion. Memberships are drained from identity "
@@ -42,6 +47,7 @@ class CohortViewSet(
     mixins.ListModelMixin,
     mixins.CreateModelMixin,
     mixins.RetrieveModelMixin,
+    mixins.UpdateModelMixin,
     mixins.DestroyModelMixin,
 ):
     serializer_class = CohortSerializer
@@ -50,6 +56,8 @@ class CohortViewSet(
     model_class = Cohort
     lookup_field = "id"
     lookup_url_kwarg = "cohort_id"
+    # PATCH only: a cohort has no meaningful full replacement.
+    http_method_names = ["get", "post", "patch", "delete", "head", "options"]
 
     def get_queryset(self) -> QuerySet[Cohort]:
         # A cohort awaiting drain-then-delete is already gone from the
@@ -59,6 +67,20 @@ def get_queryset(self) -> QuerySet[Cohort]:
             .get_queryset()
             .filter(deletion_requested_at__isnull=True)
             .select_related("segment")
+            .annotate(
+                applied_count=Count(
+                    "memberships",
+                    filter=Q(memberships__state=CohortMembershipState.APPLIED),
+                ),
+                pending_add_count=Count(
+                    "memberships",
+                    filter=Q(memberships__state=CohortMembershipState.PENDING_ADD),
+                ),
+                pending_remove_count=Count(
+                    "memberships",
+                    filter=Q(memberships__state=CohortMembershipState.PENDING_REMOVE),
+                ),
+            )
             .order_by("id")
         )