fix(cohorts): cap Mixpanel pages and reject syncs to a deleting cohort

Author: gagantrivedi

api/cohorts/serializers.py

@@ -108,7 +108,12 @@ class MixpanelParametersSerializer(serializers.Serializer[None]):
     mixpanel_cohort_id = serializers.CharField(max_length=255)
     mixpanel_cohort_name = serializers.CharField(max_length=2000)
     # An empty page is valid: a first sync of an empty cohort has no members.
-    members = MixpanelMemberSerializer(many=True, allow_empty=True)
+    # Mixpanel sends at most 1000 members per message; the cap stops anything
+    # else from posting an arbitrarily large page.
+    # The stubs don't know many=True forwards max_length to the list serialiser.
+    members = MixpanelMemberSerializer(  # type: ignore[call-arg]
+        many=True, allow_empty=True, max_length=1000
+    )
 
 
 class MixpanelWebhookSerializer(serializers.Serializer[None]):

api/cohorts/services.py

@@ -207,6 +207,22 @@ def get_cohort_for_source(
     return cohort
 
 
+def cohort_deletion_in_progress(
+    *,
+    environment: "Environment",
+    source_type: CohortSourceType,
+    external_id: str,
+) -> bool:
+    return bool(
+        Cohort.objects.filter(
+            environment=environment,
+            source_type=source_type,
+            external_id=external_id,
+            deletion_requested_at__isnull=False,
+        ).exists()
+    )
+
+
 def add_cohort_members(cohort: Cohort, identifiers: "typing.Iterable[str]") -> None:
     from cohorts.tasks import apply_cohort_membership_deltas
 

api/cohorts/sync_views.py

@@ -6,7 +6,7 @@
 from drf_spectacular.utils import extend_schema, extend_schema_view, inline_serializer
 from rest_framework import serializers, viewsets
 from rest_framework.decorators import action
-from rest_framework.exceptions import NotFound, ParseError
+from rest_framework.exceptions import NotFound, ParseError, ValidationError
 from rest_framework.request import Request
 from rest_framework.response import Response
 from rest_framework.views import APIView
@@ -157,6 +157,17 @@ def post(self, request: Request) -> Response:
         environment = typing.cast(CohortSyncKey, request.auth).environment
 
         if webhook_action == "members":
+            if services.cohort_deletion_in_progress(
+                environment=environment,
+                source_type=CohortSourceType.MIXPANEL,
+                external_id=parameters["mixpanel_cohort_id"],
+            ):
+                # Recreating the cohort while its memberships are still being
+                # drained would resurrect it. The 404 pauses the sync and
+                # emails the customer.
+                return self._failure(
+                    request, message="Cohort is being deleted.", code=404
+                )
             # A large first sync arrives as several requests, each one page
             # of members. Every page only adds; removals can't be detected
             # without seeing all pages at once.
@@ -194,6 +205,14 @@ def handle_exception(self, exc: Exception) -> Response:
             # A body that isn't valid JSON raises before post() runs, so the
             # response is shaped here to keep the envelope Mixpanel expects.
             return self._failure(self.request, message="Invalid payload.", code=400)
+        if isinstance(exc, ValidationError):
+            # Raised below the view, e.g. by the segment limit on cohort
+            # creation; shaped here for the same reason.
+            return self._failure(
+                self.request,
+                message=json.dumps(exc.detail, default=str),
+                code=400,
+            )
         return super().handle_exception(exc)
 
     def _failure(self, request: Request, *, message: str, code: int) -> Response:

api/tests/unit/cohorts/test_sync_views.py

@@ -913,3 +913,104 @@ def test_mixpanel_webhook__distinct_id_over_1024_bytes__returns_400_failure(
     assert body["status"] == "failure"
     assert "1024 bytes" in body["error"]["message"]
     assert not CohortMembership.objects.exists()
+
+
+def test_mixpanel_webhook__over_1000_members__returns_400_failure(
+    postgres_cohort_sync_key: _KeyAndPlaintext,
+    mixpanel_cohort: Cohort,
+) -> None:
+    # Given - one more member than Mixpanel's documented batch size
+    _, plaintext = postgres_cohort_sync_key
+    client = _basic_auth_client(plaintext)
+    url = reverse("api-v1:cohort-sync:mixpanel")
+
+    # When
+    response = client.post(
+        url,
+        data={
+            "action": "add_members",
+            "parameters": {
+                "mixpanel_cohort_id": "mp-42",
+                "mixpanel_cohort_name": "Power users",
+                "members": [{"mixpanel_distinct_id": f"user-{i}"} for i in range(1001)],
+            },
+        },
+        format="json",
+    )
+
+    # Then
+    assert response.status_code == status.HTTP_400_BAD_REQUEST
+    assert response.json()["status"] == "failure"
+    assert not CohortMembership.objects.exists()
+
+
+def test_mixpanel_webhook__members_action_at_segment_limit__returns_400_failure(
+    postgres_cohort_sync_key: _KeyAndPlaintext,
+    settings: SettingsWrapper,
+) -> None:
+    # Given - the project already holds as many segments as the plan allows
+    key, plaintext = postgres_cohort_sync_key
+    project = key.environment.project
+    settings.EDGE_ENABLED = True
+    project.max_segments_allowed = 1
+    project.save()
+    Segment.objects.create(name="existing", project=project)
+    client = _basic_auth_client(plaintext)
+    url = reverse("api-v1:cohort-sync:mixpanel")
+
+    # When
+    response = client.post(
+        url,
+        data={
+            "action": "members",
+            "parameters": {
+                "mixpanel_cohort_id": "mp-99",
+                "mixpanel_cohort_name": "New cohort",
+                "members": [{"mixpanel_distinct_id": "user-1"}],
+            },
+        },
+        format="json",
+    )
+
+    # Then
+    assert response.status_code == status.HTTP_400_BAD_REQUEST
+    body = response.json()
+    assert body["action"] == "members"
+    assert body["status"] == "failure"
+    assert "maximum allowed segments" in body["error"]["message"]
+    assert not Cohort.objects.filter(external_id="mp-99").exists()
+
+
+def test_mixpanel_webhook__members_action_while_cohort_draining__returns_404_failure(
+    postgres_cohort_sync_key: _KeyAndPlaintext,
+    mixpanel_cohort: Cohort,
+) -> None:
+    # Given - the cohort was deleted in Flagsmith and is still draining
+    mixpanel_cohort.deletion_requested_at = timezone.now()
+    mixpanel_cohort.save()
+    _, plaintext = postgres_cohort_sync_key
+    client = _basic_auth_client(plaintext)
+    url = reverse("api-v1:cohort-sync:mixpanel")
+
+    # When
+    response = client.post(
+        url,
+        data={
+            "action": "members",
+            "parameters": {
+                "mixpanel_cohort_id": "mp-42",
+                "mixpanel_cohort_name": "Power users",
+                "members": [{"mixpanel_distinct_id": "user-1"}],
+            },
+        },
+        format="json",
+    )
+
+    # Then
+    assert response.status_code == status.HTTP_404_NOT_FOUND
+    assert response.json() == {
+        "action": "members",
+        "status": "failure",
+        "error": {"message": "Cohort is being deleted.", "code": 404},
+    }
+    assert Cohort.objects.filter(source_type=CohortSourceType.MIXPANEL).count() == 1

docs/docs/deployment-self-hosting/observability/_events-catalogue.md

@@ -86,7 +86,7 @@ Attributes:
 ### `cohorts.cohort.deleted`
 
 Logged at `info` from:
- - `api/cohorts/services.py:421`
+ - `api/cohorts/services.py:437`
 
 Attributes:
  - `cohort.id`
@@ -95,7 +95,7 @@ Attributes:
 ### `cohorts.cohort.deletion_requested`
 
 Logged at `info` from:
- - `api/cohorts/services.py:405`
+ - `api/cohorts/services.py:421`
 
 Attributes:
  - `cohort.id`
@@ -104,7 +104,7 @@ Attributes:
 ### `cohorts.csv.synced`
 
 Logged at `info` from:
- - `api/cohorts/services.py:377`
+ - `api/cohorts/services.py:393`
 
 Attributes:
  - `adds.count`
@@ -117,7 +117,7 @@ Attributes:
 ### `cohorts.membership.adds_received`
 
 Logged at `info` from:
- - `api/cohorts/services.py:230`
+ - `api/cohorts/services.py:246`
 
 Attributes:
  - `cohort.id`
@@ -155,7 +155,7 @@ Attributes:
 ### `cohorts.membership.removals_received`
 
 Logged at `info` from:
- - `api/cohorts/services.py:255`
+ - `api/cohorts/services.py:271`
 
 Attributes:
  - `cohort.id`
@@ -166,7 +166,7 @@ Attributes:
 ### `cohorts.sync_webhook.rejected`
 
 Logged at `warning` from:
- - `api/cohorts/sync_views.py:200`
+ - `api/cohorts/sync_views.py:219`
 
 Attributes:
  - `action`