feat(experiments): Add saved metrics (#25867)

ee/clickhouse/views/experiment_saved_metrics.py

@@ -0,0 +1,63 @@
+from rest_framework import serializers, viewsets
+from rest_framework.exceptions import ValidationError
+
+
+from posthog.api.routing import TeamAndOrgViewSetMixin
+from posthog.api.shared import UserBasicSerializer
+from posthog.models.experiment import ExperimentSavedMetric, ExperimentToSavedMetric
+
+
+class ExperimentToSavedMetricSerializer(serializers.ModelSerializer):
+    class Meta:
+        model = ExperimentToSavedMetric
+        fields = [
+            "id",
+            "experiment",
+            "saved_metric",
+            "metadata",
+            "created_at",
+        ]
+        read_only_fields = [
+            "id",
+            "created_at",
+        ]
+
+
+class ExperimentSavedMetricSerializer(serializers.ModelSerializer):
+    created_by = UserBasicSerializer(read_only=True)
+
+    class Meta:
+        model = ExperimentSavedMetric
+        fields = [
+            "id",
+            "name",
+            "description",
+            "query",
+            "created_by",
+            "created_at",
+            "updated_at",
+        ]
+        read_only_fields = [
+            "id",
+            "created_by",
+            "created_at",
+            "updated_at",
+        ]
+
+    def validate_query(self, value):
+        if not value:
+            raise ValidationError("Query is required to create a saved metric")
+        return value
+
+    def create(self, validated_data):
+        request = self.context["request"]
+        validated_data["created_by"] = request.user
+        validated_data["team_id"] = self.context["team_id"]
+        return super().create(validated_data)
+
+
+class ExperimentSavedMetricViewSet(TeamAndOrgViewSetMixin, viewsets.ModelViewSet):
+    scope_object = "experiment"
+    queryset = ExperimentSavedMetric.objects.prefetch_related("created_by").all()
+    serializer_class = ExperimentSavedMetricSerializer
+    ordering = "-created_at"

ee/clickhouse/views/experiments.py

@@ -20,6 +20,7 @@
 )
 from ee.clickhouse.queries.experiments.utils import requires_flag_warning
 from ee.clickhouse.views.experiment_holdouts import ExperimentHoldoutSerializer
+from ee.clickhouse.views.experiment_saved_metrics import ExperimentToSavedMetricSerializer
 from posthog.api.cohort import CohortSerializer
 from posthog.api.feature_flag import FeatureFlagSerializer, MinimalFeatureFlagSerializer
 from posthog.api.routing import TeamAndOrgViewSetMixin
@@ -28,7 +29,7 @@
 from posthog.caching.insight_cache import update_cached_state
 from posthog.clickhouse.query_tagging import tag_queries
 from posthog.constants import INSIGHT_TRENDS
-from posthog.models.experiment import Experiment, ExperimentHoldout
+from posthog.models.experiment import Experiment, ExperimentHoldout, ExperimentSavedMetric
 from posthog.models.filters.filter import Filter
 from posthog.utils import generate_cache_key, get_safe_cache
 
@@ -163,6 +164,10 @@ class ExperimentSerializer(serializers.ModelSerializer):
     holdout_id = serializers.PrimaryKeyRelatedField(
         queryset=ExperimentHoldout.objects.all(), source="holdout", required=False, allow_null=True
     )
+    saved_metrics = ExperimentToSavedMetricSerializer(many=True, source="experimenttosavedmetric_set", read_only=True)
+    saved_metrics_ids = serializers.ListField(
+        child=serializers.JSONField(), write_only=True, required=False, allow_null=True
+    )
 
     class Meta:
         model = Experiment
@@ -179,6 +184,8 @@ class Meta:
             "exposure_cohort",
             "parameters",
             "secondary_metrics",
+            "saved_metrics",
+            "saved_metrics_ids",
             "filters",
             "archived",
             "created_by",
@@ -194,8 +201,37 @@ class Meta:
             "feature_flag",
             "exposure_cohort",
             "holdout",
+            "saved_metrics",
         ]
 
+    def validate_saved_metrics_ids(self, value):
+        if value is None:
+            return value
+
+        # check value is valid json list with id and optionally metadata param
+        if not isinstance(value, list):
+            raise ValidationError("Saved metrics must be a list")
+
+        for saved_metric in value:
+            if not isinstance(saved_metric, dict):
+                raise ValidationError("Saved metric must be an object")
+            if "id" not in saved_metric:
+                raise ValidationError("Saved metric must have an id")
+            if "metadata" in saved_metric and not isinstance(saved_metric["metadata"], dict):
+                raise ValidationError("Metadata must be an object")
+
+            # metadata is optional, but if it exists, should have type key
+            # TODO: extend with other metadata keys when known
+            if "metadata" in saved_metric and "type" not in saved_metric["metadata"]:
+                raise ValidationError("Metadata must have a type key")
+
+        # check if all saved metrics exist
+        saved_metrics = ExperimentSavedMetric.objects.filter(id__in=[saved_metric["id"] for saved_metric in value])
+        if saved_metrics.count() != len(value):
+            raise ValidationError("Saved metric does not exist")
+
+        return value
+
     def validate_parameters(self, value):
         if not value:
             return value
@@ -214,6 +250,8 @@ def create(self, validated_data: dict, *args: Any, **kwargs: Any) -> Experiment:
         if not validated_data.get("filters"):
             raise ValidationError("Filters are required to create an Experiment")
 
+        saved_metrics_data = validated_data.pop("saved_metrics_ids", [])
+
         variants = []
         aggregation_group_type_index = None
         if validated_data["parameters"]:
@@ -264,9 +302,42 @@ def create(self, validated_data: dict, *args: Any, **kwargs: Any) -> Experiment:
         experiment = Experiment.objects.create(
             team_id=self.context["team_id"], feature_flag=feature_flag, **validated_data
         )
+
+        if saved_metrics_data:
+            for saved_metric_data in saved_metrics_data:
+                saved_metric_serializer = ExperimentToSavedMetricSerializer(
+                    data={
+                        "experiment": experiment.id,
+                        "saved_metric": saved_metric_data["id"],
+                        "metadata": saved_metric_data.get("metadata"),
+                    },
+                    context=self.context,
+                )
+                saved_metric_serializer.is_valid(raise_exception=True)
+                saved_metric_serializer.save()
+                # TODO: Going the above route means we can still sometimes fail when validation fails?
+                # But this shouldn't really happen, if it does its a bug in our validation logic (validate_saved_metrics_ids)
         return experiment
 
     def update(self, instance: Experiment, validated_data: dict, *args: Any, **kwargs: Any) -> Experiment:
+        update_saved_metrics = "saved_metrics_ids" in validated_data
+        saved_metrics_data = validated_data.pop("saved_metrics_ids", []) or []
+
+        # We replace all saved metrics on update to avoid issues with partial updates
+        if update_saved_metrics:
+            instance.experimenttosavedmetric_set.all().delete()
+            for saved_metric_data in saved_metrics_data:
+                saved_metric_serializer = ExperimentToSavedMetricSerializer(
+                    data={
+                        "experiment": instance.id,
+                        "saved_metric": saved_metric_data["id"],
+                        "metadata": saved_metric_data.get("metadata"),
+                    },
+                    context=self.context,
+                )
+                saved_metric_serializer.is_valid(raise_exception=True)
+                saved_metric_serializer.save()
+
         has_start_date = validated_data.get("start_date") is not None
         feature_flag = instance.feature_flag
 
@@ -371,7 +442,9 @@ def update(self, instance: Experiment, validated_data: dict, *args: Any, **kwarg
 class EnterpriseExperimentsViewSet(TeamAndOrgViewSetMixin, viewsets.ModelViewSet):
     scope_object = "experiment"
     serializer_class = ExperimentSerializer
-    queryset = Experiment.objects.prefetch_related("feature_flag", "created_by", "holdout").all()
+    queryset = Experiment.objects.prefetch_related(
+        "feature_flag", "created_by", "holdout", "experimenttosavedmetric_set", "saved_metrics"
+    ).all()
     ordering = "-created_at"
 
     # ******************************************