perf(Segment Membership): Members list improvements (#7924)

Author: khvn26

api/segment_membership/services.py

@@ -129,7 +129,8 @@ def compute_segment_counts_for_project(
             f"SELECT {seg.id} AS segment_id, "
             f"i.environment_id AS env_key, count() AS c "
             f"FROM IDENTITIES AS i FINAL "
-            f"WHERE i.environment_id IN %(env_keys)s AND ({predicate}) "
+            f"WHERE i.environment_id IN %(env_keys)s "
+            f"AND i.is_deleted = false AND ({predicate}) "
             f"GROUP BY i.environment_id"
         )
 
@@ -186,7 +187,10 @@ def get_segment_members_page(
         )
         return []
 
-    conditions = ["i.environment_id = %(env_key)s"]
+    conditions = [
+        "i.environment_id = %(env_key)s",
+        "i.is_deleted = false",
+    ]
     params: dict[str, Any] = {"env_key": environment.api_key, "limit": limit}
     if cursor:
         conditions.append("i.identifier > %(cursor)s")
@@ -196,13 +200,29 @@ def get_segment_members_page(
         params["q"] = q
     conditions.append(f"({predicate})")
 
-    sql = (
-        "SELECT i.identifier, i.identity_key, i.traits "
-        "FROM IDENTITIES AS i FINAL "
+    # Why two queries? `traits` is a wide JSON column split over many per-path files.
+    # That fans out into thousands of object-storage requests, which is what
+    # makes a cold read slow. Delegating the `traits` read to an outer query
+    # limits them to the page.
+    #
+    # Besides that, bypass FINAL, which is heavy, in favour of LIMIT 1 BY i.identifier,
+    # which is (tolerably) less correct but substantially faster.
+    inner_sql = (
+        "SELECT i.identifier "
+        "FROM IDENTITIES AS i "
         f"WHERE {' AND '.join(conditions)} "
-        "ORDER BY i.identifier ASC "
+        "ORDER BY i.identifier ASC, i.inserted_at DESC "
+        "LIMIT 1 BY i.identifier "
         "LIMIT %(limit)s"
     )
+    sql = (
+        "SELECT m.identifier, m.identity_key, m.traits "
+        "FROM IDENTITIES AS m "
+        "WHERE m.environment_id = %(env_key)s AND m.is_deleted = false "
+        f"AND m.identifier IN ({inner_sql}) "
+        "ORDER BY m.identifier ASC, m.inserted_at DESC "
+        "LIMIT 1 BY m.identifier"
+    )
     log_comment = (
         "flagsmith:segment_membership:read"
         f":org_{segment.project.organisation_id}"

api/tests/unit/segment_membership/test_unit_segment_membership_services.py

@@ -1,7 +1,9 @@
+from datetime import datetime
 from unittest.mock import MagicMock
 
 import pytest
 from common.test_tools import RunTasksFixture
+from django.db import connections
 from flag_engine.segments.constants import EQUAL
 from pytest_django.fixtures import SettingsWrapper
 from pytest_mock import MockerFixture
@@ -186,6 +188,96 @@ def matching_segment(segment: Segment) -> Segment:
     return segment
 
 
+@pytest.mark.clickhouse
+def test_get_segment_members_page__deleted_identity__excluded(
+    segment_membership_identities: None,
+    matching_segment: Segment,
+    environment: Environment,
+    environment_api_key_str: str,
+) -> None:
+    # Given
+    with connections["clickhouse"].cursor() as cursor:
+        cursor.executemany(
+            "INSERT INTO IDENTITIES "
+            "(environment_id, identifier, identity_key, traits, is_deleted) VALUES",
+            [(environment_api_key_str, "aaron", "aaron_key", {"foo": "bar"}, True)],  # type: ignore[list-item]
+        )
+
+    # When
+    members = get_segment_members_page(
+        matching_segment, environment, cursor=None, limit=100
+    )
+
+    # Then
+    assert [member["identifier"] for member in members] == ["alice", "bob"]
+
+
+@pytest.mark.clickhouse
+def test_get_segment_members_page__duplicate_versions__returns_latest_once(
+    segment_membership_identities: None,
+    matching_segment: Segment,
+    environment: Environment,
+    environment_api_key_str: str,
+) -> None:
+    # Given
+    # a newer version of alice that still matches the segment
+    with connections["clickhouse"].cursor() as cursor:
+        cursor.executemany(
+            "INSERT INTO IDENTITIES "
+            "(environment_id, identifier, identity_key, traits, inserted_at) VALUES",
+            [
+                (  # type: ignore[list-item]
+                    environment_api_key_str,
+                    "alice",
+                    "alice_key",
+                    {"foo": "bar", "version": "new"},
+                    datetime(2099, 1, 1),
+                )
+            ],
+        )
+
+    # When
+    members = get_segment_members_page(
+        matching_segment, environment, cursor=None, limit=100
+    )
+
+    # Then
+    # alice appears once, with the latest version's traits
+    assert members == [
+        SegmentMember(
+            identifier="alice",
+            identity_key="alice_key",
+            traits={"foo": "bar", "version": "new"},
+        ),
+        SegmentMember(identifier="bob", identity_key="bob_key", traits={"foo": "bar"}),
+    ]
+
+
+@pytest.mark.clickhouse
+def test_compute_segment_counts_for_project__deleted_identity__excluded_from_count(
+    segment_membership_identities: None,
+    matching_segment: Segment,
+    project: Project,
+    environment_api_key_str: str,
+) -> None:
+    # Given
+    with connections["clickhouse"].cursor() as cursor:
+        cursor.executemany(
+            "INSERT INTO IDENTITIES "
+            "(environment_id, identifier, identity_key, traits, is_deleted) VALUES",
+            [(environment_api_key_str, "aaron", "aaron_key", {"foo": "bar"}, True)],  # type: ignore[list-item]
+        )
+
+    # When
+    with connections["clickhouse"].cursor() as cursor:
+        counts = compute_segment_counts_for_project(project, cursor)
+
+    # Then
+    assert len(counts) == 1
+    assert counts[0].segment_id == matching_segment.id
+    assert counts[0].count == 2
+
+
 @pytest.mark.clickhouse
 def test_get_segment_members_page__matching_identities__returns_members_ordered_by_identifier(
     segment_membership_identities: None,

docs/docs/deployment-self-hosting/observability/_events-catalogue.md

@@ -379,7 +379,7 @@ Attributes:
 ### `segment_membership.members.segment.skipped`
 
 Logged at `error` from:
- - `api/segment_membership/services.py:182`
+ - `api/segment_membership/services.py:183`
 
 Attributes:
  - `reason`