feat(experiments HogQL rewrite): fetch associated Experiment (#24992)

Co-authored-by: github-actions <41898282+github-actions[bot]@users.noreply.github.com>

frontend/src/queries/schema.json

@@ -4540,6 +4540,9 @@
         "ExperimentFunnelQuery": {
             "additionalProperties": false,
             "properties": {
+                "experiment_id": {
+                    "type": "integer"
+                },
                 "kind": {
                     "const": "ExperimentFunnelQuery",
                     "type": "string"
@@ -4553,15 +4556,9 @@
                 },
                 "source": {
                     "$ref": "#/definitions/FunnelsQuery"
-                },
-                "variants": {
-                    "items": {
-                        "type": "string"
-                    },
-                    "type": "array"
                 }
             },
-            "required": ["kind", "source", "variants"],
+            "required": ["experiment_id", "kind", "source"],
             "type": "object"
         },
         "ExperimentFunnelQueryResponse": {
@@ -4587,6 +4584,9 @@
                 "count_source": {
                     "$ref": "#/definitions/TrendsQuery"
                 },
+                "experiment_id": {
+                    "type": "integer"
+                },
                 "exposure_source": {
                     "$ref": "#/definitions/TrendsQuery"
                 },
@@ -4600,15 +4600,9 @@
                 },
                 "response": {
                     "$ref": "#/definitions/ExperimentTrendQueryResponse"
-                },
-                "variants": {
-                    "items": {
-                        "type": "string"
-                    },
-                    "type": "array"
                 }
             },
-            "required": ["count_source", "exposure_source", "kind", "variants"],
+            "required": ["count_source", "experiment_id", "exposure_source", "kind"],
             "type": "object"
         },
         "ExperimentTrendQueryResponse": {

frontend/src/queries/schema.ts

@@ -1577,14 +1577,14 @@ export interface ExperimentFunnelQueryResponse {
 export interface ExperimentFunnelQuery extends DataNode<ExperimentFunnelQueryResponse> {
     kind: NodeKind.ExperimentFunnelQuery
     source: FunnelsQuery
-    variants: string[]
+    experiment_id: integer
 }
 
 export interface ExperimentTrendQuery extends DataNode<ExperimentTrendQueryResponse> {
     kind: NodeKind.ExperimentTrendQuery
     count_source: TrendsQuery
     exposure_source: TrendsQuery
-    variants: string[]
+    experiment_id: integer
 }
 
 /**

latest_migrations.manifest

@@ -5,7 +5,7 @@ contenttypes: 0002_remove_content_type_name
 ee: 0016_rolemembership_organization_member
 otp_static: 0002_throttling
 otp_totp: 0002_auto_20190420_0723
-posthog: 0471_webexperiment_experiment_type_experiment_variants
+posthog: 0472_experiment_metrics
 sessions: 0001_initial
 social_django: 0010_uid_db_index
 two_factor: 0007_auto_20201201_1019

posthog/api/test/__snapshots__/test_feature_flag.ambr

@@ -1785,7 +1785,8 @@
          "posthog_experiment"."updated_at",
          "posthog_experiment"."archived",
          "posthog_experiment"."type",
-         "posthog_experiment"."variants"
+         "posthog_experiment"."variants",
+         "posthog_experiment"."metrics"
   FROM "posthog_experiment"
   WHERE "posthog_experiment"."exposure_cohort_id" = 2
   '''

posthog/api/test/__snapshots__/test_organization_feature_flag.ambr

@@ -1035,7 +1035,8 @@
          "posthog_experiment"."updated_at",
          "posthog_experiment"."archived",
          "posthog_experiment"."type",
-         "posthog_experiment"."variants"
+         "posthog_experiment"."variants",
+         "posthog_experiment"."metrics"
   FROM "posthog_experiment"
   WHERE "posthog_experiment"."feature_flag_id" = 2
   '''

posthog/hogql_queries/experiment_funnel_query_runner.py

@@ -1,5 +1,6 @@
 from posthog.hogql import ast
 from posthog.hogql_queries.query_runner import QueryRunner
+from posthog.models.experiment import Experiment
 from .insights.funnels.funnels_query_runner import FunnelsQueryRunner
 from posthog.schema import (
     ExperimentFunnelQuery,
@@ -14,6 +15,9 @@ class ExperimentFunnelQueryRunner(QueryRunner):
 
     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.query_runner = FunnelsQueryRunner(
             query=self.query.source, team=self.team, timings=self.timings, limit_context=self.limit_context
         )
@@ -24,17 +28,17 @@ def calculate(self) -> ExperimentFunnelQueryResponse:
         return ExperimentFunnelQueryResponse(insight="FUNNELS", results=results)
 
     def _process_results(self, funnels_results: list[list[dict[str, Any]]]) -> dict[str, ExperimentVariantFunnelResult]:
-        variants = self.query.variants
+        variants = self.feature_flag.variants
         processed_results = {
-            variant: ExperimentVariantFunnelResult(success_count=0, failure_count=0) for variant in variants
+            variant["key"]: ExperimentVariantFunnelResult(success_count=0, failure_count=0) for variant in variants
         }
 
         for result in funnels_results:
             first_step = result[0]
             last_step = result[-1]
             variant = first_step.get("breakdown_value")
             variant_str = variant[0] if isinstance(variant, list) else str(variant)
-            if variant_str in variants:
+            if variant_str in processed_results:
                 total_count = first_step.get("count", 0)
                 success_count = last_step.get("count", 0) if len(result) > 1 else 0
                 processed_results[variant_str].success_count = success_count

posthog/hogql_queries/experiment_trend_query_runner.py

@@ -2,6 +2,7 @@
 from posthog.hogql import ast
 from posthog.hogql_queries.insights.trends.trends_query_runner import TrendsQueryRunner
 from posthog.hogql_queries.query_runner import QueryRunner
+from posthog.models.experiment import Experiment
 from posthog.schema import (
     ExperimentTrendQuery,
     ExperimentTrendQueryResponse,
@@ -16,6 +17,9 @@ class ExperimentTrendQueryRunner(QueryRunner):
 
     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.query_runner = TrendsQueryRunner(
             query=self.query.count_source, team=self.team, timings=self.timings, limit_context=self.limit_context
         )
@@ -68,17 +72,17 @@ def run(query_runner: TrendsQueryRunner, is_parallel: bool):
     def _process_results(
         self, count_results: list[dict[str, Any]], exposure_results: list[dict[str, Any]]
     ) -> dict[str, ExperimentVariantTrendResult]:
-        variants = self.query.variants
-        processed_results = {variant: ExperimentVariantTrendResult(count=0, exposure=0) for variant in variants}
+        variants = self.feature_flag.variants
+        processed_results = {variant["key"]: ExperimentVariantTrendResult(count=0, exposure=0) for variant in variants}
 
         for result in count_results:
             variant = result.get("breakdown_value")
-            if variant in variants:
+            if variant in processed_results:
                 processed_results[variant].count += result.get("count", 0)
 
         for result in exposure_results:
             variant = result.get("breakdown_value")
-            if variant in variants:
+            if variant in processed_results:
                 processed_results[variant].exposure += result.get("count", 0)
 
         return processed_results

posthog/hogql_queries/test/test_experiment_funnel_query_runner.py

@@ -1,4 +1,6 @@
 from posthog.hogql_queries.experiment_funnel_query_runner import ExperimentFunnelQueryRunner
+from posthog.models.experiment import Experiment
+from posthog.models.feature_flag.feature_flag import FeatureFlag
 from posthog.schema import (
     BreakdownFilter,
     EventsNode,
@@ -12,11 +14,51 @@
 
 
 class TestExperimentFunnelQueryRunner(ClickhouseTestMixin, APIBaseTest):
-    def setUp(self):
-        super().setUp()
-
     def test_query_runner(self):
-        feature_flag_property = f"$feature/test-experiment"
+        feature_flag = FeatureFlag.objects.create(
+            name="Test experiment flag",
+            key="test-experiment",
+            team=self.team,
+            filters={
+                "groups": [{"properties": [], "rollout_percentage": None}],
+                "multivariate": {
+                    "variants": [
+                        {
+                            "key": "control",
+                            "name": "Control",
+                            "rollout_percentage": 50,
+                        },
+                        {
+                            "key": "test",
+                            "name": "Test",
+                            "rollout_percentage": 50,
+                        },
+                    ]
+                },
+            },
+            created_by=self.user,
+        )
+        experiment = Experiment.objects.create(
+            name="test-experiment",
+            team=self.team,
+            feature_flag=feature_flag,
+        )
+
+        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,
+            kind="ExperimentFunnelQuery",
+            source=funnels_query,
+        )
+
+        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)]:
@@ -40,17 +82,10 @@ def test_query_runner(self):
 
         flush_persons_and_events()
 
-        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),
+        query_runner = ExperimentFunnelQueryRunner(
+            query=ExperimentFunnelQuery(**experiment.metrics[0]["query"]), team=self.team
         )
-        experiment_query = ExperimentFunnelQuery(
-            kind="ExperimentFunnelQuery", source=funnels_query, variants=["control", "test"]
-        )
-
-        runner = ExperimentFunnelQueryRunner(query=experiment_query, team=self.team)
-        result = runner.calculate()
+        result = query_runner.calculate()
 
         self.assertEqual(result.insight, "FUNNELS")
         self.assertEqual(len(result.results), 2)

posthog/hogql_queries/test/test_experiment_trend_query_runner.py

@@ -1,5 +1,7 @@
 from django.test import override_settings
 from posthog.hogql_queries.experiment_trend_query_runner import ExperimentTrendQueryRunner
+from posthog.models.experiment import Experiment
+from posthog.models.feature_flag.feature_flag import FeatureFlag
 from posthog.schema import (
     BreakdownFilter,
     EventsNode,
@@ -15,7 +17,58 @@
 @override_settings(IN_UNIT_TESTING=True)
 class TestExperimentTrendQueryRunner(ClickhouseTestMixin, APIBaseTest):
     def test_query_runner(self):
-        feature_flag_property = f"$feature/test-experiment"
+        feature_flag = FeatureFlag.objects.create(
+            name="Test experiment flag",
+            key="test-experiment",
+            team=self.team,
+            filters={
+                "groups": [{"properties": [], "rollout_percentage": None}],
+                "multivariate": {
+                    "variants": [
+                        {
+                            "key": "control",
+                            "name": "Control",
+                            "rollout_percentage": 50,
+                        },
+                        {
+                            "key": "test",
+                            "name": "Test",
+                            "rollout_percentage": 50,
+                        },
+                    ]
+                },
+            },
+            created_by=self.user,
+        )
+        experiment = Experiment.objects.create(
+            name="test-experiment",
+            team=self.team,
+            feature_flag=feature_flag,
+        )
+
+        feature_flag_property = f"$feature/{feature_flag.key}"
+        count_query = TrendsQuery(
+            kind="TrendsQuery",
+            series=[EventsNode(event="$pageview")],
+            dateRange={"date_from": "2020-01-01", "date_to": "2020-01-14"},
+            breakdownFilter=BreakdownFilter(breakdown=feature_flag_property),
+        )
+        exposure_query = TrendsQuery(
+            kind="TrendsQuery",
+            series=[EventsNode(event="$feature_flag_called")],
+            dateRange={"date_from": "2020-01-01", "date_to": "2020-01-14"},
+            breakdownFilter=BreakdownFilter(breakdown=feature_flag_property),
+        )
+
+        experiment_query = ExperimentTrendQuery(
+            experiment_id=experiment.id,
+            kind="ExperimentTrendQuery",
+            count_source=count_query,
+            exposure_source=exposure_query,
+        )
+
+        experiment.metrics = [{"type": "primary", "query": experiment_query.model_dump()}]
+        experiment.save()
 
         with freeze_time("2020-01-10 12:00:00"):
             # Populate experiment events
@@ -40,28 +93,10 @@ def test_query_runner(self):
 
         flush_persons_and_events()
 
-        count_query = TrendsQuery(
-            kind="TrendsQuery",
-            series=[EventsNode(event="$pageview")],
-            dateRange={"date_from": "2020-01-01", "date_to": "2020-01-14"},
-            breakdownFilter=BreakdownFilter(breakdown=feature_flag_property),
-        )
-        exposure_query = TrendsQuery(
-            kind="TrendsQuery",
-            series=[EventsNode(event="$feature_flag_called")],
-            dateRange={"date_from": "2020-01-01", "date_to": "2020-01-14"},
-            breakdownFilter=BreakdownFilter(breakdown=feature_flag_property),
-        )
-
-        experiment_query = ExperimentTrendQuery(
-            kind="ExperimentTrendQuery",
-            count_source=count_query,
-            exposure_source=exposure_query,
-            variants=["control", "test"],
+        query_runner = ExperimentTrendQueryRunner(
+            query=ExperimentTrendQuery(**experiment.metrics[0]["query"]), team=self.team
         )
-
-        runner = ExperimentTrendQueryRunner(query=experiment_query, team=self.team)
-        result = runner.calculate()
+        result = query_runner.calculate()
 
         self.assertEqual(result.insight, "TRENDS")
         self.assertEqual(len(result.results), 2)

posthog/migrations/0472_experiment_metrics.py

@@ -0,0 +1,17 @@
+# Generated by Django 4.2.15 on 2024-09-16 12:45
+
+from django.db import migrations, models
+
+
+class Migration(migrations.Migration):
+    dependencies = [
+        ("posthog", "0471_webexperiment_experiment_type_experiment_variants"),
+    ]
+
+    operations = [
+        migrations.AddField(
+            model_name="experiment",
+            name="metrics",
+            field=models.JSONField(blank=True, default=list, null=True),
+        ),
+    ]

posthog/models/experiment.py

@@ -38,6 +38,8 @@ class ExperimentType(models.TextChoices):
     type = models.CharField(max_length=40, choices=ExperimentType.choices, null=True, blank=True, default="product")
     variants = models.JSONField(default=dict, null=True, blank=True)
 
+    metrics = models.JSONField(default=list, null=True, blank=True)
+
     def get_feature_flag_key(self):
         return self.feature_flag.key