feat(experiments HogQL): add response attributes (#25333)

PostHog · Oct 3, 2024 · d9d6296 · d9d6296
1 parent e4c7b24
commit d9d6296
Show file tree

Hide file tree

Showing 5 changed files with 158 additions and 0 deletions.
diff --git a/frontend/src/queries/schema.json b/frontend/src/queries/schema.json
@@ -1036,6 +1036,114 @@
             ],
             "type": "object"
         },
+        "CachedExperimentFunnelQueryResponse": {
+            "additionalProperties": false,
+            "properties": {
+                "cache_key": {
+                    "type": "string"
+                },
+                "cache_target_age": {
+                    "format": "date-time",
+                    "type": "string"
+                },
+                "calculation_trigger": {
+                    "description": "What triggered the calculation of the query, leave empty if user/immediate",
+                    "type": "string"
+                },
+                "insight": {
+                    "const": "FUNNELS",
+                    "type": "string"
+                },
+                "is_cached": {
+                    "type": "boolean"
+                },
+                "last_refresh": {
+                    "format": "date-time",
+                    "type": "string"
+                },
+                "next_allowed_client_refresh": {
+                    "format": "date-time",
+                    "type": "string"
+                },
+                "query_status": {
+                    "$ref": "#/definitions/QueryStatus",
+                    "description": "Query status indicates whether next to the provided data, a query is still running."
+                },
+                "results": {
+                    "additionalProperties": {
+                        "$ref": "#/definitions/ExperimentVariantFunnelResult"
+                    },
+                    "type": "object"
+                },
+                "timezone": {
+                    "type": "string"
+                }
+            },
+            "required": [
+                "cache_key",
+                "insight",
+                "is_cached",
+                "last_refresh",
+                "next_allowed_client_refresh",
+                "results",
+                "timezone"
+            ],
+            "type": "object"
+        },
+        "CachedExperimentTrendQueryResponse": {
+            "additionalProperties": false,
+            "properties": {
+                "cache_key": {
+                    "type": "string"
+                },
+                "cache_target_age": {
+                    "format": "date-time",
+                    "type": "string"
+                },
+                "calculation_trigger": {
+                    "description": "What triggered the calculation of the query, leave empty if user/immediate",
+                    "type": "string"
+                },
+                "insight": {
+                    "const": "TRENDS",
+                    "type": "string"
+                },
+                "is_cached": {
+                    "type": "boolean"
+                },
+                "last_refresh": {
+                    "format": "date-time",
+                    "type": "string"
+                },
+                "next_allowed_client_refresh": {
+                    "format": "date-time",
+                    "type": "string"
+                },
+                "query_status": {
+                    "$ref": "#/definitions/QueryStatus",
+                    "description": "Query status indicates whether next to the provided data, a query is still running."
+                },
+                "results": {
+                    "additionalProperties": {
+                        "$ref": "#/definitions/ExperimentVariantTrendResult"
+                    },
+                    "type": "object"
+                },
+                "timezone": {
+                    "type": "string"
+                }
+            },
+            "required": [
+                "cache_key",
+                "insight",
+                "is_cached",
+                "last_refresh",
+                "next_allowed_client_refresh",
+                "results",
+                "timezone"
+            ],
+            "type": "object"
+        },
         "CachedFunnelCorrelationResponse": {
             "additionalProperties": false,
             "properties": {

diff --git a/frontend/src/queries/schema.ts b/frontend/src/queries/schema.ts
@@ -1594,11 +1594,15 @@ export interface ExperimentTrendQueryResponse {
     results: Record<string, ExperimentVariantTrendResult>
 }
 
+export type CachedExperimentTrendQueryResponse = CachedQueryResponse<ExperimentTrendQueryResponse>
+
 export interface ExperimentFunnelQueryResponse {
     insight: InsightType.FUNNELS
     results: Record<string, ExperimentVariantFunnelResult>
 }
 
+export type CachedExperimentFunnelQueryResponse = CachedQueryResponse<ExperimentFunnelQueryResponse>
+
 export interface ExperimentFunnelQuery extends DataNode<ExperimentFunnelQueryResponse> {
     kind: NodeKind.ExperimentFunnelQuery
     source: FunnelsQuery

diff --git a/posthog/hogql_queries/experiments/experiment_funnel_query_runner.py b/posthog/hogql_queries/experiments/experiment_funnel_query_runner.py
@@ -3,6 +3,7 @@
 from posthog.models.experiment import Experiment
 from ..insights.funnels.funnels_query_runner import FunnelsQueryRunner
 from posthog.schema import (
+    CachedExperimentFunnelQueryResponse,
     ExperimentFunnelQuery,
     ExperimentFunnelQueryResponse,
     ExperimentVariantFunnelResult,
@@ -16,6 +17,8 @@
 
 class ExperimentFunnelQueryRunner(QueryRunner):
     query: ExperimentFunnelQuery
+    response: ExperimentFunnelQueryResponse
+    cached_response: CachedExperimentFunnelQueryResponse
 
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)

diff --git a/posthog/hogql_queries/experiments/experiment_trend_query_runner.py b/posthog/hogql_queries/experiments/experiment_trend_query_runner.py
@@ -8,6 +8,7 @@
 from posthog.schema import (
     BaseMathType,
     BreakdownFilter,
+    CachedExperimentTrendQueryResponse,
     ChartDisplayType,
     EventPropertyFilter,
     EventsNode,
@@ -25,6 +26,8 @@
 
 class ExperimentTrendQueryRunner(QueryRunner):
     query: ExperimentTrendQuery
+    response: ExperimentTrendQueryResponse
+    cached_response: CachedExperimentTrendQueryResponse
 
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)

diff --git a/posthog/schema.py b/posthog/schema.py
@@ -1793,6 +1793,46 @@ class CachedEventsQueryResponse(BaseModel):
     types: list[str]
 
 
+class CachedExperimentFunnelQueryResponse(BaseModel):
+    model_config = ConfigDict(
+        extra="forbid",
+    )
+    cache_key: str
+    cache_target_age: Optional[AwareDatetime] = None
+    calculation_trigger: Optional[str] = Field(
+        default=None, description="What triggered the calculation of the query, leave empty if user/immediate"
+    )
+    insight: Literal["FUNNELS"] = "FUNNELS"
+    is_cached: bool
+    last_refresh: AwareDatetime
+    next_allowed_client_refresh: AwareDatetime
+    query_status: Optional[QueryStatus] = Field(
+        default=None, description="Query status indicates whether next to the provided data, a query is still running."
+    )
+    results: dict[str, ExperimentVariantFunnelResult]
+    timezone: str
+
+
+class CachedExperimentTrendQueryResponse(BaseModel):
+    model_config = ConfigDict(
+        extra="forbid",
+    )
+    cache_key: str
+    cache_target_age: Optional[AwareDatetime] = None
+    calculation_trigger: Optional[str] = Field(
+        default=None, description="What triggered the calculation of the query, leave empty if user/immediate"
+    )
+    insight: Literal["TRENDS"] = "TRENDS"
+    is_cached: bool
+    last_refresh: AwareDatetime
+    next_allowed_client_refresh: AwareDatetime
+    query_status: Optional[QueryStatus] = Field(
+        default=None, description="Query status indicates whether next to the provided data, a query is still running."
+    )
+    results: dict[str, ExperimentVariantTrendResult]
+    timezone: str
+
+
 class CachedFunnelCorrelationResponse(BaseModel):
     model_config = ConfigDict(
         extra="forbid",