feat: cohort CSV sync endpoint and cohort summary on segments

api/cohorts/constants.py

@@ -1,6 +1,10 @@
 COHORT_SYSTEM_TRAIT_KEY_PREFIX = "flagsmith_cohort_"
+# Edge identifiers are DynamoDB sort keys, capped at 1024 bytes.
+COHORT_IDENTIFIER_MAX_BYTES = 1024
 COHORT_MEMBERSHIP_APPLY_BATCH_SIZE = 100
 COHORT_MEMBERSHIP_APPLY_MAX_BATCHES_PER_RUN = 10
+COHORT_CSV_MAX_FILE_SIZE_BYTES = 10 * 1024 * 1024
+COHORT_CSV_MEMBERSHIP_CREATE_BATCH_SIZE = 1000
 DYNAMODB_THROTTLING_ERROR_CODES = frozenset(
     {
         "ProvisionedThroughputExceededException",

api/cohorts/dataclasses.py

@@ -0,0 +1,25 @@
+from dataclasses import dataclass
+
+
+@dataclass
+class CsvIdentifierExtraction:
+    identifiers: list[str]
+    empty_count: int
+    duplicate_count: int
+    too_long_count: int
+
+
+@dataclass
+class CohortCsvIgnoredRows:
+    empty: int
+    duplicates: int
+    too_long: int
+
+
+@dataclass
+class CohortCsvSyncResult:
+    version: int
+    added: int
+    removed: int
+    unchanged: int
+    ignored: CohortCsvIgnoredRows

api/cohorts/exceptions.py

@@ -0,0 +1,13 @@
+from rest_framework import status
+from rest_framework.exceptions import APIException
+
+from cohorts.constants import COHORT_CSV_MAX_FILE_SIZE_BYTES
+
+
+class CsvFileTooLargeError(APIException):
+    status_code = status.HTTP_413_REQUEST_ENTITY_TOO_LARGE
+    default_detail = (
+        "CSV file exceeds the "
+        f"{COHORT_CSV_MAX_FILE_SIZE_BYTES // (1024 * 1024)}MB size limit."
+    )
+    default_code = "csv_file_too_large"

api/cohorts/metrics.py

@@ -7,3 +7,15 @@
     "The `operation` label is either `add` or `remove`.",
     ["operation"],
 )
+
+flagsmith_cohorts_csv_syncs_total = prometheus_client.Counter(
+    "flagsmith_cohorts_csv_syncs_total",
+    "Total number of accepted cohort CSV synchronisations, i.e. uploads that "
+    "yielded at least one valid identifier and enqueued a membership sync.",
+)
+
+flagsmith_cohorts_csv_sync_identifiers = prometheus_client.Histogram(
+    "flagsmith_cohorts_csv_sync_identifiers",
+    "Number of unique identifiers extracted per accepted cohort CSV synchronisation.",
+    buckets=(10, 100, 1_000, 10_000, 100_000, 1_000_000),
+)

api/cohorts/serializers.py

@@ -1,16 +1,30 @@
 import typing
 
+from django.core.files.uploadedfile import UploadedFile
 from rest_framework import serializers
 
+from cohorts.constants import COHORT_CSV_MAX_FILE_SIZE_BYTES
+from cohorts.exceptions import CsvFileTooLargeError
 from cohorts.models import Cohort
 from cohorts.services import create_cohort
+from environments.models import Environment
+from metadata.serializers import MetadataSerializer, MetadataSerializerMixin
+from segments.models import Segment
+
+
+class _SegmentMetadataHandler(MetadataSerializerMixin):
+    # The mixin derives the metadata content type from Meta.model; cohort
+    # metadata lives on the managed segment, not on the cohort itself.
+    class Meta:
+        model = Segment
 
 
 class CohortSerializer(serializers.ModelSerializer[Cohort]):
     name = serializers.CharField(max_length=2000, source="segment.name")
     description = serializers.CharField(
         source="segment.description", required=False, allow_null=True
     )
+    metadata = MetadataSerializer(required=False, many=True, write_only=True)
 
     class Meta:
         model = Cohort
@@ -19,17 +33,59 @@ class Meta:
             "uuid",
             "name",
             "description",
+            "metadata",
             "segment",
             "source_type",
             "version",
             "created_at",
         )
         read_only_fields = ("segment", "source_type", "version", "created_at")
 
+    def validate(self, attrs: dict[str, typing.Any]) -> dict[str, typing.Any]:
+        attrs = super().validate(attrs)
+        environment = Environment.objects.get(
+            api_key=self.context["view"].kwargs["environment_api_key"]
+        )
+        project = environment.project
+        _SegmentMetadataHandler()._validate_required_metadata(
+            project.organisation, attrs.get("metadata", []), project
+        )
+        return attrs
+
     def create(self, validated_data: dict[str, typing.Any]) -> Cohort:
         segment_data = validated_data["segment"]
-        return create_cohort(
+        metadata_data = validated_data.pop("metadata", [])
+        cohort = create_cohort(
             environment=validated_data["environment"],
             name=segment_data["name"],
             description=segment_data.get("description"),
         )
+        if metadata_data:
+            _SegmentMetadataHandler()._update_metadata(cohort.segment, metadata_data)
+        return cohort
+
+
+class CohortCsvSyncSerializer(serializers.Serializer):  # type: ignore[type-arg]
+    file = serializers.FileField()
+    identifier_column = serializers.IntegerField(required=False, default=0, min_value=0)
+    has_header = serializers.BooleanField(required=False, default=True)
+
+    def validate_file(self, file: UploadedFile) -> UploadedFile:
+        if file.size and file.size > COHORT_CSV_MAX_FILE_SIZE_BYTES:
+            # Deliberately not a ValidationError: propagates as a 413.
+            raise CsvFileTooLargeError()
+        return file
+
+
+class CohortCsvSyncIgnoredRowsSerializer(serializers.Serializer):  # type: ignore[type-arg]
+    empty = serializers.IntegerField(min_value=0)
+    duplicates = serializers.IntegerField(min_value=0)
+    too_long = serializers.IntegerField(min_value=0)
+
+
+class CohortCsvSyncResultSerializer(serializers.Serializer):  # type: ignore[type-arg]
+    version = serializers.IntegerField(min_value=0)
+    added = serializers.IntegerField(min_value=0)
+    removed = serializers.IntegerField(min_value=0)
+    unchanged = serializers.IntegerField(min_value=0)
+    ignored = CohortCsvSyncIgnoredRowsSerializer()

api/cohorts/services.py

@@ -1,13 +1,29 @@
+import csv
+import io
 import typing
 
 import structlog
 from django.db import transaction
 from django.db.models import QuerySet
 from django.utils import timezone
 from flag_engine.segments.constants import IS_SET
+from rest_framework.exceptions import ValidationError
 
-from cohorts.constants import COHORT_MEMBERSHIP_APPLY_BATCH_SIZE
-from cohorts.metrics import flagsmith_cohorts_membership_deltas_applied_total
+from cohorts.constants import (
+    COHORT_CSV_MEMBERSHIP_CREATE_BATCH_SIZE,
+    COHORT_IDENTIFIER_MAX_BYTES,
+    COHORT_MEMBERSHIP_APPLY_BATCH_SIZE,
+)
+from cohorts.dataclasses import (
+    CohortCsvIgnoredRows,
+    CohortCsvSyncResult,
+    CsvIdentifierExtraction,
+)
+from cohorts.metrics import (
+    flagsmith_cohorts_csv_sync_identifiers,
+    flagsmith_cohorts_csv_syncs_total,
+    flagsmith_cohorts_membership_deltas_applied_total,
+)
 from cohorts.models import Cohort, CohortMembership, CohortMembershipState
 from core.dataclasses import AuthorData
 from environments.dynamodb import DynamoIdentityWrapper
@@ -111,6 +127,130 @@ def create_cohort(
     return cohort
 
 
+def extract_identifiers_from_csv(
+    file: typing.IO[bytes],
+    *,
+    identifier_column: int = 0,
+    has_header: bool = True,
+) -> CsvIdentifierExtraction:
+    # The upload size cap keeps a full read cheap.
+    text = io.StringIO(file.read().decode("utf-8-sig", errors="replace"), newline="")
+    reader = csv.reader(text)
+    seen: set[str] = set()
+    identifiers: list[str] = []
+    empty_count = duplicate_count = too_long_count = 0
+    try:
+        for row_number, row in enumerate(reader):
+            if has_header and row_number == 0:
+                continue
+            if not row:
+                continue
+            value = (
+                row[identifier_column].strip() if identifier_column < len(row) else ""
+            )
+            if not value:
+                empty_count += 1
+            elif len(value.encode()) > COHORT_IDENTIFIER_MAX_BYTES:
+                too_long_count += 1
+            elif value in seen:
+                duplicate_count += 1
+            else:
+                seen.add(value)
+                identifiers.append(value)
+    except csv.Error as exc:
+        raise ValidationError({"file": "Could not parse the CSV file."}) from exc
+    return CsvIdentifierExtraction(
+        identifiers=identifiers,
+        empty_count=empty_count,
+        duplicate_count=duplicate_count,
+        too_long_count=too_long_count,
+    )
+
+
+def sync_cohort_memberships_from_csv(
+    *,
+    cohort: Cohort,
+    file: typing.IO[bytes],
+    identifier_column: int = 0,
+    has_header: bool = True,
+) -> CohortCsvSyncResult:
+    from cohorts.tasks import apply_cohort_membership_deltas
+
+    extraction = extract_identifiers_from_csv(
+        file, identifier_column=identifier_column, has_header=has_header
+    )
+    if not extraction.identifiers:
+        raise ValidationError({"file": "No valid identifiers found in the CSV file."})
+
+    incoming = set(extraction.identifiers)
+    added = removed = unchanged = 0
+    with transaction.atomic():
+        # Serialise concurrent syncs of the same cohort.
+        locked_cohort = Cohort.objects.select_for_update().get(id=cohort.id)
+        existing = {
+            membership.identifier: membership
+            for membership in CohortMembership.objects.filter(cohort=cohort).only(
+                "id", "identifier", "state"
+            )
+        }
+        CohortMembership.objects.bulk_create(
+            [
+                CohortMembership(cohort=cohort, identifier=identifier)
+                for identifier in extraction.identifiers
+                if identifier not in existing
+            ],
+            batch_size=COHORT_CSV_MEMBERSHIP_CREATE_BATCH_SIZE,
+        )
+        added += len(incoming - existing.keys())
+
+        readd_ids: list[int] = []
+        remove_ids: list[int] = []
+        for identifier, membership in existing.items():
+            if identifier in incoming:
+                if membership.state == CohortMembershipState.PENDING_REMOVE:
+                    readd_ids.append(membership.id)
+                else:
+                    unchanged += 1
+            elif membership.state != CohortMembershipState.PENDING_REMOVE:
+                # A pending add may have had its trait written by a concurrent
+                # applier run, so drain it via pending remove, never delete.
+                remove_ids.append(membership.id)
+
+        added += CohortMembership.objects.filter(id__in=readd_ids).update(
+            state=CohortMembershipState.PENDING_ADD, updated_at=timezone.now()
+        )
+        removed += CohortMembership.objects.filter(id__in=remove_ids).update(
+            state=CohortMembershipState.PENDING_REMOVE, updated_at=timezone.now()
+        )
+
+        locked_cohort.version += 1
+        locked_cohort.save(update_fields=["version"])
+        apply_cohort_membership_deltas.delay(kwargs={"cohort_id": cohort.id})
+
+    flagsmith_cohorts_csv_syncs_total.inc()
+    flagsmith_cohorts_csv_sync_identifiers.observe(len(incoming))
+    logger.info(
+        "csv.synced",
+        cohort__id=cohort.id,
+        environment__id=cohort.environment_id,
+        cohort__version=locked_cohort.version,
+        adds__count=added,
+        removes__count=removed,
+        unchanged__count=unchanged,
+    )
+    return CohortCsvSyncResult(
+        version=locked_cohort.version,
+        added=added,
+        removed=removed,
+        unchanged=unchanged,
+        ignored=CohortCsvIgnoredRows(
+            empty=extraction.empty_count,
+            duplicates=extraction.duplicate_count,
+            too_long=extraction.too_long_count,
+        ),
+    )
+
+
 def edge_sync_enabled(project: "Project") -> bool:
     return bool(project.enable_dynamo_db and DynamoIdentityWrapper().is_enabled)
 

api/cohorts/views.py

@@ -1,14 +1,21 @@
 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.decorators import action
+from rest_framework.parsers import MultiPartParser
 from rest_framework.permissions import IsAuthenticated
 from rest_framework.request import Request
 from rest_framework.response import Response
 
+from api.serializers import ErrorSerializer
 from cohorts import services
 from cohorts.models import Cohort
 from cohorts.permissions import CohortPermission, CohortPlanPermission
-from cohorts.serializers import CohortSerializer
+from cohorts.serializers import (
+    CohortCsvSyncResultSerializer,
+    CohortCsvSyncSerializer,
+    CohortSerializer,
+)
 from environments.views import NestedEnvironmentViewSet
 from projects.exceptions import DynamoNotEnabledError
 
@@ -62,3 +69,35 @@ 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(
+        description=(
+            "Replace the cohort's members with the identifiers found in the "
+            "uploaded CSV file and trigger a sync to identity data. "
+            "`identifier_column` is the 0-based index of the column holding "
+            "the identifiers; `has_header` skips the first row when true."
+        ),
+        request=CohortCsvSyncSerializer,
+        responses={
+            202: CohortCsvSyncResultSerializer,
+            400: ErrorSerializer,
+            413: ErrorSerializer,
+        },
+    )
+    @action(
+        detail=True,
+        methods=["post"],
+        url_path="sync-csv",
+        parser_classes=[MultiPartParser],
+    )
+    def sync_csv(self, request: Request, *args: object, **kwargs: object) -> Response:
+        cohort = self.get_object()
+        serializer = CohortCsvSyncSerializer(data=request.data)
+        serializer.is_valid(raise_exception=True)
+        result = services.sync_cohort_memberships_from_csv(
+            cohort=cohort, **serializer.validated_data
+        )
+        return Response(
+            CohortCsvSyncResultSerializer(result).data,
+            status=status.HTTP_202_ACCEPTED,
+        )