feat: cohort tests

docker/clickhouse/config.xml

@@ -220,7 +220,7 @@
         <!-- Interval of flushing data. -->
         <flush_interval_milliseconds>7500</flush_interval_milliseconds>
         <!-- Maximal size in lines for the logs. When non-flushed logs amount reaches max_size, logs dumped to the disk. -->
-        <max_size_rows>1048576</max_size_rows>
+        <max_size_rows>8388608</max_size_rows>
         <!-- Pre-allocated size in lines for the logs. -->
         <reserved_size_rows>8192</reserved_size_rows>
         <!-- Lines amount threshold, reaching it launches flushing logs to the disk in background. -->
@@ -239,7 +239,7 @@
         <table>query_thread_log</table>
         <partition_by>toYYYYMM(event_date)</partition_by>
         <flush_interval_milliseconds>7500</flush_interval_milliseconds>
-        <max_size_rows>1048576</max_size_rows>
+        <max_size_rows>8388608</max_size_rows>
         <reserved_size_rows>8192</reserved_size_rows>
         <buffer_size_rows_flush_threshold>524288</buffer_size_rows_flush_threshold>
         <flush_on_crash>false</flush_on_crash>

ee/clickhouse/models/test/test_cohort.py

@@ -7,7 +7,7 @@
 from posthog.client import sync_execute
 from posthog.hogql.hogql import HogQLContext
 from posthog.models.action import Action
-from posthog.models.cohort import Cohort
+from posthog.models.cohort import Cohort, get_and_update_pending_version
 from posthog.models.cohort.sql import GET_COHORTPEOPLE_BY_COHORT_ID
 from posthog.models.cohort.util import format_filter_query
 from posthog.models.filters import Filter
@@ -68,6 +68,21 @@ def get_person_ids_by_cohort_id(
     return [str(row[0]) for row in results]
 
 
+def calculate_cohort_hogql_test_harness(cohort: Cohort, pending_version: int):
+    version = pending_version * 2 + 2
+    cohort.calculate_people_ch(version)
+    query = f"""
+        SELECT count() FROM
+        (SELECT person_id FROM cohortpeople as cp WHERE cp.version = {version} and cp.cohort_id = {cohort.pk}) as cp1
+        FULL OUTER JOIN (SELECT person_id FROM cohortpeople as cp WHERE cp.version = {version-1} and cp.cohort_id = {cohort.pk}) as cp2
+        ON cp1.person_id = cp2.person_id
+        WHERE empty(cp1.person_id) or empty(cp2.person_id)
+    """
+    result = sync_execute(query)
+    assert 0 == result[0][0]
+    return version
+
+
 class TestCohort(ClickhouseTestMixin, BaseTest):
     def _get_cohortpeople(self, cohort: Cohort, *, team_id: Optional[int] = None):
         team_id = team_id or cohort.team_id
@@ -517,7 +532,7 @@ def test_insert_by_distinct_id_or_email(self):
         self.assertEqual(results, 3)
 
         #  If we accidentally call calculate_people it shouldn't erase people
-        cohort.calculate_people_ch(pending_version=0)
+        calculate_cohort_hogql_test_harness(cohort, 0)
         results = get_person_ids_by_cohort_id(self.team.pk, cohort.id)
         self.assertEqual(len(results), 3)
 
@@ -556,7 +571,7 @@ def test_cohortpeople_basic(self):
             name="cohort1",
         )
 
-        cohort1.calculate_people_ch(pending_version=0)
+        calculate_cohort_hogql_test_harness(cohort1, 0)
 
         results = self._get_cohortpeople(cohort1)
         self.assertEqual(len(results), 2)
@@ -592,13 +607,13 @@ def test_cohortpeople_action_basic(self):
         )
 
         cohort1 = Cohort.objects.create(team=self.team, groups=[{"action_id": action.pk, "days": 1}], name="cohort1")
-        cohort1.calculate_people_ch(pending_version=0)
+        calculate_cohort_hogql_test_harness(cohort1, 0)
 
         results = self._get_cohortpeople(cohort1)
         self.assertEqual(len(results), 2)
 
         cohort2 = Cohort.objects.create(team=self.team, groups=[{"action_id": action.pk, "days": 1}], name="cohort2")
-        cohort2.calculate_people_ch(pending_version=0)
+        calculate_cohort_hogql_test_harness(cohort2, 0)
 
         results = self._get_cohortpeople(cohort2)
         self.assertEqual(len(results), 2)
@@ -684,7 +699,7 @@ def test_cohortpeople_action_count(self):
             groups=[{"action_id": action.pk, "days": 3, "count": 2, "count_operator": "gte"}],
             name="cohort1",
         )
-        cohort1.calculate_people_ch(pending_version=0)
+        calculate_cohort_hogql_test_harness(cohort1, 0)
 
         results = self._get_cohortpeople(cohort1)
         self.assertEqual(len(results), 2)
@@ -694,7 +709,7 @@ def test_cohortpeople_action_count(self):
             groups=[{"action_id": action.pk, "days": 3, "count": 1, "count_operator": "lte"}],
             name="cohort2",
         )
-        cohort2.calculate_people_ch(pending_version=0)
+        calculate_cohort_hogql_test_harness(cohort2, 0)
 
         results = self._get_cohortpeople(cohort2)
         self.assertEqual(len(results), 1)
@@ -704,7 +719,7 @@ def test_cohortpeople_action_count(self):
             groups=[{"action_id": action.pk, "days": 3, "count": 1, "count_operator": "eq"}],
             name="cohort3",
         )
-        cohort3.calculate_people_ch(pending_version=0)
+        calculate_cohort_hogql_test_harness(cohort3, 0)
 
         results = self._get_cohortpeople(cohort3)
         self.assertEqual(len(results), 1)
@@ -738,9 +753,9 @@ def test_cohortpeople_deleted_person(self):
             name="cohort1",
         )
 
-        cohort1.calculate_people_ch(pending_version=0)
+        calculate_cohort_hogql_test_harness(cohort1, 0)
         p2.delete()
-        cohort1.calculate_people_ch(pending_version=0)
+        calculate_cohort_hogql_test_harness(cohort1, 0)
 
     def test_cohortpeople_prop_changed(self):
         with freeze_time((datetime.now() - timedelta(days=3)).strftime("%Y-%m-%d")):
@@ -776,14 +791,14 @@ def test_cohortpeople_prop_changed(self):
                 name="cohort1",
             )
 
-        cohort1.calculate_people_ch(pending_version=0)
+        calculate_cohort_hogql_test_harness(cohort1, 0)
 
         with freeze_time((datetime.now() - timedelta(days=2)).strftime("%Y-%m-%d")):
             p2.version = 1
-            p2.properties = ({"$some_prop": "another", "$another_prop": "another"},)
+            p2.properties = {"$some_prop": "another", "$another_prop": "another"}
             p2.save()
 
-        cohort1.calculate_people_ch(pending_version=1)
+        calculate_cohort_hogql_test_harness(cohort1, 1)
 
         results = self._get_cohortpeople(cohort1)
 
@@ -818,7 +833,7 @@ def test_cohort_change(self):
             ],
             name="cohort1",
         )
-        cohort1.calculate_people_ch(pending_version=0)
+        calculate_cohort_hogql_test_harness(cohort1, 0)
         results = self._get_cohortpeople(cohort1)
 
         self.assertEqual(len(results), 1)
@@ -834,7 +849,7 @@ def test_cohort_change(self):
         ]
         cohort1.save()
 
-        cohort1.calculate_people_ch(pending_version=1)
+        calculate_cohort_hogql_test_harness(cohort1, 1)
 
         results = self._get_cohortpeople(cohort1)
         self.assertEqual(len(results), 1)
@@ -851,7 +866,7 @@ def test_static_cohort_precalculated(self):
         cohort = Cohort.objects.create(team=self.team, groups=[], is_static=True, last_calculation=timezone.now())
         cohort.insert_users_by_list(["1", "123"])
 
-        cohort.calculate_people_ch(pending_version=0)
+        calculate_cohort_hogql_test_harness(cohort, 0)
 
         with self.settings(USE_PRECALCULATED_CH_COHORT_PEOPLE=True):
             sql, _ = format_filter_query(cohort, 0, HogQLContext(team_id=self.team.pk))
@@ -866,15 +881,15 @@ def test_cohortpeople_with_valid_other_cohort_filter(self):
             groups=[{"properties": [{"key": "foo", "value": "bar", "type": "person"}]}],
             name="cohort0",
         )
-        cohort0.calculate_people_ch(pending_version=0)
+        calculate_cohort_hogql_test_harness(cohort0, 0)
 
         cohort1: Cohort = Cohort.objects.create(
             team=self.team,
             groups=[{"properties": [{"key": "id", "type": "cohort", "value": cohort0.id}]}],
             name="cohort1",
         )
 
-        cohort1.calculate_people_ch(pending_version=0)
+        calculate_cohort_hogql_test_harness(cohort1, 0)
 
         res = self._get_cohortpeople(cohort1)
         self.assertEqual(len(res), 1)
@@ -914,7 +929,7 @@ def test_cohortpeople_with_not_in_cohort_operator(self):
             groups=[{"properties": [{"key": "$some_prop", "value": "something1", "type": "person"}]}],
             name="cohort0",
         )
-        cohort0.calculate_people_ch(pending_version=0)
+        calculate_cohort_hogql_test_harness(cohort0, 0)
 
         cohort1 = Cohort.objects.create(
             team=self.team,
@@ -943,7 +958,7 @@ def test_cohortpeople_with_not_in_cohort_operator(self):
             name="cohort1",
         )
 
-        cohort1.calculate_people_ch(pending_version=0)
+        calculate_cohort_hogql_test_harness(cohort1, 0)
 
         with self.settings(USE_PRECALCULATED_CH_COHORT_PEOPLE=True):
             filter = Filter(
@@ -1110,7 +1125,7 @@ def test_cohortpeople_with_not_in_cohort_operator_for_behavioural_cohorts(self):
             ],
             name="cohort0",
         )
-        cohort0.calculate_people_ch(pending_version=0)
+        calculate_cohort_hogql_test_harness(cohort0, 0)
 
         cohort1 = Cohort.objects.create(
             team=self.team,
@@ -1139,7 +1154,7 @@ def test_cohortpeople_with_not_in_cohort_operator_for_behavioural_cohorts(self):
             name="cohort1",
         )
 
-        cohort1.calculate_people_ch(pending_version=0)
+        calculate_cohort_hogql_test_harness(cohort1, 0)
 
         with self.settings(USE_PRECALCULATED_CH_COHORT_PEOPLE=True):
             filter = Filter(
@@ -1171,7 +1186,7 @@ def test_cohortpeople_with_nonexistent_other_cohort_filter(self):
             name="cohort1",
         )
 
-        cohort1.calculate_people_ch(pending_version=0)
+        calculate_cohort_hogql_test_harness(cohort1, 0)
 
         res = self._get_cohortpeople(cohort1)
         self.assertEqual(len(res), 0)
@@ -1183,7 +1198,7 @@ def test_clickhouse_empty_query(self):
             name="cohort1",
         )
 
-        cohort2.calculate_people_ch(pending_version=0)
+        calculate_cohort_hogql_test_harness(cohort2, 0)
         self.assertFalse(Cohort.objects.get().is_calculating)
 
     def test_query_with_multiple_new_style_cohorts(self):
@@ -1322,7 +1337,7 @@ def test_query_with_multiple_new_style_cohorts(self):
             name="cohort1",
         )
 
-        cohort1.calculate_people_ch(pending_version=0)
+        calculate_cohort_hogql_test_harness(cohort1, 0)
 
         result = self._get_cohortpeople(cohort1)
         self.assertCountEqual([p1.uuid, p3.uuid], [r[0] for r in result])
@@ -1350,23 +1365,23 @@ def test_update_cohort(self):
             name="cohort1",
         )
 
-        cohort1.calculate_people_ch(pending_version=0)
+        calculate_cohort_hogql_test_harness(cohort1, 0)
 
         # Should only have p1 in this cohort
         results = self._get_cohortpeople(cohort1)
         self.assertEqual(len(results), 1)
 
         cohort1.groups = [{"properties": [{"key": "$another_prop", "value": "something", "type": "person"}]}]
         cohort1.save()
-        cohort1.calculate_people_ch(pending_version=1)
+        calculate_cohort_hogql_test_harness(cohort1, 1)
 
         # Should only have p2, p3 in this cohort
         results = self._get_cohortpeople(cohort1)
         self.assertEqual(len(results), 2)
 
         cohort1.groups = [{"properties": [{"key": "$some_prop", "value": "something", "type": "person"}]}]
         cohort1.save()
-        cohort1.calculate_people_ch(pending_version=2)
+        calculate_cohort_hogql_test_harness(cohort1, 2)
 
         # Should only have p1 again in this cohort
         results = self._get_cohortpeople(cohort1)
@@ -1396,10 +1411,10 @@ def test_cohort_versioning(self):
             name="cohort1",
         )
 
-        cohort1.calculate_people_ch(pending_version=5)
+        version = calculate_cohort_hogql_test_harness(cohort1, 5)
 
-        cohort1.pending_version = 5
-        cohort1.version = 5
+        cohort1.pending_version = version
+        cohort1.version = version
         cohort1.save()
 
         # Should have p1 in this cohort even if version is different
@@ -1439,11 +1454,23 @@ def test_calculate_people_ch_in_multiteam_project(self):
             name="shared cohort",
         )
         # Calculate cohort
-        shared_cohort.calculate_people_ch(pending_version=0)
+        calculate_cohort_hogql_test_harness(shared_cohort, 0)
 
         # Verify shared_cohort is now calculated for both teams
         results_team1 = self._get_cohortpeople(shared_cohort, team_id=self.team.pk)
         results_team2 = self._get_cohortpeople(shared_cohort, team_id=team2.pk)
 
         self.assertCountEqual([r[0] for r in results_team1], [person2_team1.uuid])
         self.assertCountEqual([r[0] for r in results_team2], [person1_team2.uuid])
+
+    def test_increment_cohort(self):
+        cohort1 = Cohort.objects.create(
+            team=self.team,
+            groups=[{"properties": [{"key": "$some_prop", "value": "something", "type": "person"}]}],
+            name="cohort1",
+            pending_version=3,
+        )
+        new_version = get_and_update_pending_version(cohort1)
+        assert new_version == 6
+        new_version = get_and_update_pending_version(cohort1)
+        assert new_version == 8

ee/clickhouse/queries/test/__snapshots__/test_cohort_query.ambr

@@ -542,6 +542,28 @@
            AND (has(['[email protected]'], replaceRegexpAll(JSONExtractRaw(argMax(person.properties, version), 'email'), '^"|"$', ''))))) SETTINGS optimize_aggregation_in_order = 1
   '''
 # ---
+# name: TestCohortQuery.test_person_props_only.1
+  '''
+
+  SELECT id
+  FROM person
+  WHERE team_id = 99999
+    AND id IN
+      (SELECT id
+       FROM person
+       WHERE team_id = 99999
+         AND (((has(['[email protected]'], replaceRegexpAll(JSONExtractRaw(properties, 'email'), '^"|"$', '')))
+               OR (has(['[email protected]'], replaceRegexpAll(JSONExtractRaw(properties, 'email'), '^"|"$', ''))))
+              OR ((has(['test3'], replaceRegexpAll(JSONExtractRaw(properties, 'name'), '^"|"$', '')))
+                  AND (has(['[email protected]'], replaceRegexpAll(JSONExtractRaw(properties, 'email'), '^"|"$', ''))))) )
+  GROUP BY id
+  HAVING max(is_deleted) = 0
+  AND (((has(['[email protected]'], replaceRegexpAll(JSONExtractRaw(argMax(person.properties, version), 'email'), '^"|"$', '')))
+        OR (has(['[email protected]'], replaceRegexpAll(JSONExtractRaw(argMax(person.properties, version), 'email'), '^"|"$', ''))))
+       OR ((has(['test3'], replaceRegexpAll(JSONExtractRaw(argMax(person.properties, version), 'name'), '^"|"$', '')))
+           AND (has(['[email protected]'], replaceRegexpAll(JSONExtractRaw(argMax(person.properties, version), 'email'), '^"|"$', ''))))) SETTINGS optimize_aggregation_in_order = 1
+  '''
+# ---
 # name: TestCohortQuery.test_precalculated_cohort_filter_with_extra_filters
   '''
 
@@ -580,6 +602,50 @@
        OR (has(['test2'], replaceRegexpAll(JSONExtractRaw(argMax(person.properties, version), 'name'), '^"|"$', '')))) SETTINGS optimize_aggregation_in_order = 1
   '''
 # ---
+# name: TestCohortQuery.test_precalculated_cohort_filter_with_extra_filters.3
+  '''
+  SELECT persons.id AS id
+  FROM
+    (SELECT argMax(replaceRegexpAll(nullIf(nullIf(JSONExtractRaw(person.properties, 'name'), ''), 'null'), '^"|"$', ''), person.version) AS properties___name,
+            person.id AS id
+     FROM person
+     WHERE and(equals(person.team_id, 99999), in(id,
+                                                   (SELECT where_optimization.id AS id
+                                                    FROM person AS where_optimization
+                                                    WHERE and(equals(where_optimization.team_id, 99999), or(ifNull(equals(replaceRegexpAll(nullIf(nullIf(JSONExtractRaw(where_optimization.properties, 'name'), ''), 'null'), '^"|"$', ''), 'test'), 0), ifNull(equals(replaceRegexpAll(nullIf(nullIf(JSONExtractRaw(where_optimization.properties, 'name'), ''), 'null'), '^"|"$', ''), 'test2'), 0))))))
+     GROUP BY person.id
+     HAVING and(ifNull(equals(argMax(person.is_deleted, person.version), 0), 0), ifNull(less(argMax(toTimeZone(person.created_at, 'UTC'), person.version), plus(now64(6, 'UTC'), toIntervalDay(1))), 0))) AS persons
+  WHERE or(ifNull(equals(persons.properties___name, 'test'), 0), ifNull(equals(persons.properties___name, 'test2'), 0))
+  ORDER BY persons.id ASC
+  LIMIT 100 SETTINGS optimize_aggregation_in_order=1,
+                     join_algorithm='auto',
+                     readonly=2,
+                     max_execution_time=60,
+                     allow_experimental_object_type=1,
+                     format_csv_allow_double_quotes=0,
+                     max_ast_elements=4000000,
+                     max_expanded_ast_elements=4000000,
+                     max_bytes_before_external_group_by=0
+  '''
+# ---
+# name: TestCohortQuery.test_precalculated_cohort_filter_with_extra_filters.4
+  '''
+
+  SELECT id
+  FROM person
+  WHERE team_id = 99999
+    AND id IN
+      (SELECT id
+       FROM person
+       WHERE team_id = 99999
+         AND ((((has(['test'], replaceRegexpAll(JSONExtractRaw(properties, 'name'), '^"|"$', '')))))
+              OR (has(['test2'], replaceRegexpAll(JSONExtractRaw(properties, 'name'), '^"|"$', '')))) )
+  GROUP BY id
+  HAVING max(is_deleted) = 0
+  AND ((((has(['test'], replaceRegexpAll(JSONExtractRaw(argMax(person.properties, version), 'name'), '^"|"$', '')))))
+       OR (has(['test2'], replaceRegexpAll(JSONExtractRaw(argMax(person.properties, version), 'name'), '^"|"$', '')))) SETTINGS optimize_aggregation_in_order = 1
+  '''
+# ---
 # name: TestCohortQuery.test_static_cohort_filter
   '''