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_WRITE_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, CohortSyncKey
 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,20 +33,36 @@ 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 CohortSyncKeySerializer(serializers.ModelSerializer[CohortSyncKey]):
@@ -69,8 +99,33 @@ def _validate_identifier_byte_length(value: str) -> None:
 
 
 class CohortSyncMembersSerializer(serializers.Serializer[None]):
-    # TODO: apply the same byte-length check to CSV uploads.
     user_ids = serializers.ListField(
         child=serializers.CharField(validators=[_validate_identifier_byte_length]),
         min_length=1,
     )
+
+
+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,17 +1,32 @@
+import csv
+import io
 import typing
 
 import structlog
 from django.db import transaction
-from django.db.models import QuerySet
+from django.db.models import F, QuerySet
 from django.utils import timezone
 from flag_engine.segments.constants import IS_SET
 from rest_framework.exceptions import ValidationError
 
 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.constants import (
+    COHORT_CSV_MEMBERSHIP_WRITE_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,
@@ -37,6 +52,15 @@
     CohortMembershipState.PENDING_REMOVE,
 ]
 
+_T = typing.TypeVar("_T")
+
+
+def _batched(
+    items: list[_T], size: int = COHORT_CSV_MEMBERSHIP_WRITE_BATCH_SIZE
+) -> typing.Iterator[list[_T]]:
+    for offset in range(0, len(items), size):
+        yield items[offset : offset + size]
+
 
 def pending_memberships(cohort: Cohort) -> "QuerySet[CohortMembership]":
     return CohortMembership.objects.filter(cohort=cohort, state__in=_PENDING_STATES)
@@ -212,6 +236,141 @@ def remove_cohort_members(cohort: Cohort, identifiers: "typing.Iterable[str]") -
     )
 
 
+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."})
+
+    # A sync is a full reconciliation towards the uploaded CSV: a partially
+    # failed or interleaved run converges on the next upload, and the unique
+    # constraint absorbs concurrent inserts. Writes are therefore chunked into
+    # their own implicit transactions instead of one long transaction that
+    # would hold locks against the applier task while a 10 MB file lands.
+    incoming = set(extraction.identifiers)
+    existing = {
+        membership.identifier: membership
+        for membership in CohortMembership.objects.filter(cohort=cohort).only(
+            "id", "identifier", "state"
+        )
+    }
+    to_create = [
+        identifier
+        for identifier in extraction.identifiers
+        if identifier not in existing
+    ]
+    present = incoming & existing.keys()
+    readd_ids = [
+        existing[identifier].id
+        for identifier in present
+        if existing[identifier].state == CohortMembershipState.PENDING_REMOVE
+    ]
+    # A departed pending add may have had its trait written by a concurrent
+    # applier run, so drain it via pending remove, never delete.
+    remove_ids = [
+        membership.id
+        for identifier, membership in existing.items()
+        if identifier not in incoming
+        and membership.state != CohortMembershipState.PENDING_REMOVE
+    ]
+
+    for identifier_batch in _batched(to_create):
+        CohortMembership.objects.bulk_create(
+            [
+                CohortMembership(cohort=cohort, identifier=identifier)
+                for identifier in identifier_batch
+            ],
+            ignore_conflicts=True,
+        )
+    added = len(to_create)
+    removed = 0
+    unchanged = len(present) - len(readd_ids)
+    for id_batch in _batched(readd_ids):
+        added += CohortMembership.objects.filter(id__in=id_batch).update(
+            state=CohortMembershipState.PENDING_ADD, updated_at=timezone.now()
+        )
+    for id_batch in _batched(remove_ids):
+        removed += CohortMembership.objects.filter(id__in=id_batch).update(
+            state=CohortMembershipState.PENDING_REMOVE, updated_at=timezone.now()
+        )
+
+    Cohort.objects.filter(id=cohort.id).update(version=F("version") + 1)
+    cohort.refresh_from_db(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=cohort.version,
+        adds__count=added,
+        removes__count=removed,
+        unchanged__count=unchanged,
+    )
+    return CohortCsvSyncResult(
+        version=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 delete_cohort(cohort: Cohort) -> None:
     from cohorts.tasks import apply_cohort_membership_deltas