fix(cohorts): harden sync key auth and drop the external ID constraint

Author: gagantrivedi

api/cohorts/authentication.py

@@ -1,5 +1,4 @@
 import base64
-import typing
 from contextlib import suppress
 
 from django.contrib.auth.models import AnonymousUser
@@ -38,16 +37,8 @@ def authenticate(
         else:
             return None
 
-        if "\x00" in raw_key:
-            # Postgres refuses to run a query containing a NUL character, so
-            # the key lookup below would crash instead of returning 401.
-            raise exceptions.AuthenticationFailed("Valid cohort sync key not found.")
-
         with suppress(CohortSyncKey.DoesNotExist):
-            key = typing.cast(
-                CohortSyncKey,
-                CohortSyncKey.objects.get_from_key(raw_key),
-            )
+            key = CohortSyncKey.objects.get_from_key(raw_key)
             if not key.has_expired:
                 # No person is acting here, so no user is returned: the key
                 # alone carries authority, and audit trails record the source

api/cohorts/migrations/0004_mixpanel_source.py

@@ -30,14 +30,4 @@ class Migration(migrations.Migration):
                 max_length=50,
             ),
         ),
-        migrations.AddConstraint(
-            model_name="cohort",
-            constraint=models.UniqueConstraint(
-                condition=models.Q(
-                    ("deleted_at__isnull", True), ("external_id__isnull", False)
-                ),
-                fields=("environment", "source_type", "external_id"),
-                name="unique_active_cohort_per_source_external_id",
-            ),
-        ),
     ]

api/cohorts/models.py

@@ -1,5 +1,7 @@
+import typing
+
 from django.db import models
-from rest_framework_api_key.models import AbstractAPIKey
+from rest_framework_api_key.models import AbstractAPIKey, APIKeyManager
 
 from cohorts.constants import COHORT_SYSTEM_TRAIT_KEY_PREFIX
 from core.models import SoftDeleteExportableModel
@@ -27,9 +29,10 @@ class Cohort(SoftDeleteExportableModel):
         choices=CohortSourceType.choices,
         default=CohortSourceType.CSV,
     )
-    # The cohort's identifier in the external source (e.g. Mixpanel's cohort
-    # ID). Set for sources that push to us under their own identifier; null
-    # for sources that adopt ours (Amplitude) and for CSV cohorts.
+    # The cohort's identifier in the external source. Mixpanel pushes under
+    # its own cohort ID, so we store it to route later requests; Amplitude
+    # uses the ID we hand back at list creation, and CSV cohorts have no
+    # external system, so both leave this null.
     external_id = models.CharField(max_length=255, null=True, blank=True)
     version = models.PositiveIntegerField(default=0)
     created_at = models.DateTimeField(auto_now_add=True)
@@ -50,19 +53,22 @@ class Meta:
                 condition=models.Q(deleted_at__isnull=True),
                 name="unique_active_cohort_per_segment",
             ),
-            # One Mixpanel cohort must map to one active cohort per
-            # environment: without this, two simultaneous first-sync requests
-            # would each create their own cohort and split the members
-            # between them.
-            models.UniqueConstraint(
-                fields=["environment", "source_type", "external_id"],
-                condition=models.Q(deleted_at__isnull=True, external_id__isnull=False),
-                name="unique_active_cohort_per_source_external_id",
-            ),
         ]
 
 
+class CohortSyncKeyManager(APIKeyManager):
+    def get_from_key(self, key: str) -> "CohortSyncKey":
+        if "\x00" in key:
+            # A NUL can't travel in a raw header, but base64 credentials can
+            # decode to one, and the database driver refuses to build a query
+            # containing it. No real key holds one, so treat it as absent.
+            raise self.model.DoesNotExist("Key contains a NUL character.")
+        return typing.cast("CohortSyncKey", super().get_from_key(key))
+
+
 class CohortSyncKey(AbstractAPIKey):
+    objects: typing.ClassVar[CohortSyncKeyManager] = CohortSyncKeyManager()
+
     environment = models.ForeignKey(
         "environments.Environment",
         on_delete=models.CASCADE,

api/cohorts/services.py

@@ -1,7 +1,7 @@
 import typing
 
 import structlog
-from django.db import IntegrityError, transaction
+from django.db import transaction
 from django.db.models import QuerySet
 from django.utils import timezone
 from flag_engine.segments.constants import IS_SET
@@ -170,38 +170,6 @@ def get_cohort_for_source(
     return cohort
 
 
-def get_or_create_cohort_for_source(
-    *,
-    environment: "Environment",
-    name: str,
-    source_type: CohortSourceType,
-    external_id: str,
-) -> Cohort | None:
-    if cohort := get_cohort_for_source(
-        environment=environment, source_type=source_type, external_id=external_id
-    ):
-        return cohort
-    try:
-        return create_cohort_for_source(
-            environment=environment,
-            name=name,
-            source_type=source_type,
-            external_id=external_id,
-        )
-    except IntegrityError:
-        # Two situations end up here. A simultaneous first-sync request
-        # created the cohort between our lookup and our insert — use the one
-        # it created. Or the cohort was deleted in Flagsmith and is still
-        # draining memberships from identity data: the lookup doesn't see it,
-        # but it still occupies the external ID — nothing usable exists, so
-        # return None.
-        if cohort := get_cohort_for_source(
-            environment=environment, source_type=source_type, external_id=external_id
-        ):
-            return cohort
-        return None
-
-
 def add_cohort_members(cohort: Cohort, identifiers: "typing.Iterable[str]") -> None:
     from cohorts.tasks import apply_cohort_membership_deltas
 

api/cohorts/sync_views.py

@@ -141,17 +141,17 @@ def post(self, request: Request) -> Response:
             # 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.
-            cohort_or_none = services.get_or_create_cohort_for_source(
+            cohort = services.get_cohort_for_source(
+                environment=environment,
+                source_type=CohortSourceType.MIXPANEL,
+                external_id=parameters["mixpanel_cohort_id"],
+            ) or services.create_cohort_for_source(
                 environment=environment,
                 name=parameters["mixpanel_cohort_name"],
                 source_type=CohortSourceType.MIXPANEL,
                 external_id=parameters["mixpanel_cohort_id"],
             )
-            if cohort_or_none is None:
-                return self._failure(
-                    request, message="Cohort is being deleted.", code=404
-                )
-            services.add_cohort_members(cohort_or_none, identifiers)
+            services.add_cohort_members(cohort, identifiers)
         else:
             cohort_or_none = services.get_cohort_for_source(
                 environment=environment,

api/tests/unit/cohorts/test_services.py

@@ -1,5 +1,3 @@
-from django.db import IntegrityError
-from django.utils import timezone
 from flag_engine.segments.constants import IS_SET
 from pytest_mock import MockerFixture
 from pytest_structlog import StructuredLogCapture
@@ -8,18 +6,16 @@
     Cohort,
     CohortMembership,
     CohortMembershipState,
-    CohortSourceType,
 )
 from cohorts.services import (
     apply_pending_memberships,
     create_cohort,
     delete_cohort,
-    get_or_create_cohort_for_source,
 )
 from environments.dynamodb import DynamoIdentityWrapper
 from environments.identities.models import Identity
 from environments.models import Environment
-from segments.models import Segment, SegmentManagedBy, SegmentRule
+from segments.models import SegmentManagedBy, SegmentRule
 
 
 def test_apply_pending_memberships__no_pending_rows__returns_false(
@@ -289,64 +285,3 @@ def test_apply_pending_memberships__system_trait_already_unset__deletes_membersh
     identity.refresh_from_db()
     assert identity.system_traits == {"flagsmith_cohort_other": True}
     assert not CohortMembership.objects.filter(cohort=cohort).exists()
-
-
-def test_get_or_create_cohort_for_source__simultaneous_creation__returns_other_requests_cohort(
-    environment: Environment,
-    mocker: MockerFixture,
-) -> None:
-    # Given - creating the cohort fails because another request created its
-    # own cohort between our lookup and our insert
-    def create_winning_cohort_and_conflict(**kwargs: object) -> Cohort:
-        segment = Segment.objects.create(
-            name="Power users", project=environment.project
-        )
-        Cohort.objects.create(
-            environment=environment,
-            segment=segment,
-            source_type=CohortSourceType.MIXPANEL,
-            external_id="mp-42",
-        )
-        raise IntegrityError("unique_active_cohort_per_source_external_id")
-
-    mocker.patch(
-        "cohorts.services.create_cohort_for_source",
-        side_effect=create_winning_cohort_and_conflict,
-    )
-
-    # When
-    cohort = get_or_create_cohort_for_source(
-        environment=environment,
-        name="Power users",
-        source_type=CohortSourceType.MIXPANEL,
-        external_id="mp-42",
-    )
-
-    # Then
-    assert cohort == Cohort.objects.get(external_id="mp-42")
-
-
-def test_get_or_create_cohort_for_source__deletion_requested_same_external_id__returns_none(
-    environment: Environment,
-) -> None:
-    # Given - a cohort with the same external ID is awaiting deletion, so it
-    # is invisible to the lookup but still occupies the external ID
-    segment = Segment.objects.create(name="Power users", project=environment.project)
-    Cohort.objects.create(
-        environment=environment,
-        segment=segment,
-        source_type=CohortSourceType.MIXPANEL,
-        external_id="mp-42",
-        deletion_requested_at=timezone.now(),
-    )
-
-    # When
-    cohort = get_or_create_cohort_for_source(
-        environment=environment,
-        name="Power users",
-        source_type=CohortSourceType.MIXPANEL,
-        external_id="mp-42",
-    )
-
-    # Then
-    assert cohort is None

api/tests/unit/cohorts/test_sync_views.py

@@ -696,41 +696,6 @@ def test_mixpanel_webhook__unparseable_body__returns_400_failure(
     }
 
 
-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
-
-
 def test_mixpanel_webhook__other_environment_key__returns_404_failure(
     mixpanel_cohort: Cohort,
 ) -> None:

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:276`
+ - `api/cohorts/services.py:244`
 
 Attributes:
  - `cohort.id`
@@ -95,7 +95,7 @@ Attributes:
 ### `cohorts.cohort.deletion_requested`
 
 Logged at `info` from:
- - `api/cohorts/services.py:260`
+ - `api/cohorts/services.py:228`
 
 Attributes:
  - `cohort.id`
@@ -132,8 +132,8 @@ Attributes:
 ### `cohorts.membership.deltas_received`
 
 Logged at `info` from:
- - `api/cohorts/services.py:225`
- - `api/cohorts/services.py:244`
+ - `api/cohorts/services.py:193`
+ - `api/cohorts/services.py:212`
 
 Attributes:
  - `action`