feat(Segment Membership): Seed organisations too large for the task processor

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 <organisation_id> [--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,
+        )

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,14 +67,18 @@ 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:
         log.warning("seed.skipped", reason="clickhouse_not_configured")
         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
 

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

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,

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,16 +644,16 @@ 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`
 
 ### `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`