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
39 changes: 29 additions & 10 deletions api/api/openapi.py
Original file line number Diff line number Diff line change
Expand Up @@ -165,19 +165,38 @@ def get_security_definition(

class CohortSyncKeyAuthenticationExtension(OpenApiAuthenticationExtension): # type: ignore[no-untyped-call]
target_class = "cohorts.authentication.CohortSyncKeyAuthentication"
name = "Cohort Sync Key"
name = ["Cohort Sync Key", "Cohort Sync Key (Basic)"]

def get_security_requirement(
self, auto_schema: openapi.AutoSchema
) -> list[dict[str, list[Any]]]:
# Separate entries: the caller sends the key with either scheme,
# not both at once.
return [{name: []} for name in self.name]
Comment on lines +168 to +175

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

🔎 Supported by static analysis

🏁 Script executed:

#!/usr/bin/env bash
set -euo pipefail

files=(
  api/api/openapi.py
  openapi.yaml
  mcp/src/flagsmith_mcp/openapi.json
  sdk/openapi.yaml
)

for file in "${files[@]}"; do
  test -f "$file"
done

if rg -n -P 'name\s*=\s*\[.*Cohort Sync Key|Cohort Sync Key(?: \(Basic\))?\s*:' "${files[@]}"; then
  echo "Found non-conforming cohort security-scheme names." >&2
  exit 1
fi

Repository: Flagsmith/flagsmith

Length of output: 885


🏁 Script executed:

#!/usr/bin/env bash
set -euo pipefail

printf '%s\n' '--- api/api/openapi.py ---'
sed -n '145,205p' api/api/openapi.py

printf '%s\n' '--- generated security schemes ---'
sed -n '555,575p' sdk/openapi.yaml
sed -n '29645,29672p' openapi.yaml

printf '%s\n' '--- generated security requirements ---'
sed -n '1675,1692p' openapi.yaml

Repository: Flagsmith/flagsmith

Length of output: 4676


🏁 Script executed:

#!/usr/bin/env bash
set -euo pipefail

printf '%s\n' '--- drf-spectacular version pins ---'
rg -n -m 20 'drf-spectacular|drf_spectacular' \
  pyproject.toml poetry.lock requirements*.txt setup.cfg setup.py 2>/dev/null || true

printf '%s\n' '--- extension references ---'
rg -n -m 20 'CohortSyncKeyAuthenticationExtension|get_security_requirement|get_security_definition' api/api/openapi.py

Repository: Flagsmith/flagsmith

Length of output: 485


🌐 Web query:

drf-spectacular 0.28.0 OpenApiAuthenticationExtension name list get_security_definition implementation OpenAPI component key pattern 3.0.3

💡 Result:

In drf-spectacular 0.28.0, the OpenApiAuthenticationExtension allows you to define authentication schemes that are registered in the OpenAPI components section [1][2]. Implementation for Single vs. Multiple Security Definitions: To implement a single security definition, set the name attribute to a string and return a dictionary from get_security_definition [1][3]. For scenarios requiring multiple parameters or a group of headers (e.g., to model security requirements that include multiple headers), you should define name as a list of strings and return a corresponding list of dictionary definitions from get_security_definition [1][2][4]. Example Implementation: class MyCustomAuthExtension(OpenApiAuthenticationExtension): target_class = 'path.to.YourAuthenticationClass' name = ['HeaderOne', 'HeaderTwo'] def get_security_definition(self, auto_schema): return [ {'type': 'apiKey', 'in': 'header', 'name': 'X-Header-One'}, {'type': 'apiKey', 'in': 'header', 'name': 'X-Header-Two'} ] OpenAPI Component Key Pattern: When drf-spectacular processes these extensions, it registers them under the securitySchemes component in the OpenAPI specification [5][6]. Under OpenAPI 3.0.3, these components are structured under the components/securitySchemes key [6]. The library ensures that each name provided in your extension is registered as a unique component key within the resulting OpenAPI schema [5]. If you provide a list of names, the library iterates through them and maps each to the corresponding definition returned by get_security_definition, ensuring they are correctly placed in the schema components [5]. Key Technical Details: - The get_security_definition method is expected to return a dictionary (or a list of dictionaries) compliant with the OpenAPI 3.0.3 Security Scheme Object specification [1][3]. - Internally, drf-spectacular uses the ResolvedComponent class to handle these, where the security component type is defined as securitySchemes [6]. - The library automatically warns if it detects component naming collisions [6]. For most standard implementations, refer to the official documentation on workflow and customization [3][7].

Citations:


🌐 Web query:

OpenAPI 3.0.3 specification components securitySchemes map keys pattern ^[a-zA-Z0-9.-_]+$

💡 Result:

In the OpenAPI 3.0.3 specification, the keys used within the components/securitySchemes object must conform to the regular expression ^[a-zA-Z0-9.-_]+$ [1][2][3]. This requirement applies generally to all fixed fields declared within the components object (such as schemas, responses, parameters, examples, requestBodies, headers, links, and callbacks) [2][3]. The specification explicitly mandates that the names used as keys in these maps must match this pattern [1][4]. While these names are arbitrary and used to refer to the security schemes from other parts of the API definition (such as the global or operation-level security requirements), they must adhere to the specified character constraints to be compliant with the OpenAPI 3.0.3 standard [5]. Note that while the textual specification imposes this constraint, it has been noted in community discussions that the corresponding JSON schema files for version 3.0 historically lacked the automated validation (such as additionalProperties: false) to strictly enforce this pattern, which was considered an oversight in the 3.0 schema design [6].

Citations:


Use OpenAPI-safe names for the security schemes.

CohortSyncKeyAuthenticationExtension.name generates the components.securitySchemes keys. The current names contain spaces and parentheses, which violate the OpenAPI 3.0.3 component-key pattern. The invalid names are also present in openapi.yaml and sdk/openapi.yaml. Rename both names and regenerate these documents.

Source: MCP tools


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."
),
}
) -> list[dict[str, Any]]:
return [
{
"type": "http",
"scheme": "bearer",
"description": (
"For cohort sync endpoints called by an external cohort "
"source, such as Amplitude."
),
},
{
"type": "http",
"scheme": "basic",
"description": (
"For cohort sync endpoints called by an external cohort "
"source that can only send Basic credentials, such as "
"Mixpanel. The key is the password; the username is "
"ignored."
),
},
]


# Tag definitions controlling the order and display of sections in the Swagger UI.
Expand Down
30 changes: 24 additions & 6 deletions api/cohorts/authentication.py
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
import typing
import base64
from contextlib import suppress

from django.contrib.auth.models import AnonymousUser
Expand All @@ -9,18 +9,36 @@


class CohortSyncKeyAuthentication(authentication.BaseAuthentication):
"""
Accepts a cohort sync key sent either as a Bearer token or as the
password of Basic credentials. Amplitude sends Bearer; Mixpanel's
webhook setup only offers a username/password form, so its customers
enter any username and the key as the password. The username is
ignored.
"""

def authenticate(
self, request: Request
) -> tuple[AnonymousUser, CohortSyncKey] | None:
header = request.headers.get("Authorization", "")
if not header.startswith("Bearer "):
if header.startswith("Bearer "):
raw_key = header.removeprefix("Bearer ")
elif header.startswith("Basic "):
try:
decoded = base64.b64decode(
header.removeprefix("Basic "), validate=True
).decode()
except ValueError:
# Covers malformed base64, header bytes outside ASCII, and
# decoded credentials that are not valid UTF-8.
raise exceptions.AuthenticationFailed("Invalid Basic credentials.")
# Split at the first colon, so a key containing colons survives.
_, _, raw_key = decoded.partition(":")
else:
return None

with suppress(CohortSyncKey.DoesNotExist):
key = typing.cast(
CohortSyncKey,
CohortSyncKey.objects.get_from_key(header.removeprefix("Bearer ")),
)
key = CohortSyncKey.objects.get_from_key(raw_key)
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
Expand Down
33 changes: 33 additions & 0 deletions api/cohorts/migrations/0004_mixpanel_source.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,33 @@
# Generated by Django 5.2.16 on 2026-08-20 10:06

from django.db import migrations, models


class Migration(migrations.Migration):

dependencies = [
("cohorts", "0003_cohort_sync_key"),
("environments", "0039_use_no_ssrf_url_field"),
("segments", "0032_add_segment_rules_data"),
]

operations = [
migrations.AddField(
model_name="cohort",
name="external_id",
field=models.CharField(blank=True, max_length=255, null=True),
),
migrations.AlterField(
model_name="cohort",
name="source_type",
field=models.CharField(
choices=[
("csv", "CSV"),
("amplitude", "Amplitude"),
("mixpanel", "Mixpanel"),
],
default="csv",
max_length=50,
),
),
]
22 changes: 21 additions & 1 deletion api/cohorts/models.py
Original file line number Diff line number Diff line change
@@ -1,5 +1,7 @@
import typing

from django.db import models
from rest_framework_api_key.models import AbstractAPIKey
from rest_framework_api_key.models import AbstractAPIKey, APIKeyManager

from cohorts.constants import COHORT_SYSTEM_TRAIT_KEY_PREFIX
from core.models import SoftDeleteExportableModel
Expand All @@ -8,6 +10,7 @@
class CohortSourceType(models.TextChoices):
CSV = "csv", "CSV"
AMPLITUDE = "amplitude", "Amplitude"
MIXPANEL = "mixpanel", "Mixpanel"


class Cohort(SoftDeleteExportableModel):
Expand All @@ -26,6 +29,11 @@ class Cohort(SoftDeleteExportableModel):
choices=CohortSourceType.choices,
default=CohortSourceType.CSV,
)
# The cohort's identifier in the external source. Mixpanel pushes under
# its own cohort ID, so we store it to route later requests; Amplitude
# uses the ID we hand back at list creation, and CSV cohorts have no
# external system, so both leave this null.
external_id = models.CharField(max_length=255, null=True, blank=True)
Comment on lines +32 to +36

@coderabbitai coderabbitai Bot Aug 21, 2026

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 | 🏗️ Heavy lift

Make active cohort lookup and creation atomic.

Concurrent initial members requests can both observe no cohort and create separate active cohorts for the same environment, source type, and external ID. Later incremental requests may select only one duplicate, splitting memberships between cohorts.

Add a partial unique constraint for active cohorts with non-null external IDs, and use an atomic get-or-create flow that re-fetches after an integrity conflict.

📍 Affects 2 files
  • api/cohorts/models.py#L32-L36 (this comment)
  • api/cohorts/sync_views.py#L140-L153

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

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

Deliberately dropped to keep v1 lean — we'll add the constraint and conflict recovery if duplicate creation is observed in practice.

Copy link
Copy Markdown

Choose a reason for hiding this comment

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

@gagantrivedi, understood. The concurrency risk remains deferred from this PR.

Would you like me to create a follow-up GitHub issue for the active-cohort uniqueness constraint and atomic conflict recovery?

🐇


✏️ Learnings added
Learnt from: gagantrivedi
Repo: Flagsmith/flagsmith PR: 8338
File: api/cohorts/models.py:32-36
Timestamp: 2026-08-24T04:25:26.515Z
Learning: For the Mixpanel cohort synchronisation v1, the maintainers deliberately defer the active-cohort uniqueness constraint and atomic conflict recovery for concurrent initial `members` requests. They will add this protection if duplicate cohort creation is observed in practice.

Note: Learnings are effective only in the context of similar code segments. To apply general coding guidelines, please configure review instructions. You can manage existing learnings in the app.

You are interacting with an AI system.

version = models.PositiveIntegerField(default=0)
created_at = models.DateTimeField(auto_now_add=True)
# Deletion drains memberships from the identity store first; the cohort is
Expand All @@ -48,7 +56,19 @@ class Meta:
]


class CohortSyncKeyManager(APIKeyManager):
def get_from_key(self, key: str) -> "CohortSyncKey":
if "\x00" in key:
# A NUL can't travel in a raw header, but base64 credentials can
# decode to one, and the database driver refuses to build a query
# containing it. No real key holds one, so treat it as absent.
raise self.model.DoesNotExist("Key contains a NUL character.")
return typing.cast("CohortSyncKey", super().get_from_key(key))


class CohortSyncKey(AbstractAPIKey):
objects: typing.ClassVar[CohortSyncKeyManager] = CohortSyncKeyManager()

environment = models.ForeignKey(
"environments.Environment",
on_delete=models.CASCADE,
Expand Down
22 changes: 22 additions & 0 deletions api/cohorts/serializers.py
Original file line number Diff line number Diff line change
Expand Up @@ -98,6 +98,28 @@ def _validate_identifier_byte_length(value: str) -> None:
)


class MixpanelMemberSerializer(serializers.Serializer[None]):
mixpanel_distinct_id = serializers.CharField(
validators=[_validate_identifier_byte_length]
)


class MixpanelParametersSerializer(serializers.Serializer[None]):
mixpanel_cohort_id = serializers.CharField(max_length=255)
mixpanel_cohort_name = serializers.CharField(max_length=2000)
# An empty page is valid: a first sync of an empty cohort has no members.
members = MixpanelMemberSerializer(many=True, allow_empty=True)


class MixpanelWebhookSerializer(serializers.Serializer[None]):
# "members" carries the full membership on the first sync;
# "add_members"/"remove_members" carry changes since the last sync.
action = serializers.ChoiceField(
choices=["members", "add_members", "remove_members"]
)
parameters = MixpanelParametersSerializer()


class CohortSyncMembersSerializer(serializers.Serializer[None]):
user_ids = serializers.ListField(
child=serializers.CharField(validators=[_validate_identifier_byte_length]),
Expand Down
29 changes: 27 additions & 2 deletions api/cohorts/services.py
Original file line number Diff line number Diff line change
Expand Up @@ -116,6 +116,7 @@ def create_cohort(
name: str,
description: str | None = None,
source_type: CohortSourceType = CohortSourceType.CSV,
external_id: str | None = None,
) -> Cohort:
project = environment.project
# Mirrors the segment limit enforced by SegmentSerializer, which cohort
Expand All @@ -137,7 +138,10 @@ def create_cohort(
)
rule = SegmentRule.objects.create(segment=segment, type=SegmentRule.ALL_RULE)
cohort: Cohort = Cohort.objects.create(
environment=environment, segment=segment, source_type=source_type
environment=environment,
segment=segment,
source_type=source_type,
external_id=external_id,
)
Condition.objects.create(
rule=rule,
Expand All @@ -161,10 +165,16 @@ def create_cohort_for_source(
environment: "Environment",
name: str,
source_type: CohortSourceType,
external_id: str | None = None,
) -> 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)
cohort = create_cohort(
environment=environment,
name=name,
source_type=source_type,
external_id=external_id,
)
# 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
Expand All @@ -182,6 +192,21 @@ def create_cohort_for_source(
return cohort


def get_cohort_for_source(
*,
environment: "Environment",
source_type: CohortSourceType,
external_id: str,
) -> Cohort | None:
cohort: Cohort | None = Cohort.objects.filter(
environment=environment,
source_type=source_type,
external_id=external_id,
deletion_requested_at__isnull=True,
).first()
return cohort


def add_cohort_members(cohort: Cohort, identifiers: "typing.Iterable[str]") -> None:
from cohorts.tasks import apply_cohort_membership_deltas

Expand Down
8 changes: 6 additions & 2 deletions api/cohorts/sync_urls.py
Original file line number Diff line number Diff line change
@@ -1,11 +1,15 @@
from django.urls import path
from rest_framework.routers import SimpleRouter

from cohorts.sync_views import AmplitudeCohortSyncViewSet
from cohorts.sync_views import AmplitudeCohortSyncViewSet, MixpanelCohortSyncView

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
urlpatterns = [
path("mixpanel/webhook/", MixpanelCohortSyncView.as_view(), name="mixpanel"),
*router.urls,
]
Loading
Loading