hogql has special syntax for cohorts

PostHog · Oct 20, 2024 · d92924a · d92924a
1 parent accaf4e
commit d92924a
Show file tree

Hide file tree

Showing 3 changed files with 25 additions and 103 deletions.
diff --git a/posthog/session_recordings/queries/session_recording_list_from_filters.py b/posthog/session_recordings/queries/session_recording_list_from_filters.py
@@ -43,6 +43,10 @@ def is_static_cohort_property(p: Property) -> bool:
     return p.type == "precalculated-cohort"
 
 
+def is_cohort_property(p: Property) -> bool:
+    return is_dynamic_cohort_property(p) or is_static_cohort_property(p)
+
+
 class SessionRecordingQueryResult(NamedTuple):
     results: list
     has_more_recording: bool
@@ -238,27 +242,13 @@ def _where_predicates(self) -> Union[ast.And, ast.Or]:
                 )
             )
 
-        dynamic_cohort_subquery = DynamicCohortPropertyGroupsSubQuery(
-            self._team, self._filter, self.ttl_days
-        ).get_queries()
-        if dynamic_cohort_subquery:
+        cohort_subquery = CohortPropertyGroupsSubQuery(self._team, self._filter, self.ttl_days).get_queries()
+        if cohort_subquery:
             optional_exprs.append(
                 ast.CompareOperation(
                     op=ast.CompareOperationOp.In,
                     left=ast.Field(chain=["s", "distinct_id"]),
-                    right=dynamic_cohort_subquery,
-                )
-            )
-
-        static_cohort_subquery = StaticCohortPropertyGroupsSubQuery(
-            self._team, self._filter, self.ttl_days
-        ).get_queries()
-        if static_cohort_subquery:
-            optional_exprs.append(
-                ast.CompareOperation(
-                    op=ast.CompareOperationOp.In,
-                    left=ast.Field(chain=["s", "distinct_id"]),
-                    right=static_cohort_subquery,
+                    right=cohort_subquery,
                 )
             )
 
@@ -374,63 +364,15 @@ def _where_predicates(self) -> ast.Expr:
         )
 
 
-class StaticCohortPropertyGroupsSubQuery:
-    _team: Team
-    _filter: SessionRecordingsFilter
-    _ttl_days: int
-
-    raw_cohort_to_distinct_id = """
-    select distinct_id
-    from person_distinct_ids
-    where person_id IN (
-        SELECT person_id
-        FROM static_cohort_people
-        WHERE
-        {where_predicates}
-    )
-    """
-
-    def __init__(self, team: Team, filter: SessionRecordingsFilter, ttl_days: int):
-        self._team = team
-        self._filter = filter
-        self._ttl_days = ttl_days
-
-    def get_queries(self) -> ast.SelectQuery | ast.SelectUnionQuery | None:
-        if self.cohort_properties:
-            return parse_select(
-                self.raw_cohort_to_distinct_id,
-                {"where_predicates": 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_static_cohort_property(g)]
-        return (
-            PropertyGroup(
-                type=self._filter.property_operand,
-                values=cohort_property_groups,
-            )
-            if cohort_property_groups
-            else None
-        )
-
-
-class DynamicCohortPropertyGroupsSubQuery:
+class CohortPropertyGroupsSubQuery:
     _team: Team
     _filter: SessionRecordingsFilter
     _ttl_days: int
 
     raw_cohort_to_distinct_id = """
     select distinct_id
     from person_distinct_ids
-    where person_id IN (
-        SELECT person_id
-        FROM cohort_people
-        WHERE
-        {where_predicates}
-    )
+    where {cohort_predicate}
     """
 
     def __init__(self, team: Team, filter: SessionRecordingsFilter, ttl_days: int):
@@ -442,14 +384,14 @@ def get_queries(self) -> ast.SelectQuery | ast.SelectUnionQuery | None:
         if self.cohort_properties:
             return parse_select(
                 self.raw_cohort_to_distinct_id,
-                {"where_predicates": property_to_expr(self.cohort_properties, team=self._team, scope="replay")},
+                {"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_dynamic_cohort_property(g)]
+        cohort_property_groups = [g for g in self._filter.property_groups.flat if is_cohort_property(g)]
         return (
             PropertyGroup(
                 type=self._filter.property_operand,