feat(cohorts): Amplitude cohort sync endpoints and sync keys

api/api/openapi.py

@@ -163,6 +163,23 @@ def get_security_definition(
         }
 
 
+class CohortSyncKeyAuthenticationExtension(OpenApiAuthenticationExtension):  # type: ignore[no-untyped-call]
+    target_class = "cohorts.authentication.CohortSyncKeyAuthentication"
+    name = "Cohort Sync Key"
+
+    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."
+            ),
+        }
+
+
 # Tag definitions controlling the order and display of sections in the Swagger UI.
 TAGS: list[dict[str, str]] = [
     {

api/api/urls/v1.py

@@ -37,6 +37,7 @@
         r"^multivariate/", include("features.multivariate.urls"), name="multivariate"
     ),
     re_path(r"^segments/", include("segments.urls"), name="segments"),
+    re_path(r"^cohort-sync/", include("cohorts.sync_urls"), name="cohort-sync"),
     re_path(r"^users/", include("users.urls")),
     re_path(r"^e2etests/", include("e2etests.urls")),
     re_path(r"^audit/", include("audit.urls")),

api/cohorts/authentication.py

@@ -0,0 +1,34 @@
+import typing
+from contextlib import suppress
+
+from django.contrib.auth.models import AnonymousUser
+from rest_framework import authentication, exceptions
+from rest_framework.request import Request
+
+from cohorts.models import CohortSyncKey
+
+
+class CohortSyncKeyAuthentication(authentication.BaseAuthentication):
+    def authenticate(
+        self, request: Request
+    ) -> tuple[AnonymousUser, CohortSyncKey] | None:
+        header = request.headers.get("Authorization", "")
+        if not header.startswith("Bearer "):
+            return None
+
+        with suppress(CohortSyncKey.DoesNotExist):
+            key = typing.cast(
+                CohortSyncKey,
+                CohortSyncKey.objects.get_from_key(header.removeprefix("Bearer ")),
+            )
+            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
+                # rather than a user.
+                return AnonymousUser(), key
+
+        raise exceptions.AuthenticationFailed("Valid cohort sync key not found.")
+
+    def authenticate_header(self, request: Request) -> str:
+        # Makes missing or invalid credentials a 401 rather than DRF's default 403.
+        return "Bearer"

api/cohorts/migrations/0003_cohort_sync_key.py

@@ -0,0 +1,92 @@
+# Generated by Django 5.2.16 on 2026-08-14 08:31
+
+import django.db.models.deletion
+from django.conf import settings
+from django.db import migrations, models
+
+
+class Migration(migrations.Migration):
+
+    dependencies = [
+        ("cohorts", "0002_cohort_deletion_requested_at"),
+        ("environments", "0039_use_no_ssrf_url_field"),
+        migrations.swappable_dependency(settings.AUTH_USER_MODEL),
+    ]
+
+    operations = [
+        migrations.AlterField(
+            model_name="cohort",
+            name="source_type",
+            field=models.CharField(
+                choices=[("csv", "CSV"), ("amplitude", "Amplitude")],
+                default="csv",
+                max_length=50,
+            ),
+        ),
+        migrations.CreateModel(
+            name="CohortSyncKey",
+            fields=[
+                (
+                    "id",
+                    models.CharField(
+                        editable=False,
+                        max_length=150,
+                        primary_key=True,
+                        serialize=False,
+                        unique=True,
+                    ),
+                ),
+                ("prefix", models.CharField(editable=False, max_length=8, unique=True)),
+                ("hashed_key", models.CharField(editable=False, max_length=150)),
+                ("created", models.DateTimeField(auto_now_add=True, db_index=True)),
+                (
+                    "name",
+                    models.CharField(
+                        default=None,
+                        help_text="A free-form name for the API key. Need not be unique. 50 characters max.",
+                        max_length=50,
+                    ),
+                ),
+                (
+                    "revoked",
+                    models.BooleanField(
+                        blank=True,
+                        default=False,
+                        help_text="If the API key is revoked, clients cannot use it anymore. (This cannot be undone.)",
+                    ),
+                ),
+                (
+                    "expiry_date",
+                    models.DateTimeField(
+                        blank=True,
+                        help_text="Once API key expires, clients cannot use it anymore.",
+                        null=True,
+                        verbose_name="Expires",
+                    ),
+                ),
+                (
+                    "created_by",
+                    models.ForeignKey(
+                        blank=True,
+                        null=True,
+                        on_delete=django.db.models.deletion.SET_NULL,
+                        to=settings.AUTH_USER_MODEL,
+                    ),
+                ),
+                (
+                    "environment",
+                    models.ForeignKey(
+                        on_delete=django.db.models.deletion.CASCADE,
+                        related_name="cohort_sync_keys",
+                        to="environments.environment",
+                    ),
+                ),
+            ],
+            options={
+                "verbose_name": "cohort sync key",
+                "verbose_name_plural": "cohort sync keys",
+                "ordering": ("-created",),
+                "abstract": False,
+            },
+        ),
+    ]

api/cohorts/models.py

@@ -1,11 +1,13 @@
 from django.db import models
+from rest_framework_api_key.models import AbstractAPIKey
 
 from cohorts.constants import COHORT_SYSTEM_TRAIT_KEY_PREFIX
 from core.models import SoftDeleteExportableModel
 
 
 class CohortSourceType(models.TextChoices):
     CSV = "csv", "CSV"
+    AMPLITUDE = "amplitude", "Amplitude"
 
 
 class Cohort(SoftDeleteExportableModel):
@@ -46,6 +48,21 @@ class Meta:
         ]
 
 
+class CohortSyncKey(AbstractAPIKey):
+    environment = models.ForeignKey(
+        "environments.Environment",
+        on_delete=models.CASCADE,
+        related_name="cohort_sync_keys",
+    )
+    created_by = models.ForeignKey(
+        "users.FFAdminUser", on_delete=models.SET_NULL, null=True, blank=True
+    )
+
+    class Meta(AbstractAPIKey.Meta):
+        verbose_name = "cohort sync key"
+        verbose_name_plural = "cohort sync keys"
+
+
 class CohortMembershipState(models.TextChoices):
     PENDING_ADD = "pending_add", "Pending add"
     APPLIED = "applied", "Applied"

api/cohorts/permissions.py

@@ -7,13 +7,20 @@
 from rest_framework.request import Request
 from rest_framework.views import APIView
 
+from cohorts.models import CohortSyncKey
 from environments.models import Environment
 from organisations.subscriptions.constants import SubscriptionPlanFamily
 from organisations.subscriptions.permissions import require_minimum_plan
 from users.models import FFAdminUser
 
 _READ_ACTIONS = ("list", "retrieve")
 
+
+class HasCohortSyncKey(BasePermission):
+    def has_permission(self, request: Request, view: APIView) -> bool:
+        return isinstance(request.auth, CohortSyncKey)
+
+
 _MinimumStartupPlan = require_minimum_plan(SubscriptionPlanFamily.START_UP)
 
 

api/cohorts/serializers.py

@@ -2,7 +2,7 @@
 
 from rest_framework import serializers
 
-from cohorts.models import Cohort
+from cohorts.models import Cohort, CohortSyncKey
 from cohorts.services import create_cohort
 
 
@@ -33,3 +33,44 @@ def create(self, validated_data: dict[str, typing.Any]) -> Cohort:
             name=segment_data["name"],
             description=segment_data.get("description"),
         )
+
+
+class CohortSyncKeySerializer(serializers.ModelSerializer[CohortSyncKey]):
+    key = serializers.SerializerMethodField()
+    # The model field carries a default, which DRF would read as optional;
+    # saving without a name fails at the database instead.
+    name = serializers.CharField(max_length=50)
+
+    class Meta:
+        model = CohortSyncKey
+        fields = ("prefix", "name", "created", "key")
+        read_only_fields = ("prefix", "created")
+
+    def create(self, validated_data: dict[str, typing.Any]) -> CohortSyncKey:
+        key, self._generated_key = CohortSyncKey.objects.create_key(**validated_data)
+        return typing.cast(CohortSyncKey, key)
+
+    def get_key(self, instance: CohortSyncKey) -> str | None:
+        # The plaintext key exists only in the create response; it is
+        # unrecoverable afterwards.
+        return getattr(self, "_generated_key", None)
+
+
+class AmplitudeListSerializer(serializers.Serializer[None]):
+    name = serializers.CharField(max_length=2000)
+
+
+def _validate_identifier_byte_length(value: str) -> None:
+    # Edge identifiers are DynamoDB sort keys, capped at 1024 bytes.
+    if len(value.encode()) > 1024:
+        raise serializers.ValidationError(
+            "Ensure this identifier has no more than 1024 bytes."
+        )
+
+
+class CohortSyncMembersSerializer(serializers.Serializer[None]):
+    # TODO: apply the same byte-length check to CSV uploads.
+    user_ids = serializers.ListField(
+        child=serializers.CharField(validators=[_validate_identifier_byte_length]),
+        min_length=1,
+    )

api/cohorts/services.py

@@ -6,9 +6,17 @@
 from django.utils import timezone
 from flag_engine.segments.constants import IS_SET
 
+from audit.constants import SEGMENT_CREATED_MESSAGE
+from audit.models import AuditLog
+from audit.related_object_type import RelatedObjectType
 from cohorts.constants import COHORT_MEMBERSHIP_APPLY_BATCH_SIZE
 from cohorts.metrics import flagsmith_cohorts_membership_deltas_applied_total
-from cohorts.models import Cohort, CohortMembership, CohortMembershipState
+from cohorts.models import (
+    Cohort,
+    CohortMembership,
+    CohortMembershipState,
+    CohortSourceType,
+)
 from core.dataclasses import AuthorData
 from environments.identities.system_traits import (
     set_system_trait,
@@ -81,6 +89,7 @@ def create_cohort(
     environment: "Environment",
     name: str,
     description: str | None = None,
+    source_type: CohortSourceType = CohortSourceType.CSV,
 ) -> Cohort:
     with transaction.atomic():
         segment = Segment.objects.create(
@@ -90,7 +99,9 @@ def create_cohort(
             managed_by=SegmentManagedBy.COHORT,
         )
         rule = SegmentRule.objects.create(segment=segment, type=SegmentRule.ALL_RULE)
-        cohort: Cohort = Cohort.objects.create(environment=environment, segment=segment)
+        cohort: Cohort = Cohort.objects.create(
+            environment=environment, segment=segment, source_type=source_type
+        )
         Condition.objects.create(
             rule=rule,
             operator=IS_SET,
@@ -108,6 +119,86 @@ def create_cohort(
     return cohort
 
 
+def create_cohort_for_source(
+    *,
+    environment: "Environment",
+    name: str,
+    source_type: CohortSourceType,
+) -> 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)
+    # 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
+    # record here instead, naming the source that asked for the cohort.
+    AuditLog.objects.create(
+        environment=environment,
+        project=environment.project,
+        related_object_id=cohort.segment_id,
+        related_object_type=RelatedObjectType.SEGMENT.name,
+        log=(
+            f"{SEGMENT_CREATED_MESSAGE % cohort.segment.name} "
+            f"(via {CohortSourceType(source_type).label} cohort sync)"
+        ),
+    )
+    return cohort
+
+
+def add_cohort_members(cohort: Cohort, identifiers: "typing.Iterable[str]") -> None:
+    from cohorts.tasks import apply_cohort_membership_deltas
+
+    rows = [
+        CohortMembership(cohort=cohort, identifier=identifier)
+        for identifier in set(identifiers)
+    ]
+    with transaction.atomic():
+        # Re-adding a member is a no-op end to end: an applied row flips back
+        # to pending and the identity write it triggers is idempotent.
+        CohortMembership.objects.bulk_create(
+            rows,
+            # Postgres rejects a statement carrying more than 65535 bind
+            # parameters, which a single large batch would exceed.
+            batch_size=1000,
+            update_conflicts=True,
+            unique_fields=["cohort", "identifier"],
+            update_fields=["state", "updated_at"],
+        )
+        apply_cohort_membership_deltas.delay(kwargs={"cohort_id": cohort.id})
+    logger.info(
+        "membership.adds_received",
+        cohort__id=cohort.id,
+        environment__id=cohort.environment_id,
+        deltas__count=len(rows),
+    )
+
+
+def remove_cohort_members(cohort: Cohort, identifiers: "typing.Iterable[str]") -> None:
+    from cohorts.tasks import apply_cohort_membership_deltas
+
+    unique_identifiers = list(set(identifiers))
+    matched = 0
+    with transaction.atomic():
+        # Postgres rejects a statement carrying more than 65535 bind
+        # parameters, which a single large identifier list would exceed.
+        for start in range(0, len(unique_identifiers), 1000):
+            # Removing a non-member is a no-op: only existing rows flip.
+            matched += CohortMembership.objects.filter(
+                cohort=cohort,
+                identifier__in=unique_identifiers[start : start + 1000],
+            ).update(
+                state=CohortMembershipState.PENDING_REMOVE, updated_at=timezone.now()
+            )
+        apply_cohort_membership_deltas.delay(kwargs={"cohort_id": cohort.id})
+    logger.info(
+        "membership.removals_received",
+        cohort__id=cohort.id,
+        environment__id=cohort.environment_id,
+        deltas__count=len(unique_identifiers),
+        members__matched=matched,
+    )
+
+
 def delete_cohort(cohort: Cohort) -> None:
     from cohorts.tasks import apply_cohort_membership_deltas