feat(experiments HogQL rewrite): prepare funnel query (#25096)

posthog/hogql_queries/experiment_funnel_query_runner.py

@@ -6,8 +6,12 @@
     ExperimentFunnelQuery,
     ExperimentFunnelQueryResponse,
     ExperimentVariantFunnelResult,
+    FunnelsQuery,
+    InsightDateRange,
+    BreakdownFilter,
 )
 from typing import Any
+from zoneinfo import ZoneInfo
 
 
 class ExperimentFunnelQueryRunner(QueryRunner):
@@ -17,16 +21,51 @@ def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
         self.experiment = Experiment.objects.get(id=self.query.experiment_id)
         self.feature_flag = self.experiment.feature_flag
-
+        self.prepared_funnel_query = self._prepare_funnel_query()
         self.query_runner = FunnelsQueryRunner(
-            query=self.query.source, team=self.team, timings=self.timings, limit_context=self.limit_context
+            query=self.prepared_funnel_query, team=self.team, timings=self.timings, limit_context=self.limit_context
         )
 
     def calculate(self) -> ExperimentFunnelQueryResponse:
         response = self.query_runner.calculate()
         results = self._process_results(response.results)
         return ExperimentFunnelQueryResponse(insight="FUNNELS", results=results)
 
+    def _prepare_funnel_query(self) -> FunnelsQuery:
+        """
+        This method takes the raw funnel query and adapts it
+        for the needs of experiment analysis:
+
+        1. Set the date range to match the experiment's duration, using the project's timezone.
+        2. Configure the breakdown to use the feature flag key, which allows us
+           to separate results for different experiment variants.
+        """
+        # Clone the source query
+        prepared_funnel_query = FunnelsQuery(**self.query.source.model_dump())
+
+        # Set the date range to match the experiment's duration, using the project's timezone
+        if self.team.timezone:
+            tz = ZoneInfo(self.team.timezone)
+            start_date = self.experiment.start_date.astimezone(tz) if self.experiment.start_date else None
+            end_date = self.experiment.end_date.astimezone(tz) if self.experiment.end_date else None
+        else:
+            start_date = self.experiment.start_date
+            end_date = self.experiment.end_date
+
+        prepared_funnel_query.dateRange = InsightDateRange(
+            date_from=start_date.isoformat() if start_date else None,
+            date_to=end_date.isoformat() if end_date else None,
+            explicitDate=True,
+        )
+
+        # Configure the breakdown to use the feature flag key
+        prepared_funnel_query.breakdownFilter = BreakdownFilter(
+            breakdown=f"$feature/{self.feature_flag.key}",
+            breakdown_type="event",
+        )
+
+        return prepared_funnel_query
+
     def _process_results(self, funnels_results: list[list[dict[str, Any]]]) -> dict[str, ExperimentVariantFunnelResult]:
         variants = self.feature_flag.variants
         processed_results = {

posthog/hogql_queries/test/test_experiment_funnel_query_runner.py

@@ -2,7 +2,6 @@
 from posthog.models.experiment import Experiment
 from posthog.models.feature_flag.feature_flag import FeatureFlag
 from posthog.schema import (
-    BreakdownFilter,
     EventsNode,
     ExperimentFunnelQuery,
     ExperimentFunnelQueryResponse,
@@ -11,9 +10,12 @@
 from posthog.test.base import APIBaseTest, ClickhouseTestMixin, _create_event, _create_person, flush_persons_and_events
 from freezegun import freeze_time
 from typing import cast
+from django.utils import timezone
+from datetime import timedelta
 
 
 class TestExperimentFunnelQueryRunner(ClickhouseTestMixin, APIBaseTest):
+    @freeze_time("2020-01-01T12:00:00Z")
     def test_query_runner(self):
         feature_flag = FeatureFlag.objects.create(
             name="Test experiment flag",
@@ -38,18 +40,20 @@ def test_query_runner(self):
             },
             created_by=self.user,
         )
+
         experiment = Experiment.objects.create(
             name="test-experiment",
             team=self.team,
             feature_flag=feature_flag,
+            start_date=timezone.now(),
+            end_date=timezone.now() + timedelta(days=14),
         )
 
         feature_flag_property = f"$feature/{feature_flag.key}"
 
         funnels_query = FunnelsQuery(
             series=[EventsNode(event="$pageview"), EventsNode(event="purchase")],
             dateRange={"date_from": "2020-01-01", "date_to": "2020-01-14"},
-            breakdownFilter=BreakdownFilter(breakdown=feature_flag_property),
         )
         experiment_query = ExperimentFunnelQuery(
             experiment_id=experiment.id,
@@ -60,25 +64,24 @@ def test_query_runner(self):
         experiment.metrics = [{"type": "primary", "query": experiment_query.model_dump()}]
         experiment.save()
 
-        with freeze_time("2020-01-10 12:00:00"):
-            for variant, purchase_count in [("control", 6), ("test", 8)]:
-                for i in range(10):
-                    _create_person(distinct_ids=[f"user_{variant}_{i}"], team_id=self.team.pk)
+        for variant, purchase_count in [("control", 6), ("test", 8)]:
+            for i in range(10):
+                _create_person(distinct_ids=[f"user_{variant}_{i}"], team_id=self.team.pk)
+                _create_event(
+                    team=self.team,
+                    event="$pageview",
+                    distinct_id=f"user_{variant}_{i}",
+                    timestamp="2020-01-02T12:00:00Z",
+                    properties={feature_flag_property: variant},
+                )
+                if i < purchase_count:
                     _create_event(
                         team=self.team,
-                        event="$pageview",
+                        event="purchase",
                         distinct_id=f"user_{variant}_{i}",
-                        timestamp="2020-01-02T12:00:00Z",
+                        timestamp="2020-01-02T12:01:00Z",
                         properties={feature_flag_property: variant},
                     )
-                    if i < purchase_count:
-                        _create_event(
-                            team=self.team,
-                            event="purchase",
-                            distinct_id=f"user_{variant}_{i}",
-                            timestamp="2020-01-02T12:01:00Z",
-                            properties={feature_flag_property: variant},
-                        )
 
         flush_persons_and_events()