From 99b16152db403641659b0f6879e63f5c1c316e01 Mon Sep 17 00:00:00 2001 From: Gagan Trivedi Date: Fri, 14 Aug 2026 14:28:13 +0530 Subject: [PATCH 1/2] feat(cohorts): add Amplitude cohort sync endpoints and sync keys --- api/api/urls/v1.py | 1 + api/cohorts/authentication.py | 58 ++++ .../migrations/0003_cohort_sync_key.py | 92 +++++ api/cohorts/models.py | 22 ++ api/cohorts/serializers.py | 31 +- api/cohorts/services.py | 56 ++- api/cohorts/sync_urls.py | 10 + api/cohorts/sync_views.py | 87 +++++ api/cohorts/urls.py | 3 +- api/cohorts/views.py | 45 ++- api/tests/unit/cohorts/conftest.py | 32 +- api/tests/unit/cohorts/test_sync_views.py | 325 ++++++++++++++++++ api/tests/unit/cohorts/test_views.py | 95 ++++- .../observability/_events-catalogue.md | 20 +- 14 files changed, 864 insertions(+), 13 deletions(-) create mode 100644 api/cohorts/authentication.py create mode 100644 api/cohorts/migrations/0003_cohort_sync_key.py create mode 100644 api/cohorts/sync_urls.py create mode 100644 api/cohorts/sync_views.py create mode 100644 api/tests/unit/cohorts/test_sync_views.py diff --git a/api/api/urls/v1.py b/api/api/urls/v1.py index 8269ecbcfff8..91f9197145c2 100644 --- a/api/api/urls/v1.py +++ b/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")), diff --git a/api/cohorts/authentication.py b/api/cohorts/authentication.py new file mode 100644 index 000000000000..eb93824b2104 --- /dev/null +++ b/api/cohorts/authentication.py @@ -0,0 +1,58 @@ +import typing +from contextlib import suppress + +from rest_framework import authentication, exceptions +from rest_framework.request import Request + +from cohorts.models import CohortSyncKey + + +class CohortSyncKeyUser: + """Stand-in request user for machine calls authenticated by a + CohortSyncKey; carries no permissions of its own.""" + + # Named `sync_key`, not `key`: the audit signal that stamps history rows + # duck-types master API keys via `request.user.key`. + def __init__(self, sync_key: CohortSyncKey) -> None: + self.sync_key = sync_key + + def __str__(self) -> str: + return self.sync_key.name + + @property + def is_authenticated(self) -> bool: + return True + + @property + def pk(self) -> str: + return self.sync_key.id + + @property + def is_master_api_key_user(self) -> bool: + # The discriminator history/audit code uses for machine callers: + # without it, historical records try to store this object as the + # acting FFAdminUser and fail. + return True + + +class CohortSyncKeyAuthentication(authentication.BaseAuthentication): + def authenticate( + self, request: Request + ) -> tuple[CohortSyncKeyUser, 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: + return CohortSyncKeyUser(key), key + + raise exceptions.AuthenticationFailed("Valid cohort sync key not found.") + + def authenticate_header(self, request: Request) -> str: + # Makes missing credentials a 401 rather than DRF's default 403. + return "Bearer" diff --git a/api/cohorts/migrations/0003_cohort_sync_key.py b/api/cohorts/migrations/0003_cohort_sync_key.py new file mode 100644 index 000000000000..5b1a53e30e11 --- /dev/null +++ b/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": "API key", + "verbose_name_plural": "API keys", + "ordering": ("-created",), + "abstract": False, + }, + ), + ] diff --git a/api/cohorts/models.py b/api/cohorts/models.py index 605b67be0526..ab74704439ec 100644 --- a/api/cohorts/models.py +++ b/api/cohorts/models.py @@ -1,4 +1,5 @@ from django.db import models +from rest_framework_api_key.models import AbstractAPIKey, APIKeyManager from cohorts.constants import COHORT_SYSTEM_TRAIT_KEY_PREFIX from core.models import SoftDeleteExportableModel @@ -6,6 +7,7 @@ class CohortSourceType(models.TextChoices): CSV = "csv", "CSV" + AMPLITUDE = "amplitude", "Amplitude" class Cohort(SoftDeleteExportableModel): @@ -46,6 +48,26 @@ class Meta: ] +class CohortSyncKeyManager(APIKeyManager): + pass + + +class CohortSyncKey(AbstractAPIKey): + """Bearer credential an external cohort source uses to call the + cohort-sync endpoints; scopes every call to one environment.""" + + 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 + ) + + objects = CohortSyncKeyManager() # type: ignore[misc] + + class CohortMembershipState(models.TextChoices): PENDING_ADD = "pending_add", "Pending add" APPLIED = "applied", "Applied" diff --git a/api/cohorts/serializers.py b/api/cohorts/serializers.py index 780216bb9f4a..a74e7da5d73c 100644 --- a/api/cohorts/serializers.py +++ b/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,32 @@ 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() + + 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) + + +class CohortSyncMembersSerializer(serializers.Serializer[None]): + # Child length mirrors CohortMembership.identifier. + user_ids = serializers.ListField( + child=serializers.CharField(max_length=2000), min_length=1 + ) diff --git a/api/cohorts/services.py b/api/cohorts/services.py index c2ba9f3bd111..be9421e6d72c 100644 --- a/api/cohorts/services.py +++ b/api/cohorts/services.py @@ -8,7 +8,12 @@ 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.dynamodb import DynamoIdentityWrapper from segments.models import Condition, Segment, SegmentManagedBy, SegmentRule @@ -84,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( @@ -93,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, @@ -115,6 +123,50 @@ def edge_sync_enabled(project: "Project") -> bool: return bool(project.enable_dynamo_db and DynamoIdentityWrapper().is_enabled) +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, + 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.deltas_received", + cohort__id=cohort.id, + environment__id=cohort.environment_id, + action="add", + deltas__count=len(rows), + ) + + +def remove_cohort_members(cohort: Cohort, identifiers: "typing.Iterable[str]") -> None: + from cohorts.tasks import apply_cohort_membership_deltas + + with transaction.atomic(): + # Removing a non-member is a no-op: only existing rows flip. + updated = CohortMembership.objects.filter( + cohort=cohort, identifier__in=set(identifiers) + ).update(state=CohortMembershipState.PENDING_REMOVE, updated_at=timezone.now()) + apply_cohort_membership_deltas.delay(kwargs={"cohort_id": cohort.id}) + logger.info( + "membership.deltas_received", + cohort__id=cohort.id, + environment__id=cohort.environment_id, + action="remove", + deltas__count=updated, + ) + + def delete_cohort(cohort: Cohort) -> None: from cohorts.tasks import apply_cohort_membership_deltas diff --git a/api/cohorts/sync_urls.py b/api/cohorts/sync_urls.py new file mode 100644 index 000000000000..8c15ded24925 --- /dev/null +++ b/api/cohorts/sync_urls.py @@ -0,0 +1,10 @@ +from rest_framework.routers import DefaultRouter + +from cohorts.sync_views import AmplitudeCohortSyncViewSet + +app_name = "cohort-sync" + +router = DefaultRouter() +router.register(r"amplitude/lists", AmplitudeCohortSyncViewSet, basename="amplitude") + +urlpatterns = router.urls diff --git a/api/cohorts/sync_views.py b/api/cohorts/sync_views.py new file mode 100644 index 000000000000..a0b93bfd2b9d --- /dev/null +++ b/api/cohorts/sync_views.py @@ -0,0 +1,87 @@ +import uuid as uuid_module + +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 +from rest_framework.permissions import IsAuthenticated +from rest_framework.request import Request +from rest_framework.response import Response + +from cohorts import services +from cohorts.authentication import CohortSyncKeyAuthentication +from cohorts.models import Cohort, CohortSourceType, CohortSyncKey +from cohorts.serializers import ( + AmplitudeListSerializer, + CohortSyncMembersSerializer, +) +from projects.exceptions import DynamoNotEnabledError + +_LIST_RESPONSE = inline_serializer( + "AmplitudeListResponse", {"list_id": serializers.UUIDField()} +) + + +@extend_schema_view( + create=extend_schema( + description=( + "Called by Amplitude once per cohort sync setup; creates the " + "backing cohort and returns its identifier as the list ID." + ), + request=AmplitudeListSerializer, + responses={200: _LIST_RESPONSE}, + ), + add=extend_schema(request=CohortSyncMembersSerializer, responses={200: None}), + remove=extend_schema(request=CohortSyncMembersSerializer, responses={200: None}), +) +class AmplitudeCohortSyncViewSet(viewsets.ViewSet): + authentication_classes = [CohortSyncKeyAuthentication] + permission_classes = [IsAuthenticated] + + def create(self, request: Request) -> Response: + serializer = AmplitudeListSerializer(data=request.data) + serializer.is_valid(raise_exception=True) + environment = self._get_key(request).environment + if not services.edge_sync_enabled(environment.project): + raise DynamoNotEnabledError() + cohort = services.create_cohort( + environment=environment, + name=serializer.validated_data["name"], + source_type=CohortSourceType.AMPLITUDE, + ) + return Response({"list_id": str(cohort.uuid)}) + + @action(detail=True, methods=["POST"]) + def add(self, request: Request, pk: str) -> Response: + cohort = self._get_cohort(request, pk) + serializer = CohortSyncMembersSerializer(data=request.data) + serializer.is_valid(raise_exception=True) + services.add_cohort_members(cohort, serializer.validated_data["user_ids"]) + return Response() + + @action(detail=True, methods=["POST"]) + def remove(self, request: Request, pk: str) -> Response: + cohort = self._get_cohort(request, pk) + serializer = CohortSyncMembersSerializer(data=request.data) + serializer.is_valid(raise_exception=True) + services.remove_cohort_members(cohort, serializer.validated_data["user_ids"]) + return Response() + + def _get_key(self, request: Request) -> CohortSyncKey: + assert isinstance(request.auth, CohortSyncKey) + return request.auth + + def _get_cohort(self, request: Request, pk: str) -> Cohort: + try: + list_uuid = uuid_module.UUID(pk) + except ValueError: + raise NotFound("List not found.") + cohort: Cohort | None = Cohort.objects.filter( + uuid=list_uuid, + environment=self._get_key(request).environment, + source_type=CohortSourceType.AMPLITUDE, + deletion_requested_at__isnull=True, + ).first() + if cohort is None: + raise NotFound("List not found.") + return cohort diff --git a/api/cohorts/urls.py b/api/cohorts/urls.py index 1cb2c65460ae..62a0b963ee7f 100644 --- a/api/cohorts/urls.py +++ b/api/cohorts/urls.py @@ -1,10 +1,11 @@ from rest_framework.routers import DefaultRouter -from cohorts.views import CohortViewSet +from cohorts.views import CohortSyncKeyViewSet, CohortViewSet app_name = "cohorts" router = DefaultRouter() +router.register(r"sync-keys", CohortSyncKeyViewSet, basename="sync-keys") router.register(r"", CohortViewSet, basename="cohorts") urlpatterns = router.urls diff --git a/api/cohorts/views.py b/api/cohorts/views.py index eda7c209a3d9..f27dc3982693 100644 --- a/api/cohorts/views.py +++ b/api/cohorts/views.py @@ -1,14 +1,16 @@ from django.db.models import QuerySet from drf_spectacular.utils import extend_schema, extend_schema_view -from rest_framework import mixins, status +from rest_framework import mixins, status, viewsets from rest_framework.permissions import IsAuthenticated from rest_framework.request import Request from rest_framework.response import Response +from rest_framework.serializers import BaseSerializer from cohorts import services -from cohorts.models import Cohort +from cohorts.models import Cohort, CohortSyncKey from cohorts.permissions import CohortPermission, CohortPlanPermission -from cohorts.serializers import CohortSerializer +from cohorts.serializers import CohortSerializer, CohortSyncKeySerializer +from environments.models import Environment from environments.views import NestedEnvironmentViewSet from projects.exceptions import DynamoNotEnabledError @@ -62,3 +64,40 @@ def get_queryset(self) -> QuerySet[Cohort]: def destroy(self, request: Request, *args: object, **kwargs: object) -> Response: services.delete_cohort(self.get_object()) return Response(status=status.HTTP_202_ACCEPTED) + + +@extend_schema_view( + create=extend_schema( + description=( + "Create a cohort sync key. The response is the only time the " + "plaintext key is available." + ) + ), + destroy=extend_schema(description="Revoke a cohort sync key."), +) +class CohortSyncKeyViewSet( + mixins.ListModelMixin, + mixins.CreateModelMixin, + mixins.DestroyModelMixin, + viewsets.GenericViewSet[CohortSyncKey], +): + serializer_class = CohortSyncKeySerializer + pagination_class = None + permission_classes = [IsAuthenticated, CohortPlanPermission, CohortPermission] + lookup_field = "prefix" + + def get_queryset(self) -> QuerySet[CohortSyncKey]: + return CohortSyncKey.objects.filter( + environment__api_key=self.kwargs.get("environment_api_key"), + revoked=False, + ).order_by("-created") + + def perform_create(self, serializer: BaseSerializer[CohortSyncKey]) -> None: + environment = Environment.objects.get( + api_key=self.kwargs.get("environment_api_key") + ) + serializer.save(environment=environment, created_by=self.request.user) + + def perform_destroy(self, instance: CohortSyncKey) -> None: + instance.revoked = True + instance.save(update_fields=["revoked"]) diff --git a/api/tests/unit/cohorts/conftest.py b/api/tests/unit/cohorts/conftest.py index adf99e3cbeb7..96467c9204e8 100644 --- a/api/tests/unit/cohorts/conftest.py +++ b/api/tests/unit/cohorts/conftest.py @@ -1,6 +1,8 @@ +import typing + import pytest -from cohorts.models import Cohort +from cohorts.models import Cohort, CohortSourceType, CohortSyncKey from environments.models import Environment from projects.models import Project from segments.models import Segment @@ -12,6 +14,34 @@ def cohort(environment: Environment, segment: Segment) -> Cohort: return cohort +@pytest.fixture() +def cohort_sync_key( + dynamo_enabled_project_environment_one: Environment, +) -> typing.Tuple[CohortSyncKey, str]: + return typing.cast( + typing.Tuple[CohortSyncKey, str], + CohortSyncKey.objects.create_key( + name="test key", environment=dynamo_enabled_project_environment_one + ), + ) + + +@pytest.fixture() +def amplitude_cohort( + dynamo_enabled_project: Project, + dynamo_enabled_project_environment_one: Environment, +) -> Cohort: + segment = Segment.objects.create( + name="amplitude segment", project=dynamo_enabled_project + ) + cohort: Cohort = Cohort.objects.create( + environment=dynamo_enabled_project_environment_one, + segment=segment, + source_type=CohortSourceType.AMPLITUDE, + ) + return cohort + + @pytest.fixture() def edge_cohort( dynamo_enabled_project: Project, diff --git a/api/tests/unit/cohorts/test_sync_views.py b/api/tests/unit/cohorts/test_sync_views.py new file mode 100644 index 000000000000..d74181b158dd --- /dev/null +++ b/api/tests/unit/cohorts/test_sync_views.py @@ -0,0 +1,325 @@ +import typing + +from django.urls import reverse +from django.utils import timezone +from flag_engine.segments.constants import IS_SET +from rest_framework import status +from rest_framework.test import APIClient + +from cohorts.models import ( + Cohort, + CohortMembership, + CohortMembershipState, + CohortSourceType, + CohortSyncKey, +) +from environments.dynamodb import DynamoIdentityWrapper +from environments.models import Environment + +_KeyAndPlaintext = typing.Tuple[CohortSyncKey, str] + + +def _authenticated_client(plaintext_key: str) -> APIClient: + client = APIClient() + client.credentials(HTTP_AUTHORIZATION=f"Bearer {plaintext_key}") + return client + + +def test_amplitude_create_list__valid_key__creates_amplitude_cohort( + cohort_sync_key: _KeyAndPlaintext, + dynamodb_identity_wrapper: DynamoIdentityWrapper, +) -> None: + # Given + key, plaintext = cohort_sync_key + client = _authenticated_client(plaintext) + url = reverse("api-v1:cohort-sync:amplitude-list") + + # When + response = client.post( + url, data={"name": "[Amplitude] Beta users: 1234"}, format="json" + ) + + # Then + assert response.status_code == status.HTTP_200_OK + cohort = Cohort.objects.get(uuid=response.json()["list_id"]) + assert cohort.environment == key.environment + assert cohort.source_type == CohortSourceType.AMPLITUDE + assert cohort.segment.name == "[Amplitude] Beta users: 1234" + condition = cohort.segment.rules.get().conditions.get() + assert condition.operator == IS_SET + assert condition.property == cohort.system_trait_key + + +def test_amplitude_create_list__non_edge_environment__returns_400( + environment: Environment, +) -> None: + # Given + _, plaintext = CohortSyncKey.objects.create_key( + name="core key", environment=environment + ) + client = _authenticated_client(plaintext) + url = reverse("api-v1:cohort-sync:amplitude-list") + + # When + response = client.post(url, data={"name": "Beta users"}, format="json") + + # Then + assert response.status_code == status.HTTP_400_BAD_REQUEST + assert not Cohort.objects.exists() + + +def test_amplitude_create_list__missing_credentials__returns_401( + db: None, +) -> None: + # Given + url = reverse("api-v1:cohort-sync:amplitude-list") + + # When + response = APIClient().post(url, data={"name": "Beta users"}, format="json") + + # Then + assert response.status_code == status.HTTP_401_UNAUTHORIZED + + +def test_amplitude_create_list__unknown_key__returns_401( + db: None, +) -> None: + # Given + client = _authenticated_client("not-a-key") + url = reverse("api-v1:cohort-sync:amplitude-list") + + # When + response = client.post(url, data={"name": "Beta users"}, format="json") + + # Then + assert response.status_code == status.HTTP_401_UNAUTHORIZED + + +def test_amplitude_create_list__revoked_key__returns_401( + cohort_sync_key: _KeyAndPlaintext, +) -> None: + # Given + key, plaintext = cohort_sync_key + key.revoked = True + key.save() + client = _authenticated_client(plaintext) + url = reverse("api-v1:cohort-sync:amplitude-list") + + # When + response = client.post(url, data={"name": "Beta users"}, format="json") + + # Then + assert response.status_code == status.HTTP_401_UNAUTHORIZED + + +def test_amplitude_add_members__new_identifiers__applies_memberships( + cohort_sync_key: _KeyAndPlaintext, + amplitude_cohort: Cohort, + dynamodb_identity_wrapper: DynamoIdentityWrapper, +) -> None: + # Given + _, plaintext = cohort_sync_key + client = _authenticated_client(plaintext) + url = reverse( + "api-v1:cohort-sync:amplitude-add", kwargs={"pk": str(amplitude_cohort.uuid)} + ) + + # When + response = client.post( + url, data={"user_ids": ["user-1", "user-2", "user-1"]}, format="json" + ) + + # Then + assert response.status_code == status.HTTP_200_OK + assert sorted( + CohortMembership.objects.filter(cohort=amplitude_cohort).values_list( + "identifier", "state" + ) + ) == [ + ("user-1", CohortMembershipState.APPLIED), + ("user-2", CohortMembershipState.APPLIED), + ] + api_key = amplitude_cohort.environment.api_key + document = dynamodb_identity_wrapper.get_item(f"{api_key}_user-1") + assert document is not None + assert document["system_traits"] == {amplitude_cohort.system_trait_key: True} + + +def test_amplitude_add_members__applied_member__stays_applied( + cohort_sync_key: _KeyAndPlaintext, + amplitude_cohort: Cohort, + dynamodb_identity_wrapper: DynamoIdentityWrapper, +) -> None: + # Given + _, plaintext = cohort_sync_key + CohortMembership.objects.create( + cohort=amplitude_cohort, + identifier="user-1", + state=CohortMembershipState.APPLIED, + ) + client = _authenticated_client(plaintext) + url = reverse( + "api-v1:cohort-sync:amplitude-add", kwargs={"pk": str(amplitude_cohort.uuid)} + ) + + # When + response = client.post(url, data={"user_ids": ["user-1"]}, format="json") + + # Then + assert response.status_code == status.HTTP_200_OK + membership = CohortMembership.objects.get(cohort=amplitude_cohort) + assert (membership.identifier, membership.state) == ( + "user-1", + CohortMembershipState.APPLIED, + ) + + +def test_amplitude_remove_members__applied_member__unsets_trait_and_deletes_row( + cohort_sync_key: _KeyAndPlaintext, + amplitude_cohort: Cohort, + dynamodb_identity_wrapper: DynamoIdentityWrapper, +) -> None: + # Given + api_key = amplitude_cohort.environment.api_key + trait_key = amplitude_cohort.system_trait_key + dynamodb_identity_wrapper.put_item( + { + "composite_key": f"{api_key}_member", + "identifier": "member", + "environment_api_key": api_key, + "system_traits": {trait_key: True}, + } + ) + CohortMembership.objects.create( + cohort=amplitude_cohort, + identifier="member", + state=CohortMembershipState.APPLIED, + ) + _, plaintext = cohort_sync_key + client = _authenticated_client(plaintext) + url = reverse( + "api-v1:cohort-sync:amplitude-remove", + kwargs={"pk": str(amplitude_cohort.uuid)}, + ) + + # When + response = client.post(url, data={"user_ids": ["member"]}, format="json") + + # Then + assert response.status_code == status.HTTP_200_OK + assert not CohortMembership.objects.filter(cohort=amplitude_cohort).exists() + document = dynamodb_identity_wrapper.get_item(f"{api_key}_member") + assert document is not None + assert document["system_traits"] == {} + + +def test_amplitude_remove_members__unknown_identifier__no_rows_created( + cohort_sync_key: _KeyAndPlaintext, + amplitude_cohort: Cohort, + dynamodb_identity_wrapper: DynamoIdentityWrapper, +) -> None: + # Given + _, plaintext = cohort_sync_key + client = _authenticated_client(plaintext) + url = reverse( + "api-v1:cohort-sync:amplitude-remove", + kwargs={"pk": str(amplitude_cohort.uuid)}, + ) + + # When + response = client.post(url, data={"user_ids": ["stranger"]}, format="json") + + # Then + assert response.status_code == status.HTTP_200_OK + assert not CohortMembership.objects.filter(cohort=amplitude_cohort).exists() + + +def test_amplitude_add_members__csv_cohort__returns_404( + cohort_sync_key: _KeyAndPlaintext, + edge_cohort: Cohort, +) -> None: + # Given - an edge cohort whose source is CSV, not Amplitude + _, plaintext = cohort_sync_key + client = _authenticated_client(plaintext) + url = reverse( + "api-v1:cohort-sync:amplitude-add", kwargs={"pk": str(edge_cohort.uuid)} + ) + + # When + response = client.post(url, data={"user_ids": ["user-1"]}, format="json") + + # Then + assert response.status_code == status.HTTP_404_NOT_FOUND + + +def test_amplitude_add_members__other_environment_cohort__returns_404( + amplitude_cohort: Cohort, + dynamo_enabled_project_environment_two: Environment, +) -> None: + # Given - a valid key scoped to a different environment + _, plaintext = CohortSyncKey.objects.create_key( + name="other env", environment=dynamo_enabled_project_environment_two + ) + client = _authenticated_client(plaintext) + url = reverse( + "api-v1:cohort-sync:amplitude-add", kwargs={"pk": str(amplitude_cohort.uuid)} + ) + + # When + response = client.post(url, data={"user_ids": ["user-1"]}, format="json") + + # Then + assert response.status_code == status.HTTP_404_NOT_FOUND + + +def test_amplitude_add_members__deletion_requested_cohort__returns_404( + cohort_sync_key: _KeyAndPlaintext, + amplitude_cohort: Cohort, +) -> None: + # Given + amplitude_cohort.deletion_requested_at = timezone.now() + amplitude_cohort.save(update_fields=["deletion_requested_at"]) + _, plaintext = cohort_sync_key + client = _authenticated_client(plaintext) + url = reverse( + "api-v1:cohort-sync:amplitude-add", kwargs={"pk": str(amplitude_cohort.uuid)} + ) + + # When + response = client.post(url, data={"user_ids": ["user-1"]}, format="json") + + # Then + assert response.status_code == status.HTTP_404_NOT_FOUND + + +def test_amplitude_add_members__malformed_list_id__returns_404( + cohort_sync_key: _KeyAndPlaintext, +) -> None: + # Given + _, plaintext = cohort_sync_key + client = _authenticated_client(plaintext) + url = reverse("api-v1:cohort-sync:amplitude-add", kwargs={"pk": "not-a-uuid"}) + + # When + response = client.post(url, data={"user_ids": ["user-1"]}, format="json") + + # Then + assert response.status_code == status.HTTP_404_NOT_FOUND + + +def test_amplitude_add_members__empty_user_ids__returns_400( + cohort_sync_key: _KeyAndPlaintext, + amplitude_cohort: Cohort, +) -> None: + # Given + _, plaintext = cohort_sync_key + client = _authenticated_client(plaintext) + url = reverse( + "api-v1:cohort-sync:amplitude-add", kwargs={"pk": str(amplitude_cohort.uuid)} + ) + + # When + response = client.post(url, data={"user_ids": []}, format="json") + + # Then + assert response.status_code == status.HTTP_400_BAD_REQUEST diff --git a/api/tests/unit/cohorts/test_views.py b/api/tests/unit/cohorts/test_views.py index 488dfe7eae5c..54f9540a7cde 100644 --- a/api/tests/unit/cohorts/test_views.py +++ b/api/tests/unit/cohorts/test_views.py @@ -10,7 +10,7 @@ from rest_framework import status from rest_framework.test import APIClient -from cohorts.models import Cohort +from cohorts.models import Cohort, CohortSyncKey from environments.dynamodb import DynamoIdentityWrapper from environments.models import Environment from organisations.models import Subscription @@ -262,3 +262,96 @@ def test_create_cohort__non_edge_project__returns_400( # Then assert response.status_code == status.HTTP_400_BAD_REQUEST assert response.json()["detail"] == "Dynamo DB is not enabled for this project" + + +def test_create_sync_key__staff_with_permissions__returns_201_with_plaintext_key( + staff_client: APIClient, + environment: Environment, + with_project_permissions: WithProjectPermissionsCallable, + with_environment_permissions: WithEnvironmentPermissionsCallable, +) -> None: + # Given + with_project_permissions([MANAGE_SEGMENTS]) # type: ignore[call-arg] + with_environment_permissions( # type: ignore[call-arg] + [VIEW_ENVIRONMENT, MANAGE_SEGMENT_OVERRIDES] + ) + url = reverse( + "api-v1:environments:cohorts:sync-keys-list", args=[environment.api_key] + ) + + # When + response = staff_client.post(url, data={"name": "Amplitude prod"}, format="json") + + # Then + assert response.status_code == status.HTTP_201_CREATED + key = CohortSyncKey.objects.get(environment=environment) + assert response.json()["prefix"] == key.prefix + assert response.json()["key"].startswith(key.prefix) + assert key.name == "Amplitude prod" + + +def test_create_sync_key__staff_without_permission__returns_403( + staff_client: APIClient, + environment: Environment, +) -> None: + # Given + url = reverse( + "api-v1:environments:cohorts:sync-keys-list", args=[environment.api_key] + ) + + # When + response = staff_client.post(url, data={"name": "Amplitude prod"}, format="json") + + # Then + assert response.status_code == status.HTTP_403_FORBIDDEN + + +def test_list_sync_keys__revoked_key__excluded_and_plaintext_never_returned( + staff_client: APIClient, + environment: Environment, + with_environment_permissions: WithEnvironmentPermissionsCallable, +) -> None: + # Given + with_environment_permissions([VIEW_ENVIRONMENT]) # type: ignore[call-arg] + CohortSyncKey.objects.create_key(name="live", environment=environment) + revoked, _ = CohortSyncKey.objects.create_key( + name="revoked", environment=environment + ) + revoked.revoked = True + revoked.save() + url = reverse( + "api-v1:environments:cohorts:sync-keys-list", args=[environment.api_key] + ) + + # When + response = staff_client.get(url) + + # Then + assert response.status_code == status.HTTP_200_OK + assert [(row["name"], row["key"]) for row in response.json()] == [("live", None)] + + +def test_delete_sync_key__existing_key__revokes( + staff_client: APIClient, + environment: Environment, + with_project_permissions: WithProjectPermissionsCallable, + with_environment_permissions: WithEnvironmentPermissionsCallable, +) -> None: + # Given + with_project_permissions([MANAGE_SEGMENTS]) # type: ignore[call-arg] + with_environment_permissions( # type: ignore[call-arg] + [VIEW_ENVIRONMENT, MANAGE_SEGMENT_OVERRIDES] + ) + key, _ = CohortSyncKey.objects.create_key(name="old", environment=environment) + url = reverse( + "api-v1:environments:cohorts:sync-keys-detail", + args=[environment.api_key, key.prefix], + ) + + # When + response = staff_client.delete(url) + + # Then + assert response.status_code == status.HTTP_204_NO_CONTENT + key.refresh_from_db() + assert key.revoked is True diff --git a/docs/docs/deployment-self-hosting/observability/_events-catalogue.md b/docs/docs/deployment-self-hosting/observability/_events-catalogue.md index 6082c7fcfd1f..d09a0f6eeeaa 100644 --- a/docs/docs/deployment-self-hosting/observability/_events-catalogue.md +++ b/docs/docs/deployment-self-hosting/observability/_events-catalogue.md @@ -74,7 +74,7 @@ Attributes: ### `cohorts.cohort.created` Logged at `info` from: - - `api/cohorts/services.py:103` + - `api/cohorts/services.py:111` Attributes: - `cohort.id` @@ -86,7 +86,7 @@ Attributes: ### `cohorts.cohort.deleted` Logged at `info` from: - - `api/cohorts/services.py:140` + - `api/cohorts/services.py:192` Attributes: - `cohort.id` @@ -95,7 +95,7 @@ Attributes: ### `cohorts.cohort.deletion_requested` Logged at `info` from: - - `api/cohorts/services.py:124` + - `api/cohorts/services.py:176` Attributes: - `cohort.id` @@ -104,7 +104,7 @@ Attributes: ### `cohorts.membership.applied` Logged at `info` from: - - `api/cohorts/services.py:72` + - `api/cohorts/services.py:77` Attributes: - `adds.count` @@ -130,6 +130,18 @@ Logged at `warning` from: Attributes: - `cohort.id` +### `cohorts.membership.deltas_received` + +Logged at `info` from: + - `api/cohorts/services.py:143` + - `api/cohorts/services.py:161` + +Attributes: + - `action` + - `cohort.id` + - `deltas.count` + - `environment.id` + ### `core.encrypted_field.decrypt_failed` Logged at `warning` from: From bd48a947094e0b123cead708ad386009bf6a8148 Mon Sep 17 00:00:00 2001 From: "flagsmith-engineering[bot]" Date: Fri, 14 Aug 2026 09:03:15 +0000 Subject: [PATCH 2/2] chore: Update documentation artefacts --- openapi.yaml | 209 ++++++++++++++++++++++++++++++++++++++++++++++++++- 1 file changed, 208 insertions(+), 1 deletion(-) diff --git a/openapi.yaml b/openapi.yaml index e84747945314..db619230161e 100644 --- a/openapi.yaml +++ b/openapi.yaml @@ -1494,6 +1494,83 @@ paths: - basicAuth: [] tags: - Webhooks + /api/v1/cohort-sync/amplitude/lists/: + post: + operationId: api_v1_cohort_sync_amplitude_lists_create + description: Called by Amplitude once per cohort sync setup; creates the backing cohort and returns its identifier as the list ID. + requestBody: + required: true + content: + application/json: + schema: + $ref: '#/components/schemas/AmplitudeList' + application/x-www-form-urlencoded: + schema: + $ref: '#/components/schemas/AmplitudeList' + multipart/form-data: + schema: + $ref: '#/components/schemas/AmplitudeList' + responses: + '200': + description: '' + content: + application/json: + schema: + $ref: '#/components/schemas/AmplitudeListResponse' + tags: + - Other + '/api/v1/cohort-sync/amplitude/lists/{id}/add/': + post: + operationId: api_v1_cohort_sync_amplitude_lists_add_create + parameters: + - name: id + in: path + required: true + schema: + type: string + requestBody: + required: true + content: + application/json: + schema: + $ref: '#/components/schemas/CohortSyncMembers' + application/x-www-form-urlencoded: + schema: + $ref: '#/components/schemas/CohortSyncMembers' + multipart/form-data: + schema: + $ref: '#/components/schemas/CohortSyncMembers' + responses: + '200': + description: No response body + tags: + - Other + '/api/v1/cohort-sync/amplitude/lists/{id}/remove/': + post: + operationId: api_v1_cohort_sync_amplitude_lists_remove_create + parameters: + - name: id + in: path + required: true + schema: + type: string + requestBody: + required: true + content: + application/json: + schema: + $ref: '#/components/schemas/CohortSyncMembers' + application/x-www-form-urlencoded: + schema: + $ref: '#/components/schemas/CohortSyncMembers' + multipart/form-data: + schema: + $ref: '#/components/schemas/CohortSyncMembers' + responses: + '200': + description: No response body + tags: + - Other /api/v1/environment-document/: get: operationId: sdk_v1_environment_document @@ -2149,6 +2226,87 @@ paths: tags: - Environments x-flagsmith-minimum-plan: START_UP + '/api/v1/environments/{environment_api_key}/cohorts/sync-keys/': + get: + operationId: api_v1_environments_cohorts_sync_keys_list + parameters: + - name: environment_api_key + in: path + required: true + schema: + type: string + responses: + '200': + description: '' + content: + application/json: + schema: + type: array + items: + $ref: '#/components/schemas/CohortSyncKey' + security: + - tokenAuth: [] + - Master API Key: [] + tags: + - Environments + x-flagsmith-minimum-plan: START_UP + post: + operationId: api_v1_environments_cohorts_sync_keys_create + description: Create a cohort sync key. The response is the only time the plaintext key is available. + parameters: + - name: environment_api_key + in: path + required: true + schema: + type: string + requestBody: + content: + application/json: + schema: + $ref: '#/components/schemas/CohortSyncKey' + application/x-www-form-urlencoded: + schema: + $ref: '#/components/schemas/CohortSyncKey' + multipart/form-data: + schema: + $ref: '#/components/schemas/CohortSyncKey' + responses: + '201': + description: '' + content: + application/json: + schema: + $ref: '#/components/schemas/CohortSyncKey' + security: + - tokenAuth: [] + - Master API Key: [] + tags: + - Environments + x-flagsmith-minimum-plan: START_UP + '/api/v1/environments/{environment_api_key}/cohorts/sync-keys/{prefix}/': + delete: + operationId: api_v1_environments_cohorts_sync_keys_destroy + description: Revoke a cohort sync key. + parameters: + - name: environment_api_key + in: path + required: true + schema: + type: string + - name: prefix + in: path + required: true + schema: + type: string + responses: + '204': + description: No response body + security: + - tokenAuth: [] + - Master API Key: [] + tags: + - Environments + x-flagsmith-minimum-plan: START_UP '/api/v1/environments/{environment_api_key}/create-change-request/': post: operationId: create_environment_feature_change_request @@ -18209,6 +18367,22 @@ components: maxLength: 200 required: - api_key + AmplitudeList: + type: object + properties: + name: + type: string + maxLength: 2000 + required: + - name + AmplitudeListResponse: + type: object + properties: + list_id: + type: string + format: uuid + required: + - list_id AuditLogList: type: object properties: @@ -18831,6 +19005,36 @@ components: readOnly: true required: - name + CohortSyncKey: + type: object + properties: + prefix: + type: string + readOnly: true + name: + description: A free-form name for the API key. Need not be unique. 50 characters max. + type: string + maxLength: 50 + created: + type: string + format: date-time + readOnly: true + key: + type: + - string + - 'null' + readOnly: true + CohortSyncMembers: + type: object + properties: + user_ids: + type: array + items: + type: string + maxLength: 2000 + minItems: 1 + required: + - user_ids Condition: type: object properties: @@ -26701,10 +26905,13 @@ components: required: - channel_id SourceTypeEnum: - description: '* `csv` - CSV' + description: |- + * `csv` - CSV + * `amplitude` - Amplitude type: string enum: - csv + - amplitude StageAction: type: object properties: