Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
17 changes: 17 additions & 0 deletions api/api/openapi.py
Original file line number Diff line number Diff line change
Expand Up @@ -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]] = [
{
Expand Down
1 change: 1 addition & 0 deletions api/api/urls/v1.py
Original file line number Diff line number Diff line change
Expand Up @@ -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")),
Expand Down
34 changes: 34 additions & 0 deletions api/cohorts/authentication.py
Original file line number Diff line number Diff line change
@@ -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"
92 changes: 92 additions & 0 deletions api/cohorts/migrations/0003_cohort_sync_key.py
Original file line number Diff line number Diff line change
@@ -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,
},
),
]
17 changes: 17 additions & 0 deletions api/cohorts/models.py
Original file line number Diff line number Diff line change
@@ -1,11 +1,13 @@
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


class CohortSourceType(models.TextChoices):
CSV = "csv", "CSV"
AMPLITUDE = "amplitude", "Amplitude"


class Cohort(SoftDeleteExportableModel):
Expand Down Expand Up @@ -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"
Expand Down
7 changes: 7 additions & 0 deletions api/cohorts/permissions.py
Original file line number Diff line number Diff line change
Expand Up @@ -7,13 +7,20 @@
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
from users.models import FFAdminUser

_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)


Expand Down
38 changes: 37 additions & 1 deletion api/cohorts/serializers.py
Original file line number Diff line number Diff line change
Expand Up @@ -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


Expand Down Expand Up @@ -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
)
Comment thread
gagantrivedi marked this conversation as resolved.
90 changes: 88 additions & 2 deletions api/cohorts/services.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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(
Expand All @@ -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,
Expand All @@ -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
Comment on lines +122 to +145

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🗄️ Data Integrity & Integration | 🟠 Major | ⚡ Quick win

Make source cohort creation and audit creation atomic.

create_cohort() commits before AuditLog.objects.create() runs. If the audit write fails, the endpoint returns an error but leaves a cohort without the audit event that rebuilds the environment document. A retry can then create another cohort.

Wrap both operations in an outer transaction.atomic() block. Add a regression test that makes the audit write fail and asserts that no cohort remains.

Proposed fix
 def create_cohort_for_source(...):
-    cohort = create_cohort(...)
-    AuditLog.objects.create(...)
+    with transaction.atomic():
+        cohort = create_cohort(...)
+        AuditLog.objects.create(...)
     return cohort
📝 Committable suggestion

‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.

Suggested change
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 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."""
with transaction.atomic():
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})
Comment thread
coderabbitai[bot] marked this conversation as resolved.
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})
Comment on lines +177 to +186

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🎯 Functional Correctness | 🟠 Major | ⚡ Quick win

🧩 Analysis chain

🏁 Script executed:

#!/bin/bash
set -eu

printf '%s\n' '--- candidate files ---'
git ls-files 'api/cohorts/services.py' 'api/cohorts/tests*' 'api/**' | sed -n '1,120p'

printf '%s\n' '--- service structure ---'
ast-grep outline api/cohorts/services.py --view expanded | sed -n '1,220p'

printf '%s\n' '--- relevant service code ---'
sed -n '1,240p' api/cohorts/services.py

printf '%s\n' '--- references to removal and batching ---'
rg -n -C 4 'remove_cohort_members|batch_size|identifier__in|apply_cohort_membership_deltas' api

Repository: Flagsmith/flagsmith

Length of output: 35347


🏁 Script executed:

#!/bin/bash
set -eu

printf '%s\n' '--- cohort tests and API validation ---'
fd -i 'cohort' api/tests api/cohorts | sed -n '1,160p'
rg -n -C 5 'CohortSyncMembersSerializer|user_ids|add_cohort_members|remove_cohort_members' api/tests api/cohorts

printf '%s\n' '--- model and database configuration ---'
rg -n -C 8 'class CohortMembership|identifier\s*=|DATABASES|ENGINE|postgres|PostgreSQL|psycopg' api/cohorts api/app api/requirements* api/pyproject.toml api/setup.cfg api/tox.ini 2>/dev/null || true

printf '%s\n' '--- relevant constants and task behaviour ---'
cat -n api/cohorts/constants.py
sed -n '1,90p' api/cohorts/tasks.py

printf '%s\n' '--- tracked dependency/config files ---'
git ls-files | rg '(^|/)(requirements[^/]*|pyproject\.toml|setup\.cfg|tox\.ini|Pipfile|.*settings.*\.py)$' | sed -n '1,160p'

Repository: Flagsmith/flagsmith

Length of output: 49639


🏁 Script executed:

#!/bin/bash
set -eu

python3 - <<'PY'
from pathlib import Path
import re

services = Path("api/cohorts/services.py").read_text()
serializers = Path("api/cohorts/serializers.py").read_text()
settings = Path("api/app/settings/common.py").read_text()
models = Path("api/cohorts/models.py").read_text()

print("serializer_list_max_length=",
      re.search(r"user_ids\s*=\s*serializers\.ListField\((.*?)\n\s*\)", serializers, re.S).group(1).find("max_length=") >= 0)
print("service_uses_set=", "unique_identifiers = set(identifiers)" in services)
print("service_uses_single_identifier_in=",
      "identifier__in=unique_identifiers" in services)
print("default_database_is_postgresql=",
      '"ENGINE": "django.db.backends.postgresql"' in settings)
print("membership_identifier_max_length=",
      re.search(r"identifier\s*=\s*models\.CharField\(max_length=(\d+)\)", models).group(1))
print("add_batch_size=",
      re.search(r"bulk_create\(\s*rows,.*?batch_size=(\d+)", services, re.S).group(1))

# PostgreSQL's extended-query protocol uses an Int16 parameter count.
postgresql_bind_limit = 2**16 - 1
# The update binds at least one value per identifier, plus the cohort
# predicate and update values. This lower bound is sufficient to show that
# an unbounded IN list can exceed the protocol limit.
fixed_update_parameters = 3  # cohort, state, updated_at
for identifiers in (postgresql_bind_limit - fixed_update_parameters,
                    postgresql_bind_limit - fixed_update_parameters + 1):
    total = identifiers + fixed_update_parameters
    print(f"identifiers={identifiers} total_update_parameters={total} "
          f"within_limit={total <= postgresql_bind_limit}")
PY

printf '%s\n' '--- exact removal tests ---'
sed -n '180,250p' api/tests/unit/cohorts/test_sync_views.py

printf '%s\n' '--- model uniqueness and database routers ---'
sed -n '66,105p' api/cohorts/models.py
rg -n -C 5 'class .*Router|db_for_write|allow_migrate|default' api/app/routers.py api/task_processor/routers.py api/cohorts 2>/dev/null || true

Repository: Flagsmith/flagsmith

Length of output: 17487


🏁 Script executed:

#!/bin/bash
set -eu

python3 - <<'PY'
from pathlib import Path
import re

services = Path("api/cohorts/services.py").read_text()
serializers = Path("api/cohorts/serializers.py").read_text()
settings = Path("api/app/settings/common.py").read_text()
models = Path("api/cohorts/models.py").read_text()

print("serializer_list_max_length=",
      re.search(r"user_ids\s*=\s*serializers\.ListField\((.*?)\n\s*\)", serializers, re.S).group(1).find("max_length=") >= 0)
print("service_uses_set=", "unique_identifiers = set(identifiers)" in services)
print("service_uses_single_identifier_in=",
      "identifier__in=unique_identifiers" in services)
print("default_database_is_postgresql=",
      '"ENGINE": "django.db.backends.postgresql"' in settings)
print("membership_identifier_max_length=",
      re.search(r"identifier\s*=\s*models\.CharField\(max_length=(\d+)\)", models).group(1))
print("add_batch_size=",
      re.search(r"bulk_create\(\s*rows,.*?batch_size=(\d+)", services, re.S).group(1))

postgresql_bind_limit = 2**16 - 1
fixed_update_parameters = 3
for identifiers in (postgresql_bind_limit - fixed_update_parameters,
                    postgresql_bind_limit - fixed_update_parameters + 1):
    total = identifiers + fixed_update_parameters
    print(f"identifiers={identifiers} total_update_parameters={total} "
          f"within_limit={total <= postgresql_bind_limit}")
PY

printf '%s\n' '--- exact removal tests ---'
sed -n '180,250p' api/tests/unit/cohorts/test_sync_views.py

printf '%s\n' '--- model uniqueness and database routers ---'
sed -n '66,105p' api/cohorts/models.py
rg -n -C 5 'class .*Router|db_for_write|allow_migrate|default' api/app/routers.py api/task_processor/routers.py api/cohorts 2>/dev/null || true

Repository: Flagsmith/flagsmith

Length of output: 17487


Batch membership removals.

The unbounded identifier__in query can exceed PostgreSQL’s 65,535-parameter limit and return HTTP 500. Split unique_identifiers into bounded chunks, such as 1,000, within the transaction. Accumulate matched across updates and dispatch the task once. Add a regression test with more than one chunk.

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

Expand Down
11 changes: 11 additions & 0 deletions api/cohorts/sync_urls.py
Original file line number Diff line number Diff line change
@@ -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
Loading
Loading