test(retention): add test for retention breakdown with materialized property (#7505)

* test(retention): add test for retention breakdown with materialized property

* remove unused imports

* Join tables based on breakdowns properties

ee/clickhouse/queries/clickhouse_retention.py

@@ -101,7 +101,9 @@ class ClickhouseRetention(Retention):
         trunc_func = get_trunc_func_ch(period)
 
         returning_event_query_templated, returning_event_params = RetentionEventsQuery(
-            filter=filter, team_id=team.pk, event_query_type=RetentionQueryType.RETURNING,
+            filter=filter.with_data({"breakdowns": []}),  # Avoid pulling in breakdown values from reterning event query
+            team_id=team.pk,
+            event_query_type=RetentionQueryType.RETURNING,
         ).get_query()
 
         returning_event_query = substitute_params(returning_event_query_templated, returning_event_params)

ee/clickhouse/queries/column_optimizer.py

@@ -96,6 +96,11 @@ class ColumnOptimizer:
                     (self.filter.breakdown, self.filter.breakdown_type, self.filter.breakdown_group_type_index)
                 ] += 1
 
+            # If we have a breakdowns attribute then make sure we pull in everything we
+            # need to calculate it
+            for breakdown in self.filter.breakdowns or []:
+                counter[(breakdown["property"], breakdown["type"], self.filter.breakdown_group_type_index)] += 1
+
         # Both entities and funnel exclusions can contain nested property filters
         for entity in self.filter.entities + cast(List[Entity], self.filter.exclusions):
             counter += extract_tables_and_properties(entity.properties)

ee/clickhouse/queries/retention/retention_event_query.py

@@ -25,14 +25,7 @@ class RetentionEventsQuery(ClickhouseEventQuery):
     def __init__(self, filter: RetentionFilter, event_query_type: RetentionQueryType, *args, **kwargs):
         self._event_query_type = event_query_type
         super().__init__(
-            filter=filter,
-            *args,
-            extra_person_fields=(
-                ["person_props"]
-                if event_query_type != RetentionQueryType.RETURNING and filter.breakdown_type == "person"
-                else []
-            ),
-            **kwargs,
+            filter=filter, *args, **kwargs,
         )
 
         self._trunc_func = get_trunc_func_ch(self._filter.period)
@@ -53,29 +46,30 @@ class RetentionEventsQuery(ClickhouseEventQuery):
             ),
         ]
 
-        if self._event_query_type != RetentionQueryType.RETURNING:
-            if self._filter.breakdowns and self._filter.breakdown_type:
-                # NOTE: `get_single_or_multi_property_string_expr` doesn't
-                # support breakdowns with different types e.g. a person property
-                # then an event property, so for now we just take the type of
-                # the self._filter.breakdown_type.
-                # TODO: update 'get_single_or_multi_property_string_expr` to take
-                # `Breakdown` type
-                breakdown_type = self._filter.breakdown_type
-                table = "events"
+        if (
+            self._event_query_type != RetentionQueryType.RETURNING
+            and self._filter.breakdowns
+            and self._filter.breakdown_type
+        ):
+            # NOTE: `get_single_or_multi_property_string_expr` doesn't
+            # support breakdowns with different types e.g. a person property
+            # then an event property, so for now we just take the type of
+            # the self._filter.breakdown_type.
+            # TODO: update 'get_single_or_multi_property_string_expr` to take
+            # `Breakdown` type
+            breakdown_type = self._filter.breakdown_type
+            table = "events"
 
-                if breakdown_type == "person":
-                    table = "person"
+            if breakdown_type == "person":
+                table = "person"
 
-                breakdown_values_expression = get_single_or_multi_property_string_expr(
-                    breakdown=[breakdown["property"] for breakdown in self._filter.breakdowns],
-                    table=table,
-                    query_alias=None,
-                )
+            breakdown_values_expression = get_single_or_multi_property_string_expr(
+                breakdown=[breakdown["property"] for breakdown in self._filter.breakdowns],
+                table=table,
+                query_alias=None,
+            )
 
-                _fields += [
-                    f"argMin({breakdown_values_expression}, {self._trunc_func}(e.timestamp)) AS breakdown_values"
-                ]
+            _fields += [f"argMin({breakdown_values_expression}, {self._trunc_func}(e.timestamp)) AS breakdown_values"]
 
         date_query, date_params = self._get_date_filter()
         self.params.update(date_params)

ee/clickhouse/views/test/test_clickhouse_retention.py

@@ -1,5 +1,3 @@
-import json
-import numbers
 from dataclasses import asdict, dataclass
 from typing import List, Literal, Optional, TypedDict, Union
 
@@ -7,14 +5,16 @@ from django.test import TestCase
 from django.test.client import Client
 
 from ee.clickhouse.test.test_journeys import _create_all_events, update_or_create_person
+from ee.clickhouse.util import ClickhouseTestMixin
 from ee.clickhouse.views.test.funnel.util import EventPattern
 from posthog.api.test.test_organization import create_organization
 from posthog.api.test.test_team import create_team
 from posthog.api.test.test_user import create_user
+from posthog.test.base import test_with_materialized_columns
 from posthog.utils import encode_get_request_params
 
 
-class RetentionTests(TestCase):
+class RetentionTests(TestCase, ClickhouseTestMixin):
     def test_can_get_retention_cohort_breakdown(self):
         organization = create_organization(name="test")
         team = create_team(organization=organization)
@@ -59,6 +59,7 @@ class RetentionTests(TestCase):
             },
         )
 
+    @test_with_materialized_columns(person_properties=["os"])
     def test_can_specify_breakdown_person_property(self):
         """
         By default, we group users together by the first time they perform the
@@ -118,6 +119,7 @@ class RetentionTests(TestCase):
             },
         )
 
+    @test_with_materialized_columns(event_properties=["os"])
     def test_can_specify_breakdown_event_property(self):
         """
         By default, we group users together by the first time they perform the