feat(hogql): stickiness query runner (#19676)

* WIP * Added stickiness query runner
PostHog · Jan 10, 2024 · e4591cb · e4591cb
1 parent 1c62159
commit e4591cb
Show file tree

Hide file tree

Showing 12 changed files with 321 additions and 6 deletions.
diff --git a/frontend/src/lib/constants.tsx b/frontend/src/lib/constants.tsx
@@ -167,6 +167,7 @@ export const FEATURE_FLAGS = {
     HOGQL_INSIGHTS_LIFECYCLE: 'hogql-insights-lifecycle', // owner: @mariusandra
     HOGQL_INSIGHTS_RETENTION: 'hogql-insights-retention', // owner: @webjunkie
     HOGQL_INSIGHTS_TRENDS: 'hogql-insights-trends', // owner: @Gilbert09
+    HOGQL_INSIGHTS_STICKINESS: 'hogql-insights-stickiness', // owner: @Gilbert09
     HOGQL_INSIGHT_LIVE_COMPARE: 'hogql-insight-live-compare', // owner: @mariusandra
     BI_VIZ: 'bi_viz', // owner: @Gilbert09
     WEBHOOKS_DENYLIST: 'webhooks-denylist', // owner: #team-pipeline

diff --git a/frontend/src/queries/nodes/DataNode/dataNodeLogic.ts b/frontend/src/queries/nodes/DataNode/dataNodeLogic.ts
@@ -48,6 +48,7 @@ import {
     isInsightQueryNode,
     isLifecycleQuery,
     isPersonsNode,
+    isStickinessQuery,
     isTrendsQuery,
 } from '~/queries/utils'
 
@@ -126,6 +127,7 @@ export const dataNodeLogic = kea<dataNodeLogicType>([
                         isInsightQueryNode(props.query) &&
                         !(values.hogQLInsightsLifecycleFlagEnabled && isLifecycleQuery(props.query)) &&
                         !(values.hogQLInsightsTrendsFlagEnabled && isTrendsQuery(props.query)) &&
+                        !(values.hogQLInsightsStickinessFlagEnabled && isStickinessQuery(props.query)) &&
                         props.cachedResults &&
                         props.cachedResults['id'] &&
                         props.cachedResults['filters'] &&
@@ -357,6 +359,10 @@ export const dataNodeLogic = kea<dataNodeLogicType>([
             (s) => [s.featureFlags],
             (featureFlags) => !!featureFlags[FEATURE_FLAGS.HOGQL_INSIGHTS_TRENDS],
         ],
+        hogQLInsightsStickinessFlagEnabled: [
+            (s) => [s.featureFlags],
+            (featureFlags) => !!featureFlags[FEATURE_FLAGS.HOGQL_INSIGHTS_STICKINESS],
+        ],
         query: [(_, p) => [p.query], (query) => query],
         newQuery: [
             (s, p) => [p.query, s.response],

diff --git a/frontend/src/queries/query.ts b/frontend/src/queries/query.ts
@@ -31,6 +31,7 @@ import {
     isLifecycleQuery,
     isPersonsNode,
     isRetentionQuery,
+    isStickinessQuery,
     isTimeToSeeDataQuery,
     isTimeToSeeDataSessionsNode,
     isTimeToSeeDataSessionsQuery,
@@ -155,6 +156,9 @@ export async function query<N extends DataNode = DataNode>(
     const hogQLInsightsTrendsFlagEnabled = Boolean(
         featureFlagLogic.findMounted()?.values.featureFlags?.[FEATURE_FLAGS.HOGQL_INSIGHTS_TRENDS]
     )
+    const hogQLInsightsStickinessFlagEnabled = Boolean(
+        featureFlagLogic.findMounted()?.values.featureFlags?.[FEATURE_FLAGS.HOGQL_INSIGHTS_STICKINESS]
+    )
     const hogQLInsightsLiveCompareEnabled = Boolean(
         featureFlagLogic.findMounted()?.values.featureFlags?.[FEATURE_FLAGS.HOGQL_INSIGHT_LIVE_COMPARE]
     )
@@ -198,7 +202,8 @@ export async function query<N extends DataNode = DataNode>(
             if (
                 (hogQLInsightsLifecycleFlagEnabled && isLifecycleQuery(queryNode)) ||
                 (hogQLInsightsRetentionFlagEnabled && isRetentionQuery(queryNode)) ||
-                (hogQLInsightsTrendsFlagEnabled && isTrendsQuery(queryNode))
+                (hogQLInsightsTrendsFlagEnabled && isTrendsQuery(queryNode)) ||
+                (hogQLInsightsStickinessFlagEnabled && isStickinessQuery(queryNode))
             ) {
                 if (hogQLInsightsLiveCompareEnabled) {
                     let legacyResponse

diff --git a/frontend/src/queries/schema.json b/frontend/src/queries/schema.json
@@ -3053,6 +3053,37 @@
             "required": ["kind", "series"],
             "type": "object"
         },
+        "StickinessQueryResponse": {
+            "additionalProperties": false,
+            "properties": {
+                "hogql": {
+                    "type": "string"
+                },
+                "is_cached": {
+                    "type": "boolean"
+                },
+                "last_refresh": {
+                    "type": "string"
+                },
+                "next_allowed_client_refresh": {
+                    "type": "string"
+                },
+                "results": {
+                    "items": {
+                        "type": "object"
+                    },
+                    "type": "array"
+                },
+                "timings": {
+                    "items": {
+                        "$ref": "#/definitions/QueryTiming"
+                    },
+                    "type": "array"
+                }
+            },
+            "required": ["results"],
+            "type": "object"
+        },
         "TimeToSeeDataJSONNode": {
             "additionalProperties": false,
             "properties": {

diff --git a/frontend/src/queries/schema.ts b/frontend/src/queries/schema.ts
@@ -571,6 +571,11 @@ export type StickinessFilter = Omit<
     StickinessFilterType & { hidden_legend_indexes?: number[] },
     keyof FilterType | 'hidden_legend_keys' | 'stickiness_days' | 'shown_as'
 >
+
+export interface StickinessQueryResponse extends QueryResponse {
+    results: Record<string, any>[]
+}
+
 export interface StickinessQuery extends Omit<InsightsQueryBase, 'aggregation_group_type_index'> {
     kind: NodeKind.StickinessQuery
     /** Granularity of the response. Can be one of `hour`, `day`, `week` or `month` */

diff --git a/mypy-baseline.txt b/mypy-baseline.txt
@@ -365,7 +365,6 @@ posthog/hogql/query.py:0: error: Argument 1 to "get_default_limit_for_context" h
 posthog/hogql/query.py:0: error: "SelectQuery" has no attribute "select_queries"  [attr-defined]
 posthog/hogql/query.py:0: error: Subclass of "SelectQuery" and "SelectUnionQuery" cannot exist: would have incompatible method signatures  [unreachable]
 posthog/hogql_queries/query_runner.py:0: error: Unused "type: ignore" comment  [unused-ignore]
-posthog/hogql_queries/query_runner.py:0: error: Argument "limit_context" to "TrendsQueryRunner" has incompatible type "LimitContext | None"; expected "bool | None"  [arg-type]
 posthog/hogql_queries/insights/trends/breakdown_values.py:0: error: Argument "chain" to "Field" has incompatible type "list[str]"; expected "list[str | int]"  [arg-type]
 posthog/hogql_queries/insights/trends/breakdown_values.py:0: note: "List" is invariant -- see https://mypy.readthedocs.io/en/stable/common_issues.html#variance
 posthog/hogql_queries/insights/trends/breakdown_values.py:0: note: Consider using "Sequence" instead, which is covariant
@@ -439,7 +438,6 @@ posthog/api/person.py:0: error: Unused "type: ignore" comment  [unused-ignore]
 posthog/api/person.py:0: error: Unused "type: ignore" comment  [unused-ignore]
 posthog/api/person.py:0: error: Unused "type: ignore" comment  [unused-ignore]
 posthog/hogql_queries/web_analytics/web_analytics_query_runner.py:0: error: Argument 1 to "append" of "list" has incompatible type "EventPropertyFilter"; expected "Expr"  [arg-type]
-posthog/hogql_queries/insights/trends/trends_query_runner.py:0: error: Argument "limit_context" to "__init__" of "QueryRunner" has incompatible type "bool | None"; expected "LimitContext | None"  [arg-type]
 posthog/hogql_queries/insights/trends/trends_query_runner.py:0: error: Return type "list[SelectQuery]" of "to_query" incompatible with return type "SelectQuery" in supertype "QueryRunner"  [override]
 posthog/hogql_queries/insights/trends/trends_query_runner.py:0: error: Incompatible return value type (got "list[SelectQuery | SelectUnionQuery]", expected "list[SelectQuery]")  [return-value]
 posthog/hogql_queries/insights/trends/trends_query_runner.py:0: error: Need type annotation for "timings" (hint: "timings: List[<type>] = ...")  [var-annotated]

diff --git a/posthog/api/services/query.py b/posthog/api/services/query.py
@@ -21,6 +21,7 @@
     "LifecycleQuery",
     "RetentionQuery",
     "TrendsQuery",
+    "StickinessQuery",
     "WebOverviewQuery",
     "WebTopSourcesQuery",
     "WebTopClicksQuery",

diff --git a/posthog/hogql/ast.py b/posthog/hogql/ast.py
@@ -389,8 +389,8 @@ class ArithmeticOperation(Expr):
 
 @dataclass(kw_only=True)
 class And(Expr):
-    exprs: List[Expr]
     type: Optional[ConstantType] = None
+    exprs: List[Expr]
 
 
 @dataclass(kw_only=True)
@@ -566,8 +566,8 @@ class SelectQuery(Expr):
 
 @dataclass(kw_only=True)
 class SelectUnionQuery(Expr):
-    select_queries: List[SelectQuery]
     type: Optional[SelectUnionQueryType] = None
+    select_queries: List[SelectQuery]
 
 
 @dataclass(kw_only=True)