Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

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

Merged
merged 2 commits into from
Sep 17, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
22 changes: 8 additions & 14 deletions frontend/src/queries/schema.json
Original file line number Diff line number Diff line change
Expand Up @@ -4536,6 +4536,9 @@
"ExperimentFunnelQuery": {
"additionalProperties": false,
"properties": {
"experiment_id": {
"type": "integer"
},
"kind": {
"const": "ExperimentFunnelQuery",
"type": "string"
Expand All @@ -4549,15 +4552,9 @@
},
"source": {
"$ref": "#/definitions/FunnelsQuery"
},
"variants": {
"items": {
"type": "string"
},
"type": "array"
}
},
"required": ["kind", "source", "variants"],
"required": ["experiment_id", "kind", "source"],
"type": "object"
},
"ExperimentFunnelQueryResponse": {
Expand All @@ -4583,6 +4580,9 @@
"count_source": {
"$ref": "#/definitions/TrendsQuery"
},
"experiment_id": {
"type": "integer"
},
"exposure_source": {
"$ref": "#/definitions/TrendsQuery"
},
Expand All @@ -4596,15 +4596,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": {
Expand Down
4 changes: 2 additions & 2 deletions frontend/src/queries/schema.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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
}

/**
Expand Down
2 changes: 1 addition & 1 deletion latest_migrations.manifest
Original file line number Diff line number Diff line change
Expand Up @@ -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
3 changes: 2 additions & 1 deletion posthog/api/test/__snapshots__/test_feature_flag.ambr
Original file line number Diff line number Diff line change
Expand Up @@ -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
'''
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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
'''
Expand Down
10 changes: 7 additions & 3 deletions posthog/hogql_queries/experiment_funnel_query_runner.py
Original file line number Diff line number Diff line change
@@ -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,
Expand All @@ -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
)
Expand All @@ -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
Comment on lines -27 to +31
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

forgive me if I'm out of the loop, but what's the ramifications/rationale behind moving the variants to come from the feature flag instead of the query?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The way I think about this - the experiment (and its linked flag) is the source of truth for experiment variants. We need to retrieve it in the query runner for other things too, so might as well retrieve the variants from there, no need to pass via the query.

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
Expand Down
12 changes: 8 additions & 4 deletions posthog/hogql_queries/experiment_trend_query_runner.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand All @@ -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
)
Expand Down Expand Up @@ -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
Expand Down
63 changes: 49 additions & 14 deletions posthog/hogql_queries/test/test_experiment_funnel_query_runner.py
Original file line number Diff line number Diff line change
@@ -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,
Expand All @@ -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)]:
Expand All @@ -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)
Expand Down
79 changes: 57 additions & 22 deletions posthog/hogql_queries/test/test_experiment_trend_query_runner.py
Original file line number Diff line number Diff line change
@@ -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,
Expand All @@ -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
Expand All @@ -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)
Expand Down
17 changes: 17 additions & 0 deletions posthog/migrations/0472_experiment_metrics.py
Original file line number Diff line number Diff line change
@@ -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),
),
]
2 changes: 2 additions & 0 deletions posthog/models/experiment.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down
Loading
Loading