feat(cohorts): Amplitude cohort sync endpoints and sync keys - #8290
feat(cohorts): Amplitude cohort sync endpoints and sync keys#8290gagantrivedi wants to merge 10 commits into
Conversation
|
The latest updates on your projects. Learn more about Vercel for GitHub. 3 Skipped Deployments
|
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Organization UI Review profile: ASSERTIVE Plan: Pro Plus Run ID: 📒 Files selected for processing (2)
Included review availability: Your plan provides up to 8 included reviews per hour; 7 remain after this review. 📝 WalkthroughWalkthroughAdds Bearer-token authentication and environment-scoped lifecycle management for cohort sync keys. Adds Amplitude cohort creation and member add/remove endpoints. Persists source types, records audit events, updates membership state, and documents the API and security scheme. Adds unit tests for authentication, validation, environment isolation, membership updates, audit records, and key revocation. Estimated code review effort: 4 (Complex) | ~60 minutes Merge Risk: 🟠 High · up to The new Amplitude cohort-sync endpoints can acknowledge requests that later fail to apply membership changes, especially for oversized identifiers or removals, unavailable synchronization services, transaction timing failures, or audit-write failures; the API schema also does not fully match runtime requirements. These are high-impact correctness and data-consistency risks, so the PR is not ready to merge without fixes or explicit acceptance. Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
github-actions
Bot
added
api
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## main #8290 +/- ##
==========================================
+ Coverage 98.64% 98.79% +0.14%
==========================================
Files 1609 1614 +5
Lines 64570 64934 +364
==========================================
+ Hits 63696 64150 +454
+ Misses 874 784 -90 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
github-actions
Bot
added
feature
github-actions
Bot
added
feature
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 4
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: ASSERTIVE
Plan: Pro Plus
Run ID: 5882fd93-bbb9-44ac-bfa1-bfd21628caaf
📒 Files selected for processing (16)
api/api/urls/v1.pyapi/cohorts/authentication.pyapi/cohorts/migrations/0003_cohort_sync_key.pyapi/cohorts/models.pyapi/cohorts/permissions.pyapi/cohorts/serializers.pyapi/cohorts/services.pyapi/cohorts/sync_urls.pyapi/cohorts/sync_views.pyapi/cohorts/urls.pyapi/cohorts/views.pyapi/tests/unit/cohorts/conftest.pyapi/tests/unit/cohorts/test_sync_views.pyapi/tests/unit/cohorts/test_views.pydocs/docs/deployment-self-hosting/observability/_events-catalogue.mdopenapi.yaml
Included review availability: Your plan includes up to 8 reviews per rolling hour; 7 remain after this review.
✅ oss · depot-ubuntu-latest-arm-16 — run #19533 (attempt 1)Playwright Test Results (oss - depot-ubuntu-latest-arm-16)
Details
🗂️ Previous results✅ oss · depot-ubuntu-latest-16 — run #19533 (attempt 1)Playwright Test Results (oss - depot-ubuntu-latest-16)
Details
✅ private-cloud · depot-ubuntu-latest-arm-16 — run #19532 (attempt 1)Playwright Test Results (private-cloud - depot-ubuntu-latest-arm-16)
Details
✅ private-cloud · depot-ubuntu-latest-16 — run #19532 (attempt 1)Playwright Test Results (private-cloud - depot-ubuntu-latest-16)
Details
✅ oss · depot-ubuntu-latest-arm-16 — run #19532 (attempt 1)Playwright Test Results (oss - depot-ubuntu-latest-arm-16)
Details
✅ oss · depot-ubuntu-latest-16 — run #19532 (attempt 1)Playwright Test Results (oss - depot-ubuntu-latest-16)
Details
❌ private-cloud · depot-ubuntu-latest-16 — run #19458 (attempt 1)Playwright Test Results (private-cloud - depot-ubuntu-latest-16)
Details
Failed testsfirefox › tests/environment-permission-test.pw.ts › Environment Permission Tests › Environment-level permissions control access to features, identities, and segments @enterprise ✅ private-cloud · depot-ubuntu-latest-arm-16 — run #19458 (attempt 1)Playwright Test Results (private-cloud - depot-ubuntu-latest-arm-16)
Details
✅ oss · depot-ubuntu-latest-arm-16 — run #19458 (attempt 1)Playwright Test Results (oss - depot-ubuntu-latest-arm-16)
Details
✅ oss · depot-ubuntu-latest-arm-16 — run #19459 (attempt 1)Playwright Test Results (oss - depot-ubuntu-latest-arm-16)
Details
✅ oss · depot-ubuntu-latest-16 — run #19458 (attempt 1)Playwright Test Results (oss - depot-ubuntu-latest-16)
Details
✅ oss · depot-ubuntu-latest-16 — run #19459 (attempt 1)Playwright Test Results (oss - depot-ubuntu-latest-16)
Details
|
Visual Regression19 screenshots compared. See report for details. |
github-actions
Bot
added
docs
There was a problem hiding this comment.
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (2)
openapi.yaml (2)
1528-1533: 🎯 Functional Correctness | 🟡 Minor | ⚡ Quick winDeclare list IDs as UUIDs.
AmplitudeCohortSyncViewSet._get_cohortparses the path value withuuid.UUIDand returns 404 for non-UUID values, as shown inapi/cohorts/sync_views.py, Lines [38-88]. The add and remove parameters currently declare onlytype: string. Addformat: uuidto both schemas so generated clients use the server's actual contract.Proposed OpenAPI fix
schema: type: string + format: uuidApply this to both
idpath parameters.Also applies to: 1556-1560
19015-19045: 🗄️ Data Integrity & Integration | 🟠 Major | ⚡ Quick winAlign the public identifier limit with persistence.
CohortSyncMembers.user_ids[*]permits up to 2000 characters, whileCohortMembership.identifierhas a 1024-byte limit. The suppliedapi/cohorts/serializers.py, Lines [63-71], confirms that a multibyte identifier can pass validation and fail when the membership row is written. A documented request can therefore fail after acceptance.Enforce the byte-length limit before creating the pending membership row, then publish the same effective constraint in this schema. Do not fix this only by reducing the character limit.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: ASSERTIVE
Plan: Pro Plus
Run ID: 91f1770b-4706-4a6f-8e3b-db3ad55bad7d
📒 Files selected for processing (2)
mcp/src/flagsmith_mcp/openapi.jsonopenapi.yaml
Included review availability: Your plan includes up to 8 reviews per rolling hour; 7 remain after this review.
b439ee4 to
3964c41
Compare
github-actions
Bot
added
feature
There was a problem hiding this comment.
Actionable comments posted: 4
♻️ Duplicate comments (1)
openapi.yaml (1)
1661-1737: 🔒 Security & Privacy | 🟠 Major | ⚡ Quick winBearer authentication is still undocumented for the Amplitude cohort-sync operations.
The three operations
api_v1_cohort_sync_amplitude_lists_create,api_v1_cohort_sync_amplitude_lists_add_create, andapi_v1_cohort_sync_amplitude_lists_remove_createhave nosecurityfield.CohortSyncKeyAuthenticationrequiresAuthorization: Bearer <sync-key>for these endpoints, per the PR objectives.A prior review thread on this file flagged the identical issue and recorded it as resolved by a commit that added a "Cohort Sync Key" security-scheme extension in
api/api/openapi.py, with the note that CI regenerates this rootopenapi.yamlautomatically from that source. In this file version, the fix has not taken effect:
- Lines 1661-1737 still show no
securityfield on the three operations.- The
securitySchemesblock (around line 29354) still lists onlyEnvironment API Key,Master API Key,basicAuth, andtokenAuth. No "Cohort Sync Key" scheme is defined.Confirm that the authentication extension still exists and is wired to these operations, and that the generation step that produces this file actually runs before merge.
#!/bin/bash # Description: Verify the Cohort Sync Key OpenAPI authentication extension exists and is applied. set -euo pipefail echo '--- Check for the authentication extension in api/api/openapi.py ---' fd -a openapi.py api/api | xargs -r rg -n -C 5 'CohortSyncKeyAuthentication|OpenApiAuthenticationExtension|Cohort Sync Key' echo '--- Check CohortSyncKeyAuthentication usage in cohorts sync views ---' fd -a sync_views.py api/cohorts | xargs -r rg -n -C 3 'authentication_classes|CohortSyncKeyAuthentication' echo '--- Check sdk/openapi.yaml for the scheme ---' rg -n -C 3 'Cohort Sync Key' sdk/openapi.yaml || echo 'not found in sdk/openapi.yaml'
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: ASSERTIVE
Plan: Pro Plus
Run ID: 94df6b2d-a068-4e8b-a583-933430a093e8
📒 Files selected for processing (7)
api/cohorts/services.pyapi/cohorts/sync_views.pyapi/cohorts/views.pyapi/tests/unit/cohorts/test_sync_views.pyapi/tests/unit/cohorts/test_views.pydocs/docs/deployment-self-hosting/observability/_events-catalogue.mdopenapi.yaml
Included review availability: Your plan provides up to 8 included reviews per hour; 7 remain after this review.
| 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 |
There was a problem hiding this comment.
🗄️ 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.
| 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 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}) |
There was a problem hiding this comment.
🎯 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' apiRepository: 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 || trueRepository: 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 || trueRepository: 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.
| serializer = CohortSyncMembersSerializer(data=request.data) | ||
| serializer.is_valid(raise_exception=True) | ||
| services.add_cohort_members(cohort, serializer.validated_data["user_ids"]) | ||
| return Response() | ||
|
|
||
| @action(detail=True, methods=["POST"]) | ||
| def remove(self, request: Request, pk: str) -> Response: | ||
| cohort = self._get_cohort(request, pk) | ||
| serializer = CohortSyncMembersSerializer(data=request.data) | ||
| serializer.is_valid(raise_exception=True) |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
Reject identifiers above 1,024 bytes before queueing membership work.
CohortSyncMembersSerializer accepts up to 2,000 characters. A multibyte user_id can therefore pass Lines 55 and 63, persist in CohortMembership, and return HTTP 200 before the asynchronous Edge write rejects it.
Validate UTF-8 byte length in the shared serializer and return HTTP 400. Apply the same validation to CSV membership input. Add a regression test with a value above 1,024 bytes.
Based on learnings, cohort identifiers must be rejected above 1,024 bytes with HTTP 400 to match Edge DynamoDB constraints.
Source: Learnings
| ### `cohorts.membership.deltas_received` | ||
|
|
||
| Logged at `info` from: | ||
| - `api/cohorts/services.py:168` | ||
| - `api/cohorts/services.py:187` | ||
|
|
||
| Attributes: | ||
| - `action` | ||
| - `cohort.id` | ||
| - `deltas.count` | ||
| - `environment.id` | ||
| - `members.matched` |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
Document members.matched as a remove-only attribute.
add_cohort_members does not emit members.matched. Only remove_cohort_members emits that attribute. Mark it as remove-only, or split the add and remove event contracts.
1 participant
docs/if required so people know about the feature.Changes
Contributes to Flagsmith/flagsmith-private#260.
The receiving side of Amplitude's list-based cohort sync contract, feeding the existing membership ledger. Spec copied from Create a cohort sync integration and Receiving behavioral cohorts — routes, field names and status codes are fixed by Amplitude.
CohortSyncKey: per-environment bearer credential, hashed at rest, plaintext returned once; managed at/environments/<key>/cohorts/sync-keys/./cohort-sync/amplitude/lists/creates the cohort and returnslist_id;lists/<uuid>/add|remove/turn{"user_ids": [...]}into pending ledger rows for the existing apply task.How did you test this code?
Unit tests cover the auth matrix, list creation, add/remove reaching (moto) DynamoDB via the apply task, the 404 fences, and key management.
Beyond that, this can only really be tested once Amplitude grants us Integration Portal access: the endpoints are exercised for real by Amplitude's own calls, first through the portal's testing tab and
Release internallymode, then in production.