fix(cohorts): queue membership count refreshes for cohort segments

Author: gagantrivedi

api/cohorts/services.py

@@ -1,9 +1,11 @@
 import csv
 import io
 import typing
+from datetime import timedelta
 from uuid import UUID
 
 import structlog
+from django.conf import settings
 from django.db import transaction
 from django.db.models import F, QuerySet
 from django.utils import timezone
@@ -101,6 +103,21 @@ def apply_pending_memberships(cohort: Cohort) -> bool:
         removed_count
     )
     if added_count or removed_count:
+        if settings.CLICKHOUSE_ENABLED:
+            from segment_membership.services import enqueue_membership_refresh
+
+            # Cohort membership is the one way a segment's members change
+            # without a segment or identity edit, so nothing else queues a
+            # count refresh. The delay lets the identity changes reach
+            # ClickHouse before the recount; repeated batches collapse into
+            # the one pending refresh.
+            enqueue_membership_refresh(
+                environment.project,
+                delay_until=timezone.now()
+                + timedelta(
+                    seconds=settings.SEGMENT_MEMBERSHIP_DELETE_REFRESH_DELAY_SECONDS
+                ),
+            )
         logger.info(
             "membership.applied",
             cohort__id=cohort.id,
@@ -158,6 +175,13 @@ def create_cohort(
         project__id=environment.project_id,
         organisation__id=environment.project.organisation_id,
     )
+    if settings.CLICKHOUSE_ENABLED:
+        from segment_membership.services import enqueue_membership_refresh
+
+        # The segment is created directly rather than through the segment
+        # serializer, which is where membership count refreshes are normally
+        # queued from.
+        enqueue_membership_refresh(environment.project)
     return cohort
 
 

api/tests/unit/cohorts/test_services.py

@@ -1,7 +1,9 @@
 import io
 
 import pytest
+from django.utils import timezone
 from flag_engine.segments.constants import IS_SET
+from pytest_django.fixtures import SettingsWrapper
 from pytest_mock import MockerFixture
 from pytest_structlog import StructuredLogCapture
 from rest_framework.exceptions import ValidationError
@@ -546,3 +548,44 @@ def test_sync_cohort_memberships_from_csv__edge_cohort__applies_traits(
     assert document["system_traits"] == {edge_cohort.system_trait_key: True}
     membership = CohortMembership.objects.get(cohort=edge_cohort)
     assert membership.state == CohortMembershipState.APPLIED
+
+
+def test_create_cohort__clickhouse_enabled__queues_membership_refresh(
+    environment: Environment,
+    settings: SettingsWrapper,
+    mocker: MockerFixture,
+) -> None:
+    # Given
+    settings.CLICKHOUSE_ENABLED = True
+    enqueue_mock = mocker.patch(
+        "segment_membership.services.enqueue_membership_refresh"
+    )
+
+    # When
+    create_cohort(environment=environment, name="Beta users")
+
+    # Then
+    enqueue_mock.assert_called_once_with(environment.project)
+
+
+def test_apply_pending_memberships__clickhouse_enabled__queues_delayed_refresh(
+    cohort: Cohort,
+    settings: SettingsWrapper,
+    mocker: MockerFixture,
+) -> None:
+    # Given
+    settings.CLICKHOUSE_ENABLED = True
+    settings.SEGMENT_MEMBERSHIP_DELETE_REFRESH_DELAY_SECONDS = 60
+    enqueue_mock = mocker.patch(
+        "segment_membership.services.enqueue_membership_refresh"
+    )
+    CohortMembership.objects.create(cohort=cohort, identifier="user-1")
+
+    # When
+    apply_pending_memberships(cohort)
+
+    # Then
+    enqueue_mock.assert_called_once()
+    args, kwargs = enqueue_mock.call_args
+    assert args == (cohort.environment.project,)
+    assert kwargs["delay_until"] > timezone.now()