feat(cohorts): generic cohort sync webhook

Author: gagantrivedi

api/cohorts/serializers.py

@@ -156,3 +156,11 @@ class CohortCsvSyncResultSerializer(serializers.Serializer):  # type: ignore[typ
     removed = serializers.IntegerField(min_value=0)
     unchanged = serializers.IntegerField(min_value=0)
     ignored = CohortCsvSyncIgnoredRowsSerializer()
+
+
+class WebhookSyncMembersSerializer(serializers.Serializer[None]):
+    identifiers = serializers.ListField(
+        child=serializers.CharField(validators=[_validate_identifier_byte_length]),
+        min_length=1,
+        max_length=10000,
+    )

api/cohorts/services.py

@@ -1,6 +1,7 @@
 import csv
 import io
 import typing
+from uuid import UUID
 
 import structlog
 from django.db import transaction
@@ -192,6 +193,25 @@ def create_cohort_for_source(
     return cohort
 
 
+def get_active_cohort(
+    *,
+    environment: "Environment",
+    source_type: CohortSourceType,
+    uuid: str,
+) -> Cohort | None:
+    try:
+        cohort_uuid = UUID(uuid)
+    except ValueError:
+        return None
+    cohort: Cohort | None = Cohort.objects.filter(
+        uuid=cohort_uuid,
+        environment=environment,
+        source_type=source_type,
+        deletion_requested_at__isnull=True,
+    ).first()
+    return cohort
+
+
 def get_cohort_for_source(
     *,
     environment: "Environment",

api/cohorts/sync_urls.py

@@ -1,13 +1,18 @@
 from django.urls import path
 from rest_framework.routers import SimpleRouter
 
-from cohorts.sync_views import AmplitudeCohortSyncViewSet, MixpanelCohortSyncView
+from cohorts.sync_views import (
+    AmplitudeCohortSyncViewSet,
+    MixpanelCohortSyncView,
+    WebhookCohortSyncViewSet,
+)
 
 app_name = "cohort-sync"
 
 # SimpleRouter: nothing here is browsed by a person.
 router = SimpleRouter()
 router.register(r"amplitude/lists", AmplitudeCohortSyncViewSet, basename="amplitude")
+router.register(r"webhook/cohorts", WebhookCohortSyncViewSet, basename="webhook")
 
 urlpatterns = [
     path("mixpanel/webhook/", MixpanelCohortSyncView.as_view(), name="mixpanel"),

api/cohorts/sync_views.py

@@ -1,6 +1,5 @@
 import json
 import typing
-import uuid as uuid_module
 
 import structlog
 from drf_spectacular.utils import extend_schema, extend_schema_view, inline_serializer
@@ -19,6 +18,7 @@
     AmplitudeListSerializer,
     CohortSyncMembersSerializer,
     MixpanelWebhookSerializer,
+    WebhookSyncMembersSerializer,
 )
 
 _LIST_RESPONSE = inline_serializer(
@@ -96,16 +96,11 @@ def _get_key(self, request: Request) -> CohortSyncKey:
         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,
+        cohort = services.get_active_cohort(
             environment=self._get_key(request).environment,
             source_type=CohortSourceType.AMPLITUDE,
-            deletion_requested_at__isnull=True,
-        ).first()
+            uuid=pk,
+        )
         if cohort is None:
             raise NotFound("List not found.")
         return cohort
@@ -241,3 +236,55 @@ def _echo_action(self, request: Request) -> str | None:
         ):
             return action_value
         return None
+
+
+@extend_schema_view(
+    add=extend_schema(
+        description=(
+            "Add members to a CSV cohort. Accepted deltas are applied to "
+            "identity data asynchronously. Re-adding a member is a no-op, "
+            "so retries are safe."
+        ),
+        request=WebhookSyncMembersSerializer,
+        responses={200: None},
+    ),
+    remove=extend_schema(
+        description=(
+            "Remove members from a CSV cohort. Accepted deltas are applied "
+            "to identity data asynchronously. Removing a non-member is a "
+            "no-op, so retries are safe."
+        ),
+        request=WebhookSyncMembersSerializer,
+        responses={200: None},
+    ),
+)
+class WebhookCohortSyncViewSet(viewsets.ViewSet):
+    authentication_classes = [CohortSyncKeyAuthentication]
+    permission_classes = [HasCohortSyncKey, CohortSyncPlanPermission]
+
+    @action(detail=True, methods=["POST"], url_path="members/add")
+    def add(self, request: Request, pk: str) -> Response:
+        cohort = self._get_cohort(request, pk)
+        serializer = WebhookSyncMembersSerializer(data=request.data)
+        serializer.is_valid(raise_exception=True)
+        services.add_cohort_members(cohort, serializer.validated_data["identifiers"])
+        return Response()
+
+    @action(detail=True, methods=["POST"], url_path="members/remove")
+    def remove(self, request: Request, pk: str) -> Response:
+        cohort = self._get_cohort(request, pk)
+        serializer = WebhookSyncMembersSerializer(data=request.data)
+        serializer.is_valid(raise_exception=True)
+        services.remove_cohort_members(cohort, serializer.validated_data["identifiers"])
+        return Response()
+
+    def _get_cohort(self, request: Request, pk: str) -> Cohort:
+        cohort = services.get_active_cohort(
+            # HasCohortSyncKey has already established the type.
+            environment=typing.cast(CohortSyncKey, request.auth).environment,
+            source_type=CohortSourceType.CSV,
+            uuid=pk,
+        )
+        if cohort is None:
+            raise NotFound("Cohort not found.")
+        return cohort

api/tests/unit/cohorts/test_sync_views.py

@@ -1041,3 +1041,227 @@ def test_mixpanel_webhook__saas_free_plan__returns_403(
 
     # Then
     assert response.status_code == status.HTTP_403_FORBIDDEN
+
+
+def test_webhook_add_members__csv_cohort__applies_memberships(
+    cohort: Cohort,
+) -> None:
+    # Given
+    _, plaintext = CohortSyncKey.objects.create_key(
+        name="test key", environment=cohort.environment
+    )
+    client = _authenticated_client(plaintext)
+    url = reverse("api-v1:cohort-sync:webhook-add", kwargs={"pk": str(cohort.uuid)})
+
+    # When
+    response = client.post(
+        url, data={"identifiers": ["user-1", "user-2"]}, format="json"
+    )
+
+    # Then
+    assert response.status_code == status.HTTP_200_OK
+    assert sorted(
+        CohortMembership.objects.filter(cohort=cohort).values_list(
+            "identifier", "state"
+        )
+    ) == [
+        ("user-1", CohortMembershipState.APPLIED),
+        ("user-2", CohortMembershipState.APPLIED),
+    ]
+    identity = Identity.objects.get(environment=cohort.environment, identifier="user-1")
+    assert identity.system_traits == {cohort.system_trait_key: True}
+
+
+def test_webhook_remove_members__applied_member__removes_membership_and_trait(
+    cohort: Cohort,
+) -> None:
+    # Given
+    Identity.objects.create(
+        environment=cohort.environment,
+        identifier="member",
+        system_traits={cohort.system_trait_key: True},
+    )
+    CohortMembership.objects.create(
+        cohort=cohort, identifier="member", state=CohortMembershipState.APPLIED
+    )
+    _, plaintext = CohortSyncKey.objects.create_key(
+        name="test key", environment=cohort.environment
+    )
+    client = _authenticated_client(plaintext)
+    url = reverse("api-v1:cohort-sync:webhook-remove", kwargs={"pk": str(cohort.uuid)})
+
+    # When
+    response = client.post(url, data={"identifiers": ["member"]}, format="json")
+
+    # Then
+    assert response.status_code == status.HTTP_200_OK
+    assert not CohortMembership.objects.filter(cohort=cohort).exists()
+    identity = Identity.objects.get(environment=cohort.environment, identifier="member")
+    assert identity.system_traits == {}
+
+
+def test_webhook_add_members__non_csv_cohort__returns_404(
+    cohort_sync_key: _KeyAndPlaintext,
+    amplitude_cohort: Cohort,
+) -> None:
+    # Given
+    _, plaintext = cohort_sync_key
+    client = _authenticated_client(plaintext)
+    url = reverse(
+        "api-v1:cohort-sync:webhook-add", kwargs={"pk": str(amplitude_cohort.uuid)}
+    )
+
+    # When
+    response = client.post(url, data={"identifiers": ["user-1"]}, format="json")
+
+    # Then
+    assert response.status_code == status.HTTP_404_NOT_FOUND
+
+
+def test_webhook_add_members__other_environment_key__returns_404(
+    cohort: Cohort,
+) -> None:
+    # Given - a key scoped to a different environment than the cohort's
+    other_environment = Environment.objects.create(
+        name="Other environment", project=cohort.environment.project
+    )
+    _, plaintext = CohortSyncKey.objects.create_key(
+        name="other key", environment=other_environment
+    )
+    client = _authenticated_client(plaintext)
+    url = reverse("api-v1:cohort-sync:webhook-add", kwargs={"pk": str(cohort.uuid)})
+
+    # When
+    response = client.post(url, data={"identifiers": ["user-1"]}, format="json")
+
+    # Then
+    assert response.status_code == status.HTTP_404_NOT_FOUND
+    assert not CohortMembership.objects.exists()
+
+
+def test_webhook_add_members__identifier_over_1024_bytes__returns_400(
+    cohort: Cohort,
+) -> None:
+    # Given - 512 three-byte characters: few characters, too many bytes
+    multibyte_identifier = "€" * 512
+    _, plaintext = CohortSyncKey.objects.create_key(
+        name="test key", environment=cohort.environment
+    )
+    client = _authenticated_client(plaintext)
+    url = reverse("api-v1:cohort-sync:webhook-add", kwargs={"pk": str(cohort.uuid)})
+
+    # When
+    response = client.post(
+        url, data={"identifiers": [multibyte_identifier]}, format="json"
+    )
+
+    # Then
+    assert response.status_code == status.HTTP_400_BAD_REQUEST
+    assert "1024 bytes" in str(response.json())
+    assert not CohortMembership.objects.exists()
+
+
+def test_webhook_add_members__over_10000_identifiers__returns_400(
+    cohort: Cohort,
+) -> None:
+    # Given
+    _, plaintext = CohortSyncKey.objects.create_key(
+        name="test key", environment=cohort.environment
+    )
+    client = _authenticated_client(plaintext)
+    url = reverse("api-v1:cohort-sync:webhook-add", kwargs={"pk": str(cohort.uuid)})
+
+    # When
+    response = client.post(
+        url,
+        data={"identifiers": [f"user-{i}" for i in range(10001)]},
+        format="json",
+    )
+
+    # Then
+    assert response.status_code == status.HTTP_400_BAD_REQUEST
+    assert not CohortMembership.objects.exists()
+
+
+def test_webhook_add_members__empty_identifiers__returns_400(
+    cohort: Cohort,
+) -> None:
+    # Given
+    _, plaintext = CohortSyncKey.objects.create_key(
+        name="test key", environment=cohort.environment
+    )
+    client = _authenticated_client(plaintext)
+    url = reverse("api-v1:cohort-sync:webhook-add", kwargs={"pk": str(cohort.uuid)})
+
+    # When
+    response = client.post(url, data={"identifiers": []}, format="json")
+
+    # Then
+    assert response.status_code == status.HTTP_400_BAD_REQUEST
+
+
+def test_webhook_add_members__malformed_uuid__returns_404(
+    cohort: Cohort,
+) -> None:
+    # Given
+    _, plaintext = CohortSyncKey.objects.create_key(
+        name="test key", environment=cohort.environment
+    )
+    client = _authenticated_client(plaintext)
+    url = reverse("api-v1:cohort-sync:webhook-add", kwargs={"pk": "not-a-uuid"})
+
+    # When
+    response = client.post(url, data={"identifiers": ["user-1"]}, format="json")
+
+    # Then
+    assert response.status_code == status.HTTP_404_NOT_FOUND
+
+
+def test_webhook_add_members__deletion_requested_cohort__returns_404(
+    cohort: Cohort,
+) -> None:
+    # Given
+    cohort.deletion_requested_at = timezone.now()
+    cohort.save()
+    _, plaintext = CohortSyncKey.objects.create_key(
+        name="test key", environment=cohort.environment
+    )
+    client = _authenticated_client(plaintext)
+    url = reverse("api-v1:cohort-sync:webhook-add", kwargs={"pk": str(cohort.uuid)})
+
+    # When
+    response = client.post(url, data={"identifiers": ["user-1"]}, format="json")
+
+    # Then
+    assert response.status_code == status.HTTP_404_NOT_FOUND
+
+
+def test_webhook_add_members__missing_credentials__returns_401(
+    cohort: Cohort,
+) -> None:
+    # Given
+    url = reverse("api-v1:cohort-sync:webhook-add", kwargs={"pk": str(cohort.uuid)})
+
+    # When
+    response = APIClient().post(url, data={"identifiers": ["user-1"]}, format="json")
+
+    # Then
+    assert response.status_code == status.HTTP_401_UNAUTHORIZED
+
+
+@pytest.mark.saas_mode
+def test_webhook_add_members__saas_free_plan__returns_403(
+    cohort: Cohort,
+) -> None:
+    # Given
+    _, plaintext = CohortSyncKey.objects.create_key(
+        name="test key", environment=cohort.environment
+    )
+    client = _authenticated_client(plaintext)
+    url = reverse("api-v1:cohort-sync:webhook-add", kwargs={"pk": str(cohort.uuid)})
+
+    # When
+    response = client.post(url, data={"identifiers": ["user-1"]}, format="json")
+
+    # Then
+    assert response.status_code == status.HTTP_403_FORBIDDEN