From 54b0134f98168e53edab36c90b4db4d37094308d Mon Sep 17 00:00:00 2001 From: Juraj Majerik Date: Sun, 22 Sep 2024 22:57:44 +0200 Subject: [PATCH] feat(experiments HogQL rewrite): prepare funnel query (#25096) --- .../experiment_funnel_query_runner.py | 43 ++++++++++++++++++- .../test_experiment_funnel_query_runner.py | 35 ++++++++------- 2 files changed, 60 insertions(+), 18 deletions(-) diff --git a/posthog/hogql_queries/experiment_funnel_query_runner.py b/posthog/hogql_queries/experiment_funnel_query_runner.py index 36bdc42f0baf5..7931ce23f7fe2 100644 --- a/posthog/hogql_queries/experiment_funnel_query_runner.py +++ b/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,9 +21,9 @@ 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: @@ -27,6 +31,41 @@ def calculate(self) -> ExperimentFunnelQueryResponse: 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 = { diff --git a/posthog/hogql_queries/test/test_experiment_funnel_query_runner.py b/posthog/hogql_queries/test/test_experiment_funnel_query_runner.py index da9d14fb511be..8d8a9be6a9fd7 100644 --- a/posthog/hogql_queries/test/test_experiment_funnel_query_runner.py +++ b/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,10 +40,13 @@ 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}" @@ -49,7 +54,6 @@ def test_query_runner(self): 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()