From 79966d731b06318f639b37287a751270d1e557aa Mon Sep 17 00:00:00 2001 From: Matthew Elwell Date: Mon, 14 Sep 2026 10:41:09 +0100 Subject: [PATCH] feat(Segment Membership): Seed organisations too large for the task processor A seed that outlives its 4h task timeout is marked failed while its thread keeps running, retried alongside that thread, and re-enqueued hourly once the retries are exhausted. A management command calls the task inline instead, so it can run as a standalone ECS task that no deployment interrupts and nothing retries. `--ignore-feature-flag` populates ClickHouse before an organisation is given the feature, so the first membership counts it sees are complete. Co-Authored-By: Claude Opus 5 --- api/segment_membership/management/__init__.py | 0 .../management/commands/__init__.py | 0 .../commands/seed_segment_membership.py | 50 ++++++++++++++++ api/segment_membership/tasks.py | 11 +++- .../test_unit_segment_membership_commands.py | 60 +++++++++++++++++++ .../test_unit_segment_membership_tasks.py | 22 +++++++ .../observability/_events-catalogue.md | 18 +++--- 7 files changed, 150 insertions(+), 11 deletions(-) create mode 100644 api/segment_membership/management/__init__.py create mode 100644 api/segment_membership/management/commands/__init__.py create mode 100644 api/segment_membership/management/commands/seed_segment_membership.py create mode 100644 api/tests/unit/segment_membership/test_unit_segment_membership_commands.py diff --git a/api/segment_membership/management/__init__.py b/api/segment_membership/management/__init__.py new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/api/segment_membership/management/commands/__init__.py b/api/segment_membership/management/commands/__init__.py new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/api/segment_membership/management/commands/seed_segment_membership.py b/api/segment_membership/management/commands/seed_segment_membership.py new file mode 100644 index 000000000000..3caa6b593ae0 --- /dev/null +++ b/api/segment_membership/management/commands/seed_segment_membership.py @@ -0,0 +1,50 @@ +"""Run an organisation's identity seed out of band. + +The task processor is a poor host for an organisation large enough that the +seed outlives its task timeout: the run is marked failed while its thread keeps +going, retried alongside the thread that is still running, and re-enqueued +hourly once the retries are exhausted. Calling the task directly sidesteps all +of that, so the seed can be run as a standalone ECS task that no deployment +interrupts and nothing retries. + + flagsmith seed_segment_membership [--ignore-feature-flag] + +Base the ECS task definition on `flagsmith-task-processor`: it carries the +ClickHouse URL, the Dynamo table names, and the Flagsmith-on-Flagsmith server +key that the feature flag check needs. +""" + +from argparse import ArgumentParser +from typing import Any + +from django.core.management import BaseCommand + +from segment_membership.tasks import seed_organisation_identities + + +class Command(BaseCommand): + help = "Mirror an organisation's Dynamo identities into ClickHouse." + + def add_arguments(self, parser: ArgumentParser) -> None: + parser.add_argument("organisation_id", type=int) + parser.add_argument( + "--ignore-feature-flag", + action="store_true", + help=( + "Seed even though segment_membership_inspection is off for the " + "organisation, so ClickHouse can be populated before the " + "feature is exposed to it." + ), + ) + + def handle( + self, + *args: Any, + organisation_id: int, + ignore_feature_flag: bool, + **options: Any, + ) -> None: + seed_organisation_identities( + organisation_id, + ignore_feature_flag=ignore_feature_flag, + ) diff --git a/api/segment_membership/tasks.py b/api/segment_membership/tasks.py index bab40ac1717e..3e334d23360c 100644 --- a/api/segment_membership/tasks.py +++ b/api/segment_membership/tasks.py @@ -55,7 +55,10 @@ # 4h fits several large environments back-to-back at SaaS scale. timeout=timedelta(hours=4), ) -def seed_organisation_identities(organisation_id: int) -> None: +def seed_organisation_identities( + organisation_id: int, + ignore_feature_flag: bool = False, +) -> None: """Mirror one organisation's current Dynamo identities into IDENTITIES, dispatching a refresh per project as each completes. @@ -64,6 +67,10 @@ def seed_organisation_identities(organisation_id: int) -> None: Identities carrying no traits are skipped as they carry little to no value for segment membership. + + `ignore_feature_flag` lets an operator populate ClickHouse before the + organisation is let into the beta, so the first counts it sees are complete; + see the `seed_segment_membership` management command. """ log = logger.bind(organisation__id=organisation_id) if not settings.CLICKHOUSE_ENABLED: @@ -71,7 +78,7 @@ def seed_organisation_identities(organisation_id: int) -> None: return organisation = Organisation.objects.get(pk=organisation_id) - if not is_membership_enabled(organisation): + if not ignore_feature_flag and not is_membership_enabled(organisation): log.info("seed.skipped", reason="ff_disabled") return diff --git a/api/tests/unit/segment_membership/test_unit_segment_membership_commands.py b/api/tests/unit/segment_membership/test_unit_segment_membership_commands.py new file mode 100644 index 000000000000..333747c435ad --- /dev/null +++ b/api/tests/unit/segment_membership/test_unit_segment_membership_commands.py @@ -0,0 +1,60 @@ +from io import StringIO +from unittest.mock import MagicMock + +from django.core.management import call_command +from mypy_boto3_dynamodb.service_resource import Table +from pytest_django.fixtures import SettingsWrapper +from pytest_mock import MockerFixture + +from projects.models import Project +from segment_membership.models import SegmentMembershipSeed +from segments.models import Segment + +COMMAND = "seed_segment_membership" + + +def test_seed_segment_membership__organisation_id__runs_the_seed_inline( + mocker: MockerFixture, + project: Project, +) -> None: + # Given + seed = mocker.patch( + "segment_membership.management.commands." + "seed_segment_membership.seed_organisation_identities" + ) + + # When + call_command(COMMAND, project.organisation_id, stdout=StringIO()) + + # Then the task runs in this process rather than being queued, so no + # task timeout applies and nothing retries it + seed.assert_called_once_with(project.organisation_id, ignore_feature_flag=False) + seed.delay.assert_not_called() + + +def test_seed_segment_membership__ignore_feature_flag__seeds_with_flag_off( + mocker: MockerFixture, + settings: SettingsWrapper, + project: Project, + segment: Segment, + flagsmith_identities_table: Table, +) -> None: + # Given the organisation has not been let into the beta yet + settings.CLICKHOUSE_ENABLED = True + cursor = MagicMock() + open_cursor = mocker.patch("segment_membership.tasks.open_clickhouse_cursor") + open_cursor.return_value.__enter__.return_value = cursor + mocker.patch("segment_membership.tasks.enqueue_membership_refresh") + + # When + call_command( + COMMAND, + project.organisation_id, + "--ignore-feature-flag", + stdout=StringIO(), + ) + + # Then the seed runs to completion regardless + assert SegmentMembershipSeed.objects.filter( + organisation=project.organisation, seeded_at__isnull=False + ).exists() diff --git a/api/tests/unit/segment_membership/test_unit_segment_membership_tasks.py b/api/tests/unit/segment_membership/test_unit_segment_membership_tasks.py index 1e66ed240873..6c6203622dd7 100644 --- a/api/tests/unit/segment_membership/test_unit_segment_membership_tasks.py +++ b/api/tests/unit/segment_membership/test_unit_segment_membership_tasks.py @@ -115,6 +115,28 @@ def test_seed_organisation_identities__flag_off__skips( ).exists() +def test_seed_organisation_identities__flag_off_but_ignored__seeds( + mocker: MockerFixture, + settings: SettingsWrapper, + project: Project, + segment: Segment, + flagsmith_identities_table: Table, +) -> None: + # Given the flag is off for the organisation + settings.CLICKHOUSE_ENABLED = True + spy = mocker.patch.object(tasks, "open_clickhouse_cursor") + mocker.patch.object(tasks, "enqueue_membership_refresh") + + # When + seed_organisation_identities(project.organisation_id, ignore_feature_flag=True) + + # Then + spy.assert_called_once() + assert SegmentMembershipSeed.objects.filter( + organisation=project.organisation, seeded_at__isnull=False + ).exists() + + def test_seed_organisation_identities__insert_fails__logs_and_continues( mocker: MockerFixture, settings: SettingsWrapper, diff --git a/docs/docs/deployment-self-hosting/observability/_events-catalogue.md b/docs/docs/deployment-self-hosting/observability/_events-catalogue.md index 3826b5823873..af4e9c888c25 100644 --- a/docs/docs/deployment-self-hosting/observability/_events-catalogue.md +++ b/docs/docs/deployment-self-hosting/observability/_events-catalogue.md @@ -634,7 +634,7 @@ Attributes: ### `segment_membership.refresh.project.completed` Logged at `info` from: - - `api/segment_membership/tasks.py:272` + - `api/segment_membership/tasks.py:279` Attributes: - `membership_counts.count` @@ -644,7 +644,7 @@ Attributes: ### `segment_membership.refresh.project.failed` Logged at `exception` from: - - `api/segment_membership/tasks.py:245` + - `api/segment_membership/tasks.py:252` Attributes: - `project.id` @@ -652,8 +652,8 @@ Attributes: ### `segment_membership.refresh.project.skipped` Logged at `info` from: - - `api/segment_membership/tasks.py:212` - - `api/segment_membership/tasks.py:224` + - `api/segment_membership/tasks.py:219` + - `api/segment_membership/tasks.py:231` Attributes: - `project.id` @@ -663,7 +663,7 @@ Attributes: ### `segment_membership.seed.environment.completed` Logged at `info` from: - - `api/segment_membership/tasks.py:127` + - `api/segment_membership/tasks.py:134` Attributes: - `environment.id` @@ -674,7 +674,7 @@ Attributes: ### `segment_membership.seed.environment.failed` Logged at `exception` from: - - `api/segment_membership/tasks.py:120` + - `api/segment_membership/tasks.py:127` Attributes: - `environment.id` @@ -684,9 +684,9 @@ Attributes: ### `segment_membership.seed.skipped` Logged at `warning` from: - - `api/segment_membership/tasks.py:70` - - `api/segment_membership/tasks.py:75` - - `api/segment_membership/tasks.py:80` + - `api/segment_membership/tasks.py:77` + - `api/segment_membership/tasks.py:82` + - `api/segment_membership/tasks.py:87` Attributes: - `organisation.id`