diff --git a/api/api/openapi.py b/api/api/openapi.py index d01a2ffdff54..2784dd1270b4 100644 --- a/api/api/openapi.py +++ b/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]] = [ { 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..175a29b2f005 --- /dev/null +++ b/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" 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..1d75794d8d5b --- /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": "cohort sync key", + "verbose_name_plural": "cohort sync keys", + "ordering": ("-created",), + "abstract": False, + }, + ), + ] diff --git a/api/cohorts/models.py b/api/cohorts/models.py index 605b67be0526..b15927b6fcd0 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 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,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" diff --git a/api/cohorts/permissions.py b/api/cohorts/permissions.py index a33d801397ba..c85d43b000da 100644 --- a/api/cohorts/permissions.py +++ b/api/cohorts/permissions.py @@ -7,6 +7,7 @@ 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 @@ -14,6 +15,12 @@ _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) diff --git a/api/cohorts/serializers.py b/api/cohorts/serializers.py index 780216bb9f4a..a5e5bc39915d 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,39 @@ 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) + + +class CohortSyncMembersSerializer(serializers.Serializer[None]): + # Child length mirrors CohortMembership.identifier. + # TODO: this counts characters, but identity data is stored with a + # 1024-byte identifier limit, so a multibyte identifier is accepted here + # and only fails once we try to write it. Check the byte length, together + # with the same check for CSV uploads. + 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 4f4a6c568ac1..0b0688433f83 100644 --- a/api/cohorts/services.py +++ b/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,81 @@ 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.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 + + unique_identifiers = set(identifiers) + with transaction.atomic(): + # Removing a non-member is a no-op: only existing rows flip. + matched = CohortMembership.objects.filter( + cohort=cohort, identifier__in=unique_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=len(unique_identifiers), + members__matched=matched, + ) + + 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..6afc517a0b6c --- /dev/null +++ b/api/cohorts/sync_urls.py @@ -0,0 +1,11 @@ +from rest_framework.routers import SimpleRouter + +from cohorts.sync_views import AmplitudeCohortSyncViewSet + +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 diff --git a/api/cohorts/sync_views.py b/api/cohorts/sync_views.py new file mode 100644 index 000000000000..eae69581f2ab --- /dev/null +++ b/api/cohorts/sync_views.py @@ -0,0 +1,85 @@ +import typing +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.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.permissions import HasCohortSyncKey +from cohorts.serializers import ( + AmplitudeListSerializer, + CohortSyncMembersSerializer, +) + +_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 = [HasCohortSyncKey] + + def create(self, request: Request) -> Response: + serializer = AmplitudeListSerializer(data=request.data) + serializer.is_valid(raise_exception=True) + environment = self._get_key(request).environment + cohort = services.create_cohort_for_source( + 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: + # HasCohortSyncKey has already established the type. + return typing.cast(CohortSyncKey, 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 864b80af24a2..532cc4e59304 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 @@ -54,3 +56,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..f8ddd2a0de68 --- /dev/null +++ b/api/tests/unit/cohorts/test_sync_views.py @@ -0,0 +1,402 @@ +import typing + +from django.urls import reverse +from django.utils import timezone +from flag_engine.segments.constants import IS_SET +from pytest_django.fixtures import SettingsWrapper +from rest_framework import status +from rest_framework.test import APIClient + +from audit.models import AuditLog +from audit.related_object_type import RelatedObjectType +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__postgres_environment__creates_cohort( + environment: Environment, +) -> None: + # Given - an environment whose identities live in Postgres + _, plaintext = CohortSyncKey.objects.create_key( + name="postgres 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_200_OK + cohort = Cohort.objects.get(uuid=response.json()["list_id"]) + assert cohort.environment == environment + + +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 + + +def test_amplitude_create_list__valid_key__audits_and_queues_environment_update( + 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": "Beta users"}, format="json") + + # Then - the audit record carries no user, names the source, and is the + # hook that rebuilds the environment document. + assert response.status_code == status.HTTP_200_OK + cohort = Cohort.objects.get(uuid=response.json()["list_id"]) + audit_log = AuditLog.objects.get(related_object_id=cohort.segment_id) + assert audit_log.author is None + assert audit_log.master_api_key is None + assert audit_log.environment == key.environment + assert audit_log.related_object_type == RelatedObjectType.SEGMENT.name + assert audit_log.log == ( + "New Segment created: Beta users (via Amplitude cohort sync)" + ) + assert audit_log.environment_document_updated is True + + +def test_amplitude_create_list__valid_key__history_records_no_user( + cohort_sync_key: _KeyAndPlaintext, + dynamodb_identity_wrapper: DynamoIdentityWrapper, +) -> None: + # Given + _, plaintext = cohort_sync_key + client = _authenticated_client(plaintext) + url = reverse("api-v1:cohort-sync:amplitude-list") + + # When + response = client.post(url, data={"name": "Beta users"}, format="json") + + # Then - a machine caller leaves no user on historical records; stamping + # one would fail, since the sync key is not a Flagsmith user. + assert response.status_code == status.HTTP_200_OK + cohort = Cohort.objects.get(uuid=response.json()["list_id"]) + history_record = cohort.segment.history.get() + assert history_record.history_user is None + assert history_record.master_api_key is None + + +def test_amplitude_add_members__master_api_key_throttle_enabled__succeeds( + cohort_sync_key: _KeyAndPlaintext, + amplitude_cohort: Cohort, + dynamodb_identity_wrapper: DynamoIdentityWrapper, + settings: SettingsWrapper, +) -> None: + # Given - the throttle that reads master API key attributes off the caller + settings.REST_FRAMEWORK = { + **settings.REST_FRAMEWORK, + "DEFAULT_THROTTLE_CLASSES": ["core.throttling.MasterAPIKeyUserRateThrottle"], + "DEFAULT_THROTTLE_RATES": {"master_api_key": "1000/minute"}, + } + _, 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_200_OK diff --git a/api/tests/unit/cohorts/test_views.py b/api/tests/unit/cohorts/test_views.py index 64b218a71f5f..b128a776e7d4 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,131 @@ def test_create_cohort__non_edge_project__returns_201( # Then assert response.status_code == status.HTTP_201_CREATED assert Cohort.objects.get(id=response.json()["id"]).environment == environment + + +def test_create_sync_key__staff_with_permissions__returns_201_with_plaintext_key( + staff_client: APIClient, + dynamo_enabled_project: Project, + dynamo_enabled_project_environment_one: Environment, + dynamodb_identity_wrapper: DynamoIdentityWrapper, + with_project_permissions: WithProjectPermissionsCallable, + with_environment_permissions: WithEnvironmentPermissionsCallable, +) -> None: + # Given + environment = dynamo_enabled_project_environment_one + with_project_permissions( # type: ignore[call-arg] + [MANAGE_SEGMENTS], project_id=dynamo_enabled_project.id + ) + with_environment_permissions( # type: ignore[call-arg] + [VIEW_ENVIRONMENT, MANAGE_SEGMENT_OVERRIDES], + environment_id=environment.id, + ) + 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 + + +def test_create_sync_key__missing_name__returns_400( + staff_client: APIClient, + dynamo_enabled_project: Project, + dynamo_enabled_project_environment_one: Environment, + dynamodb_identity_wrapper: DynamoIdentityWrapper, + with_project_permissions: WithProjectPermissionsCallable, + with_environment_permissions: WithEnvironmentPermissionsCallable, +) -> None: + # Given + with_project_permissions( # type: ignore[call-arg] + [MANAGE_SEGMENTS], project_id=dynamo_enabled_project.id + ) + with_environment_permissions( # type: ignore[call-arg] + [VIEW_ENVIRONMENT, MANAGE_SEGMENT_OVERRIDES], + environment_id=dynamo_enabled_project_environment_one.id, + ) + url = reverse( + "api-v1:environments:cohorts:sync-keys-list", + args=[dynamo_enabled_project_environment_one.api_key], + ) + + # When + response = staff_client.post(url, data={}, format="json") + + # Then + assert response.status_code == status.HTTP_400_BAD_REQUEST + assert not CohortSyncKey.objects.exists() diff --git a/docs/docs/deployment-self-hosting/observability/_events-catalogue.md b/docs/docs/deployment-self-hosting/observability/_events-catalogue.md index d28308c413ad..742e613bf0b3 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:100` + - `api/cohorts/services.py:111` Attributes: - `cohort.id` @@ -86,7 +86,7 @@ Attributes: ### `cohorts.cohort.deleted` Logged at `info` from: - - `api/cohorts/services.py:133` + - `api/cohorts/services.py:219` Attributes: - `cohort.id` @@ -95,7 +95,7 @@ Attributes: ### `cohorts.cohort.deletion_requested` Logged at `info` from: - - `api/cohorts/services.py:117` + - `api/cohorts/services.py:203` Attributes: - `cohort.id` @@ -104,7 +104,7 @@ Attributes: ### `cohorts.membership.applied` Logged at `info` from: - - `api/cohorts/services.py:69` + - `api/cohorts/services.py:77` Attributes: - `adds.count` @@ -129,6 +129,19 @@ Logged at `warning` from: Attributes: - `cohort.id` +### `cohorts.membership.deltas_received` + +Logged at `info` from: + - `api/cohorts/services.py:168` + - `api/cohorts/services.py:187` + +Attributes: + - `action` + - `cohort.id` + - `deltas.count` + - `environment.id` + - `members.matched` + ### `core.encrypted_field.decrypt_failed` Logged at `warning` from: diff --git a/mcp/src/flagsmith_mcp/openapi.json b/mcp/src/flagsmith_mcp/openapi.json index cdb4319246ba..ef2dbc10e145 100644 --- a/mcp/src/flagsmith_mcp/openapi.json +++ b/mcp/src/flagsmith_mcp/openapi.json @@ -7756,6 +7756,11 @@ } }, "securitySchemes": { + "Cohort Sync Key": { + "type": "http", + "scheme": "bearer", + "description": "For cohort sync endpoints called by an external cohort source, such as Amplitude." + }, "Environment API Key": { "type": "apiKey", "in": "header", diff --git a/openapi.yaml b/openapi.yaml index 5eaa7d92ec94..8b03b8f842d7 100644 --- a/openapi.yaml +++ b/openapi.yaml @@ -1658,6 +1658,89 @@ 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' + security: + - Cohort Sync Key: [] + 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 + security: + - Cohort Sync Key: [] + 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 + security: + - Cohort Sync Key: [] + tags: + - Other /api/v1/environment-document/: get: operationId: sdk_v1_environment_document @@ -2313,6 +2396,88 @@ 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: + required: true + 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 @@ -18550,6 +18715,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: @@ -19187,6 +19368,37 @@ components: readOnly: true required: - name + CohortSyncKey: + type: object + properties: + prefix: + type: string + readOnly: true + name: + type: string + maxLength: 50 + created: + type: string + format: date-time + readOnly: true + key: + type: + - string + - 'null' + readOnly: true + required: + - name + CohortSyncMembers: + type: object + properties: + user_ids: + type: array + items: + type: string + maxLength: 2000 + minItems: 1 + required: + - user_ids Condition: type: object properties: @@ -27131,10 +27343,13 @@ components: required: - channel_id SourceTypeEnum: - description: '* `csv` - CSV' + description: |- + * `csv` - CSV + * `amplitude` - Amplitude type: string enum: - csv + - amplitude StageAction: type: object properties: @@ -29143,6 +29358,10 @@ components: required: - type securitySchemes: + Cohort Sync Key: + type: http + scheme: bearer + description: 'For cohort sync endpoints called by an external cohort source, such as Amplitude.' Environment API Key: type: apiKey in: header diff --git a/sdk/openapi.yaml b/sdk/openapi.yaml index 0c67801d259e..e534853228a9 100644 --- a/sdk/openapi.yaml +++ b/sdk/openapi.yaml @@ -545,6 +545,10 @@ components: - feature_state_value title: V1Flag securitySchemes: + Cohort Sync Key: + type: http + scheme: bearer + description: 'For cohort sync endpoints called by an external cohort source, such as Amplitude.' Environment API Key: type: apiKey in: header