feat(cohorts): Backwards compatibility of groups and properties

ee/clickhouse/models/cohort.py

@@ -1,6 +1,6 @@
 import uuid
 from datetime import datetime, timedelta
-from typing import Any, Dict, List, Optional, Tuple, Union
+from typing import Any, Dict, List, Optional, Set, Tuple, Union
 
 import structlog
 from dateutil import parser
@@ -38,38 +38,57 @@
 
 
 def format_person_query(
-    cohort: Cohort, index: int, *, custom_match_field: str = "person_id"
+    cohort: Cohort,
+    index: int,
+    *,
+    custom_match_field: str = "person_id",
+    cohorts_seen: Optional[Set[int]] = None,
+    using_new_query: bool = False,
 ) -> Tuple[str, Dict[str, Any]]:
-    filters = []
-    params: Dict[str, Any] = {}
-
     if cohort.is_static:
         return format_static_cohort_query(cohort.pk, index, prepend="", custom_match_field=custom_match_field)
 
-    or_queries = []
-    groups = cohort.groups
+    if using_new_query:
+        if not cohort.properties.values:
+            # No person can match an empty cohort
+            return "0 = 19", {}
 
-    if not groups:
-        # No person can match a cohort that has no match groups
-        return "0 = 19", {}
+        from ee.clickhouse.queries.cohort_query import CohortQuery
 
-    for group_idx, group in enumerate(groups):
-        if group.get("action_id") or group.get("event_id"):
-            entity_query, entity_params = get_entity_cohort_subquery(cohort, group, group_idx)
-            params = {**params, **entity_params}
-            filters.append(entity_query)
+        query, params = CohortQuery(
+            Filter(data={"properties": cohort.properties}), cohort.team, cohort_pk=cohort.pk, cohorts_seen=cohorts_seen
+        ).get_query()
 
-        elif group.get("properties"):
-            prop_query, prop_params = get_properties_cohort_subquery(cohort, group, group_idx)
-            or_queries.append(prop_query)
-            params = {**params, **prop_params}
+        return f"{custom_match_field} IN ({query})", params
+
+    else:
+        filters = []
+        params = {}
+
+        or_queries = []
+        groups = cohort.groups
+
+        if not groups:
+            # No person can match a cohort that has no match groups
+            return "0 = 19", {}
+
+        for group_idx, group in enumerate(groups):
+            if group.get("action_id") or group.get("event_id"):
+                entity_query, entity_params = get_entity_cohort_subquery(cohort, group, group_idx)
+                params = {**params, **entity_params}
+                filters.append(entity_query)
+
+            elif group.get("properties"):
+                prop_query, prop_params = get_properties_cohort_subquery(cohort, group, group_idx)
+                or_queries.append(prop_query)
+                params = {**params, **prop_params}
 
-    if len(or_queries) > 0:
-        query = "AND ({})".format(" OR ".join(or_queries))
-        filters.append("{} IN {}".format(custom_match_field, GET_LATEST_PERSON_ID_SQL.format(query=query)))
+        if len(or_queries) > 0:
+            query = "AND ({})".format(" OR ".join(or_queries))
+            filters.append("{} IN {}".format(custom_match_field, GET_LATEST_PERSON_ID_SQL.format(query=query)))
 
-    joined_filter = " OR ".join(filters)
-    return joined_filter, params
+        joined_filter = " OR ".join(filters)
+        return joined_filter, params
 
 
 def format_static_cohort_query(
@@ -239,8 +258,16 @@ def is_precalculated_query(cohort: Cohort) -> bool:
         return False
 
 
-def format_filter_query(cohort: Cohort, index: int = 0, id_column: str = "distinct_id") -> Tuple[str, Dict[str, Any]]:
-    person_query, params = format_cohort_subquery(cohort, index)
+def format_filter_query(
+    cohort: Cohort,
+    index: int = 0,
+    id_column: str = "distinct_id",
+    cohorts_seen: Optional[Set[int]] = None,
+    using_new_query: bool = False,
+) -> Tuple[str, Dict[str, Any]]:
+    person_query, params = format_cohort_subquery(
+        cohort, index, cohorts_seen=cohorts_seen, using_new_query=using_new_query
+    )
 
     person_id_query = CALCULATE_COHORT_PEOPLE_SQL.format(
         query=person_query,
@@ -250,12 +277,24 @@ def format_filter_query(cohort: Cohort, index: int = 0, id_column: str = "distin
     return person_id_query, params
 
 
-def format_cohort_subquery(cohort: Cohort, index: int, custom_match_field="person_id") -> Tuple[str, Dict[str, Any]]:
+def format_cohort_subquery(
+    cohort: Cohort,
+    index: int,
+    custom_match_field="person_id",
+    cohorts_seen: Optional[Set[int]] = None,
+    using_new_query: bool = False,
+) -> Tuple[str, Dict[str, Any]]:
     is_precalculated = is_precalculated_query(cohort)
     person_query, params = (
         format_precalculated_cohort_query(cohort.pk, index, custom_match_field=custom_match_field)
         if is_precalculated
-        else format_person_query(cohort, index, custom_match_field=custom_match_field)
+        else format_person_query(
+            cohort,
+            index,
+            custom_match_field=custom_match_field,
+            cohorts_seen=cohorts_seen,
+            using_new_query=using_new_query,
+        )
     )
     return person_query, params
 
@@ -296,6 +335,21 @@ def insert_static_cohort(person_uuids: List[Optional[uuid.UUID]], cohort_id: int
     sync_execute(INSERT_PERSON_STATIC_COHORT, persons)
 
 
+def recalculate_cohortpeople_with_new_query(cohort: Cohort) -> Optional[int]:
+    cohort_filter, cohort_params = format_person_query(cohort, 0, custom_match_field="id", using_new_query=True)
+
+    count = sync_execute(
+        f"""
+        SELECT COUNT(1)
+        FROM person
+        WHERE {cohort_filter}
+        """,
+        {**cohort_params, "team_id": cohort.team_id, "cohort_id": cohort.pk},
+    )[0][0]
+
+    return count
+
+
 def recalculate_cohortpeople(cohort: Cohort) -> Optional[int]:
     cohort_filter, cohort_params = format_person_query(cohort, 0, custom_match_field="id")
 

ee/clickhouse/models/test/__snapshots__/test_cohort.ambr

@@ -31,16 +31,30 @@
 # name: TestCohort.test_cohortpeople_action_count.10
   '
 
-  SELECT count(*)
-  FROM
-    (SELECT 1
-     FROM cohortpeople
-     WHERE team_id = 2
-       AND cohort_id = 2
-     GROUP BY person_id,
-              cohort_id,
-              team_id
-     HAVING sum(sign) > 0)
+  SELECT COUNT(1)
+  FROM person
+  WHERE id IN
+      (SELECT behavior_query.person_id AS id
+       FROM
+         (SELECT pdi.person_id AS person_id,
+                 countIf(timestamp > now() - INTERVAL 3 day
+                         AND timestamp < now()
+                         AND ((event = '$pageview'))) = 1 AS performed_event_multiple_condition_level_level_level_0
+          FROM events e
+          INNER JOIN
+            (SELECT distinct_id,
+                    argMax(person_id, version) as person_id
+             FROM person_distinct_id2
+             WHERE team_id = 2
+             GROUP BY distinct_id
+             HAVING argMax(is_deleted, version) = 0) AS pdi ON e.distinct_id = pdi.distinct_id
+          WHERE team_id = 2
+            AND event IN ['$pageview']
+            AND timestamp <= now()
+            AND timestamp >= now() - INTERVAL 3 day
+          GROUP BY person_id) behavior_query
+       WHERE 1 = 1
+         AND (((performed_event_multiple_condition_level_level_level_0))) )
   '
 ---
 # name: TestCohort.test_cohortpeople_action_count.11
@@ -50,14 +64,7 @@
   where cohort_id = 2
   '
 ---
-# name: TestCohort.test_cohortpeople_action_count.2
-  '
-  SELECT person_id
-  FROM cohortpeople
-  where cohort_id = 2
-  '
----
-# name: TestCohort.test_cohortpeople_action_count.3
+# name: TestCohort.test_cohortpeople_action_count.12
   '
 
   SELECT count(*)
@@ -72,7 +79,7 @@
      HAVING sum(sign) > 0)
   '
 ---
-# name: TestCohort.test_cohortpeople_action_count.4
+# name: TestCohort.test_cohortpeople_action_count.13
   '
 
   SELECT count(*)
@@ -87,14 +94,79 @@
      HAVING sum(sign) > 0)
   '
 ---
-# name: TestCohort.test_cohortpeople_action_count.5
+# name: TestCohort.test_cohortpeople_action_count.14
+  '
+
+  SELECT COUNT(1)
+  FROM person
+  WHERE id IN
+      (SELECT behavior_query.person_id AS id
+       FROM
+         (SELECT pdi.person_id AS person_id,
+                 countIf(timestamp > now() - INTERVAL 3 day
+                         AND timestamp < now()
+                         AND ((event = '$pageview'))) > 0 AS performed_event_condition_level_level_level_0
+          FROM events e
+          INNER JOIN
+            (SELECT distinct_id,
+                    argMax(person_id, version) as person_id
+             FROM person_distinct_id2
+             WHERE team_id = 2
+             GROUP BY distinct_id
+             HAVING argMax(is_deleted, version) = 0) AS pdi ON e.distinct_id = pdi.distinct_id
+          WHERE team_id = 2
+            AND event IN ['$pageview']
+            AND timestamp <= now()
+            AND timestamp >= now() - INTERVAL 3 day
+          GROUP BY person_id) behavior_query
+       WHERE 1 = 1
+         AND (((performed_event_condition_level_level_level_0))) )
+  '
+---
+# name: TestCohort.test_cohortpeople_action_count.15
   '
   SELECT person_id
   FROM cohortpeople
   where cohort_id = 2
   '
 ---
-# name: TestCohort.test_cohortpeople_action_count.6
+# name: TestCohort.test_cohortpeople_action_count.2
+  '
+
+  SELECT COUNT(1)
+  FROM person
+  WHERE id IN
+      (SELECT behavior_query.person_id AS id
+       FROM
+         (SELECT pdi.person_id AS person_id,
+                 countIf(timestamp > now() - INTERVAL 3 day
+                         AND timestamp < now()
+                         AND ((event = '$pageview'))) >= 2 AS performed_event_multiple_condition_level_level_level_0
+          FROM events e
+          INNER JOIN
+            (SELECT distinct_id,
+                    argMax(person_id, version) as person_id
+             FROM person_distinct_id2
+             WHERE team_id = 2
+             GROUP BY distinct_id
+             HAVING argMax(is_deleted, version) = 0) AS pdi ON e.distinct_id = pdi.distinct_id
+          WHERE team_id = 2
+            AND event IN ['$pageview']
+            AND timestamp <= now()
+            AND timestamp >= now() - INTERVAL 3 day
+          GROUP BY person_id) behavior_query
+       WHERE 1 = 1
+         AND (((performed_event_multiple_condition_level_level_level_0))) )
+  '
+---
+# name: TestCohort.test_cohortpeople_action_count.3
+  '
+  SELECT person_id
+  FROM cohortpeople
+  where cohort_id = 2
+  '
+---
+# name: TestCohort.test_cohortpeople_action_count.4
   '
 
   SELECT count(*)
@@ -109,7 +181,7 @@
      HAVING sum(sign) > 0)
   '
 ---
-# name: TestCohort.test_cohortpeople_action_count.7
+# name: TestCohort.test_cohortpeople_action_count.5
   '
 
   SELECT count(*)
@@ -124,13 +196,57 @@
      HAVING sum(sign) > 0)
   '
 ---
-# name: TestCohort.test_cohortpeople_action_count.8
+# name: TestCohort.test_cohortpeople_action_count.6
+  '
+
+  SELECT COUNT(1)
+  FROM person
+  WHERE id IN
+      (SELECT behavior_query.person_id AS id
+       FROM
+         (SELECT pdi.person_id AS person_id,
+                 countIf(timestamp > now() - INTERVAL 3 day
+                         AND timestamp < now()
+                         AND ((event = '$pageview'))) <= 1 AS performed_event_multiple_condition_level_level_level_0
+          FROM events e
+          INNER JOIN
+            (SELECT distinct_id,
+                    argMax(person_id, version) as person_id
+             FROM person_distinct_id2
+             WHERE team_id = 2
+             GROUP BY distinct_id
+             HAVING argMax(is_deleted, version) = 0) AS pdi ON e.distinct_id = pdi.distinct_id
+          WHERE team_id = 2
+            AND event IN ['$pageview']
+            AND timestamp <= now()
+            AND timestamp >= now() - INTERVAL 3 day
+          GROUP BY person_id) behavior_query
+       WHERE 1 = 1
+         AND (((performed_event_multiple_condition_level_level_level_0))) )
+  '
+---
+# name: TestCohort.test_cohortpeople_action_count.7
   '
   SELECT person_id
   FROM cohortpeople
   where cohort_id = 2
   '
 ---
+# name: TestCohort.test_cohortpeople_action_count.8
+  '
+
+  SELECT count(*)
+  FROM
+    (SELECT 1
+     FROM cohortpeople
+     WHERE team_id = 2
+       AND cohort_id = 2
+     GROUP BY person_id,
+              cohort_id,
+              team_id
+     HAVING sum(sign) > 0)
+  '
+---
 # name: TestCohort.test_cohortpeople_action_count.9
   '
 
@@ -163,4 +279,4 @@
        WHERE cohort_id = %(_cohort_id_0)s
          AND team_id = %(team_id)s)
   '
----
+---