Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
39 changes: 29 additions & 10 deletions api/api/openapi.py
Original file line number Diff line number Diff line change
Expand Up @@ -165,19 +165,38 @@ def get_security_definition(

class CohortSyncKeyAuthenticationExtension(OpenApiAuthenticationExtension): # type: ignore[no-untyped-call]
target_class = "cohorts.authentication.CohortSyncKeyAuthentication"
name = "Cohort Sync Key"
name = ["Cohort Sync Key", "Cohort Sync Key (Basic)"]

def get_security_requirement(
self, auto_schema: openapi.AutoSchema
) -> list[dict[str, list[Any]]]:
# Separate entries: the caller sends the key with either scheme,
# not both at once.
return [{name: []} for name in self.name]

def get_security_definition(
self, auto_schema: openapi.AutoSchema | None = None
) -> dict[str, Any]:
return {
"type": "http",
"scheme": "bearer",
"description": (
"For cohort sync endpoints called by an external cohort "
"source, such as Amplitude."
),
}
) -> list[dict[str, Any]]:
return [
{
"type": "http",
"scheme": "bearer",
"description": (
"For cohort sync endpoints called by an external cohort "
"source, such as Amplitude."
),
},
{
"type": "http",
"scheme": "basic",
"description": (
"For cohort sync endpoints called by an external cohort "
"source that can only send Basic credentials, such as "
"Mixpanel. The key is the password; the username is "
"ignored."
),
},
]


# Tag definitions controlling the order and display of sections in the Swagger UI.
Expand Down
30 changes: 24 additions & 6 deletions api/cohorts/authentication.py
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
import typing
import base64
from contextlib import suppress

from django.contrib.auth.models import AnonymousUser
Expand All @@ -9,18 +9,36 @@


class CohortSyncKeyAuthentication(authentication.BaseAuthentication):
"""
Accepts a cohort sync key sent either as a Bearer token or as the
password of Basic credentials. Amplitude sends Bearer; Mixpanel's
webhook setup only offers a username/password form, so its customers
enter any username and the key as the password. The username is
ignored.
"""

def authenticate(
self, request: Request
) -> tuple[AnonymousUser, CohortSyncKey] | None:
header = request.headers.get("Authorization", "")
if not header.startswith("Bearer "):
if header.startswith("Bearer "):
raw_key = header.removeprefix("Bearer ")
elif header.startswith("Basic "):
try:
decoded = base64.b64decode(
header.removeprefix("Basic "), validate=True
).decode()
except ValueError:
# Covers malformed base64, header bytes outside ASCII, and
# decoded credentials that are not valid UTF-8.
raise exceptions.AuthenticationFailed("Invalid Basic credentials.")
# Split at the first colon, so a key containing colons survives.
_, _, raw_key = decoded.partition(":")
else:
return None

with suppress(CohortSyncKey.DoesNotExist):
key = typing.cast(
CohortSyncKey,
CohortSyncKey.objects.get_from_key(header.removeprefix("Bearer ")),
)
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
Expand Down
33 changes: 33 additions & 0 deletions api/cohorts/migrations/0004_mixpanel_source.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,33 @@
# Generated by Django 5.2.16 on 2026-08-20 10:06

from django.db import migrations, models


class Migration(migrations.Migration):

dependencies = [
("cohorts", "0003_cohort_sync_key"),
("environments", "0039_use_no_ssrf_url_field"),
("segments", "0032_add_segment_rules_data"),
]

operations = [
migrations.AddField(
model_name="cohort",
name="external_id",
field=models.CharField(blank=True, max_length=255, null=True),
),
migrations.AlterField(
model_name="cohort",
name="source_type",
field=models.CharField(
choices=[
("csv", "CSV"),
("amplitude", "Amplitude"),
("mixpanel", "Mixpanel"),
],
default="csv",
max_length=50,
),
),
]
22 changes: 21 additions & 1 deletion api/cohorts/models.py
Original file line number Diff line number Diff line change
@@ -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
Expand All @@ -8,6 +10,7 @@
class CohortSourceType(models.TextChoices):
CSV = "csv", "CSV"
AMPLITUDE = "amplitude", "Amplitude"
MIXPANEL = "mixpanel", "Mixpanel"


class Cohort(SoftDeleteExportableModel):
Expand All @@ -26,6 +29,11 @@ class Cohort(SoftDeleteExportableModel):
choices=CohortSourceType.choices,
default=CohortSourceType.CSV,
)
# 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)
Comment on lines +32 to +36

@coderabbitai coderabbitai Bot Aug 21, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🗄️ Data Integrity & Integration | 🟠 Major | 🏗️ Heavy lift

Make active cohort lookup and creation atomic.

Concurrent initial members requests can both observe no cohort and create separate active cohorts for the same environment, source type, and external ID. Later incremental requests may select only one duplicate, splitting memberships between cohorts.

Add a partial unique constraint for active cohorts with non-null external IDs, and use an atomic get-or-create flow that re-fetches after an integrity conflict.

📍 Affects 2 files
  • api/cohorts/models.py#L32-L36 (this comment)
  • api/cohorts/sync_views.py#L140-L153

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Deliberately dropped to keep v1 lean — we'll add the constraint and conflict recovery if duplicate creation is observed in practice.

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@gagantrivedi, understood. The concurrency risk remains deferred from this PR.

Would you like me to create a follow-up GitHub issue for the active-cohort uniqueness constraint and atomic conflict recovery?

🐇


✏️ Learnings added
Learnt from: gagantrivedi
Repo: Flagsmith/flagsmith PR: 8338
File: api/cohorts/models.py:32-36
Timestamp: 2026-08-24T04:25:26.515Z
Learning: For the Mixpanel cohort synchronisation v1, the maintainers deliberately defer the active-cohort uniqueness constraint and atomic conflict recovery for concurrent initial `members` requests. They will add this protection if duplicate cohort creation is observed in practice.

Note: Learnings are effective only in the context of similar code segments. To apply general coding guidelines, please configure review instructions. You can manage existing learnings in the app.

You are interacting with an AI system.

version = models.PositiveIntegerField(default=0)
created_at = models.DateTimeField(auto_now_add=True)
# Deletion drains memberships from the identity store first; the cohort is
Expand All @@ -48,7 +56,19 @@ class Meta:
]


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,
Expand Down
21 changes: 21 additions & 0 deletions api/cohorts/serializers.py
Original file line number Diff line number Diff line change
Expand Up @@ -60,6 +60,27 @@ class AmplitudeListSerializer(serializers.Serializer[None]):
name = serializers.CharField(max_length=2000)


class MixpanelMemberSerializer(serializers.Serializer[None]):
# Length mirrors CohortMembership.identifier.
mixpanel_distinct_id = serializers.CharField(max_length=2000)
Comment on lines +63 to +65

@coderabbitai coderabbitai Bot Aug 21, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🗄️ Data Integrity & Integration | 🟠 Major | ⚡ Quick win

Reject identifiers that exceed the identity-store byte limit.

Line 65 validates 2000 characters. The downstream identity store permits only 1024 bytes. A multibyte identifier can pass validation, receive a success response, and then fail during identity processing.

Validate the UTF-8 byte length before calling the membership service.

Proposed validation
 class MixpanelMemberSerializer(serializers.Serializer[None]):
     # Length mirrors CohortMembership.identifier.
     mixpanel_distinct_id = serializers.CharField(max_length=2000)
+
+    def validate_mixpanel_distinct_id(self, value: str) -> str:
+        if len(value.encode("utf-8")) > 1024:
+            raise serializers.ValidationError(
+                "Ensure this field is no longer than 1024 bytes."
+            )
+        return value
📝 Committable suggestion

‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.

Suggested change
class MixpanelMemberSerializer(serializers.Serializer[None]):
# Length mirrors CohortMembership.identifier.
mixpanel_distinct_id = serializers.CharField(max_length=2000)
class MixpanelMemberSerializer(serializers.Serializer[None]):
# Length mirrors CohortMembership.identifier.
mixpanel_distinct_id = serializers.CharField(max_length=2000)
def validate_mixpanel_distinct_id(self, value: str) -> str:
if len(value.encode("utf-8")) > 1024:
raise serializers.ValidationError(
"Ensure this field is no longer than 1024 bytes."
)
return value

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Deferred for now — tracked as a TODO together with the same byte-length check for Amplitude user_ids and CSV uploads.

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@gagantrivedi, acknowledged. The existing TODO tracks the shared 1024-byte validation requirement for Mixpanel identifiers, Amplitude user_ids, and CSV uploads. This finding remains deferred for this pull request.

🐇


🧠 Learnings used
Learnt from: gagantrivedi
Repo: Flagsmith/flagsmith PR: 8211
File: api/cohorts/migrations/0001_initial.py:121-126
Timestamp: 2026-08-04T09:34:42.123Z
Learning: For the Flagsmith Django cohort API, validate cohort identifiers before persistence and reject identifiers exceeding 1024 bytes with HTTP 400, matching Edge DynamoDB constraints. Do not introduce digest-based uniqueness for cohort identifiers: Identity already enforces unique_together(environment, identifier) on the 2,000-character field, and digesting would unnecessarily complicate identity lookup and evaluation without addressing a cohort-specific risk.

You are interacting with an AI system.



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)


class MixpanelWebhookSerializer(serializers.Serializer[None]):
# "members" carries the full membership on the first sync;
# "add_members"/"remove_members" carry changes since the last sync.
action = serializers.ChoiceField(
choices=["members", "add_members", "remove_members"]
)
parameters = MixpanelParametersSerializer()


class CohortSyncMembersSerializer(serializers.Serializer[None]):
# Child length mirrors CohortMembership.identifier.
# TODO: this counts characters, but identity data is stored with a
Expand Down
29 changes: 27 additions & 2 deletions api/cohorts/services.py
Original file line number Diff line number Diff line change
Expand Up @@ -90,6 +90,7 @@ def create_cohort(
name: str,
description: str | None = None,
source_type: CohortSourceType = CohortSourceType.CSV,
external_id: str | None = None,
) -> Cohort:
with transaction.atomic():
segment = Segment.objects.create(
Expand All @@ -100,7 +101,10 @@ def create_cohort(
)
rule = SegmentRule.objects.create(segment=segment, type=SegmentRule.ALL_RULE)
cohort: Cohort = Cohort.objects.create(
environment=environment, segment=segment, source_type=source_type
environment=environment,
segment=segment,
source_type=source_type,
external_id=external_id,
)
Condition.objects.create(
rule=rule,
Expand All @@ -124,10 +128,16 @@ def create_cohort_for_source(
environment: "Environment",
name: str,
source_type: CohortSourceType,
external_id: str | None = None,
) -> Cohort:
"""Create a cohort on behalf of an external source, where no Flagsmith
user is acting."""
cohort = create_cohort(environment=environment, name=name, source_type=source_type)
cohort = create_cohort(
environment=environment,
name=name,
source_type=source_type,
external_id=external_id,
)
# Nothing records a user for these calls, so the audit log that Flagsmith
# derives from historical records is skipped — and with it the environment
# document rebuild that makes the new segment visible to SDKs. Write the
Expand All @@ -145,6 +155,21 @@ def create_cohort_for_source(
return cohort


def get_cohort_for_source(
*,
environment: "Environment",
source_type: CohortSourceType,
external_id: str,
) -> Cohort | None:
cohort: Cohort | None = Cohort.objects.filter(
environment=environment,
source_type=source_type,
external_id=external_id,
deletion_requested_at__isnull=True,
).first()
return cohort


def add_cohort_members(cohort: Cohort, identifiers: "typing.Iterable[str]") -> None:
from cohorts.tasks import apply_cohort_membership_deltas

Expand Down
8 changes: 6 additions & 2 deletions api/cohorts/sync_urls.py
Original file line number Diff line number Diff line change
@@ -1,11 +1,15 @@
from django.urls import path
from rest_framework.routers import SimpleRouter

from cohorts.sync_views import AmplitudeCohortSyncViewSet
from cohorts.sync_views import AmplitudeCohortSyncViewSet, MixpanelCohortSyncView

app_name = "cohort-sync"

# SimpleRouter: nothing here is browsed by a person.
router = SimpleRouter()
router.register(r"amplitude/lists", AmplitudeCohortSyncViewSet, basename="amplitude")

urlpatterns = router.urls
urlpatterns = [
path("mixpanel/webhook/", MixpanelCohortSyncView.as_view(), name="mixpanel"),
*router.urls,
]
Loading
Loading