fix: smaller cohort resource usage for replay filters (#25699)

Co-authored-by: github-actions <41898282+github-actions[bot]@users.noreply.github.com>

posthog/session_recordings/queries/session_recording_list_from_filters.py

@@ -2,6 +2,8 @@
 from typing import Any, NamedTuple, cast, Optional, Union
 from datetime import datetime, timedelta
 
+import posthoganalytics
+
 from posthog.hogql import ast
 from posthog.hogql.ast import CompareOperation
 from posthog.hogql.parser import parse_select
@@ -33,12 +35,22 @@ def is_group_property(p: Property) -> bool:
     return p.type == "group"
 
 
+def is_cohort_property(p: Property) -> bool:
+    return "cohort" in p.type
+
+
 class SessionRecordingQueryResult(NamedTuple):
     results: list
     has_more_recording: bool
     timings: list[QueryTiming] | None = None
 
 
+class UnexpectedQueryProperties(Exception):
+    def __init__(self, remaining_properties: PropertyGroup | None):
+        self.remaining_properties = remaining_properties
+        super().__init__(f"Unexpected properties in query: {remaining_properties}")
+
+
 class SessionRecordingListFromFilters:
     SESSION_RECORDINGS_DEFAULT_LIMIT = 50
 
@@ -224,11 +236,19 @@ def _where_predicates(self) -> Union[ast.And, ast.Or]:
                 )
             )
 
-        remaining_properties = self._strip_person_and_event_properties(self._filter.property_groups)
-        if remaining_properties:
-            logger.info(
-                "session_replay_query_builder has unhandled properties", unhandled_properties=remaining_properties
+        cohort_subquery = CohortPropertyGroupsSubQuery(self._team, self._filter, self.ttl_days).get_query()
+        if cohort_subquery:
+            optional_exprs.append(
+                ast.CompareOperation(
+                    op=ast.CompareOperationOp.In,
+                    left=ast.Field(chain=["s", "distinct_id"]),
+                    right=cohort_subquery,
+                )
             )
+
+        remaining_properties = self._strip_person_and_event_and_cohort_properties(self._filter.property_groups)
+        if remaining_properties:
+            posthoganalytics.capture_exception(UnexpectedQueryProperties(remaining_properties))
             optional_exprs.append(property_to_expr(remaining_properties, team=self._team, scope="replay"))
 
         if self._filter.console_log_filters.values:
@@ -267,11 +287,14 @@ def _where_predicates(self) -> Union[ast.And, ast.Or]:
     def _having_predicates(self) -> ast.Expr:
         return property_to_expr(self._filter.having_predicates, team=self._team, scope="replay")
 
-    def _strip_person_and_event_properties(self, property_group: PropertyGroup) -> PropertyGroup | None:
+    def _strip_person_and_event_and_cohort_properties(self, property_group: PropertyGroup) -> PropertyGroup | None:
         property_groups_to_keep = [
             g
             for g in property_group.flat
-            if not is_event_property(g) and not is_person_property(g) and not is_group_property(g)
+            if not is_event_property(g)
+            and not is_person_property(g)
+            and not is_group_property(g)
+            and not is_cohort_property(g)
         ]
 
         return (
@@ -334,6 +357,44 @@ def _where_predicates(self) -> ast.Expr:
         )
 
 
+class CohortPropertyGroupsSubQuery:
+    _team: Team
+    _filter: SessionRecordingsFilter
+    _ttl_days: int
+
+    raw_cohort_to_distinct_id = """
+    select distinct_id
+    from person_distinct_ids
+    where {cohort_predicate}
+    """
+
+    def __init__(self, team: Team, filter: SessionRecordingsFilter, ttl_days: int):
+        self._team = team
+        self._filter = filter
+        self._ttl_days = ttl_days
+
+    def get_query(self) -> ast.SelectQuery | ast.SelectUnionQuery | None:
+        if self.cohort_properties:
+            return parse_select(
+                self.raw_cohort_to_distinct_id,
+                {"cohort_predicate": property_to_expr(self.cohort_properties, team=self._team, scope="replay")},
+            )
+
+        return None
+
+    @cached_property
+    def cohort_properties(self) -> PropertyGroup | None:
+        cohort_property_groups = [g for g in self._filter.property_groups.flat if is_cohort_property(g)]
+        return (
+            PropertyGroup(
+                type=self._filter.property_operand,
+                values=cohort_property_groups,
+            )
+            if cohort_property_groups
+            else None
+        )
+
+
 class PersonsIdCompareOperation:
     _team: Team
     _filter: SessionRecordingsFilter