feat(hogql-queries): modifiers in query runners (#18680)

posthog/hogql_queries/events_query_runner.py

@@ -1,6 +1,6 @@
 import json
 from datetime import timedelta
-from typing import Dict, List, Optional, Any
+from typing import Dict, List, Optional
 
 from dateutil.parser import isoparse
 from django.db.models import Prefetch
@@ -15,7 +15,7 @@
 from posthog.hogql.query import execute_hogql_query
 from posthog.hogql.timings import HogQLTimings
 from posthog.hogql_queries.query_runner import QueryRunner
-from posthog.models import Action, Person, Team
+from posthog.models import Action, Person
 from posthog.models.element import chain_to_elements
 from posthog.models.person.person import get_distinct_ids_for_subquery
 from posthog.models.person.util import get_persons_by_distinct_ids
@@ -39,19 +39,6 @@ class EventsQueryRunner(QueryRunner):
     query: EventsQuery
     query_type = EventsQuery
 
-    def __init__(
-        self,
-        query: EventsQuery | Dict[str, Any],
-        team: Team,
-        timings: Optional[HogQLTimings] = None,
-        in_export_context: Optional[bool] = False,
-    ):
-        super().__init__(query, team, timings, in_export_context)
-        if isinstance(query, EventsQuery):
-            self.query = query
-        else:
-            self.query = EventsQuery.model_validate(query)
-
     def to_query(self) -> ast.SelectQuery:
         # Note: This code is inefficient and problematic, see https://github.com/PostHog/posthog/issues/13485 for details.
         if self.timings is None:
@@ -199,6 +186,7 @@ def calculate(self) -> EventsQueryResponse:
             workload=Workload.ONLINE,
             query_type="EventsQuery",
             timings=self.timings,
+            modifiers=self.modifiers,
             in_export_context=self.in_export_context,
         )
 

posthog/hogql_queries/hogql_query_runner.py

@@ -1,5 +1,4 @@
 from datetime import timedelta
-from typing import Dict, Optional, Any
 
 from posthog.clickhouse.client.connection import Workload
 from posthog.hogql import ast
@@ -9,7 +8,6 @@
 from posthog.hogql.query import execute_hogql_query
 from posthog.hogql.timings import HogQLTimings
 from posthog.hogql_queries.query_runner import QueryRunner
-from posthog.models import Team
 from posthog.schema import (
     HogQLQuery,
     HogQLQueryResponse,
@@ -23,19 +21,6 @@ class HogQLQueryRunner(QueryRunner):
     query: HogQLQuery
     query_type = HogQLQuery
 
-    def __init__(
-        self,
-        query: HogQLQuery | Dict[str, Any],
-        team: Team,
-        timings: Optional[HogQLTimings] = None,
-        in_export_context: Optional[bool] = False,
-    ):
-        super().__init__(query, team, timings, in_export_context)
-        if isinstance(query, HogQLQuery):
-            self.query = query
-        else:
-            self.query = HogQLQuery.model_validate(query)
-
     def to_query(self) -> ast.SelectQuery:
         if self.timings is None:
             self.timings = HogQLTimings()
@@ -60,7 +45,7 @@ def calculate(self) -> HogQLQueryResponse:
             query_type="HogQLQuery",
             query=self.to_query(),
             filters=self.query.filters,
-            modifiers=self.query.modifiers,
+            modifiers=self.query.modifiers or self.modifiers,
             team=self.team,
             workload=Workload.ONLINE,
             timings=self.timings,

posthog/hogql_queries/insights/insight_persons_query_runner.py

@@ -1,13 +1,11 @@
 from datetime import timedelta
-from typing import Dict, Optional, Any, cast
+from typing import cast
 
 from posthog.hogql import ast
 from posthog.hogql.query import execute_hogql_query
-from posthog.hogql.timings import HogQLTimings
 from posthog.hogql_queries.insights.lifecycle_query_runner import LifecycleQueryRunner
 from posthog.hogql_queries.insights.trends.trends_query_runner import TrendsQueryRunner
 from posthog.hogql_queries.query_runner import QueryRunner, get_query_runner
-from posthog.models import Team
 from posthog.models.filters.mixins.utils import cached_property
 from posthog.schema import InsightPersonsQuery, HogQLQueryResponse
 
@@ -16,19 +14,6 @@ class InsightPersonsQueryRunner(QueryRunner):
     query: InsightPersonsQuery
     query_type = InsightPersonsQuery
 
-    def __init__(
-        self,
-        query: InsightPersonsQuery | Dict[str, Any],
-        team: Team,
-        timings: Optional[HogQLTimings] = None,
-        in_export_context: Optional[bool] = False,
-    ):
-        super().__init__(query, team, timings, in_export_context)
-        if isinstance(query, InsightPersonsQuery):
-            self.query = query
-        else:
-            self.query = InsightPersonsQuery.model_validate(query)
-
     @cached_property
     def source_runner(self) -> QueryRunner:
         return get_query_runner(self.query.source, self.team, self.timings, self.in_export_context)
@@ -54,6 +39,7 @@ def calculate(self) -> HogQLQueryResponse:
             query=self.to_query(),
             team=self.team,
             timings=self.timings,
+            modifiers=self.modifiers,
         )
 
     def _is_stale(self, cached_result_package):

posthog/hogql_queries/insights/lifecycle_query_runner.py

@@ -1,6 +1,6 @@
 from datetime import timedelta
 from math import ceil
-from typing import Optional, Any, Dict, List
+from typing import Optional, List
 
 from django.utils.timezone import datetime
 from posthog.caching.insights_api import (
@@ -14,9 +14,8 @@
 from posthog.hogql.printer import to_printed_hogql
 from posthog.hogql.property import property_to_expr, action_to_expr
 from posthog.hogql.query import execute_hogql_query
-from posthog.hogql.timings import HogQLTimings
 from posthog.hogql_queries.query_runner import QueryRunner
-from posthog.models import Team, Action
+from posthog.models import Action
 from posthog.hogql_queries.utils.query_date_range import QueryDateRange
 from posthog.models.filters.mixins.utils import cached_property
 from posthog.schema import (
@@ -31,15 +30,6 @@ class LifecycleQueryRunner(QueryRunner):
     query: LifecycleQuery
     query_type = LifecycleQuery
 
-    def __init__(
-        self,
-        query: LifecycleQuery | Dict[str, Any],
-        team: Team,
-        timings: Optional[HogQLTimings] = None,
-        in_export_context: Optional[bool] = False,
-    ):
-        super().__init__(query, team, timings, in_export_context)
-
     def to_query(self) -> ast.SelectQuery | ast.SelectUnionQuery:
         if self.query.samplingFactor == 0:
             counts_with_sampling = ast.Constant(value=0)
@@ -139,6 +129,7 @@ def calculate(self) -> LifecycleQueryResponse:
             query=query,
             team=self.team,
             timings=self.timings,
+            modifiers=self.modifiers,
         )
 
         # TODO: can we move the data conversion part into the query as well? It would make it easier to swap

posthog/hogql_queries/insights/trends/trends_query_runner.py

@@ -35,6 +35,7 @@
     HogQLQueryResponse,
     TrendsQuery,
     TrendsQueryResponse,
+    HogQLQueryModifiers,
 )
 
 
@@ -48,9 +49,10 @@ def __init__(
         query: TrendsQuery | Dict[str, Any],
         team: Team,
         timings: Optional[HogQLTimings] = None,
+        modifiers: Optional[HogQLQueryModifiers] = None,
         in_export_context: Optional[bool] = None,
     ):
-        super().__init__(query, team, timings, in_export_context)
+        super().__init__(query, team=team, timings=timings, modifiers=modifiers, in_export_context=in_export_context)
         self.series = self.setup_series()
 
     def _is_stale(self, cached_result_package):
@@ -129,6 +131,7 @@ def calculate(self):
                 query=query,
                 team=self.team,
                 timings=self.timings,
+                modifiers=self.modifiers,
             )
 
             timings.extend(response.timings)

posthog/hogql_queries/persons_query_runner.py

@@ -1,15 +1,13 @@
 import json
 from datetime import timedelta
-from typing import Optional, Any, Dict, List, cast, Literal
+from typing import List, cast, Literal
 
 from posthog.hogql import ast
 from posthog.hogql.constants import DEFAULT_RETURNED_ROWS, MAX_SELECT_RETURNED_ROWS
 from posthog.hogql.parser import parse_expr, parse_order_expr
 from posthog.hogql.property import property_to_expr, has_aggregation
 from posthog.hogql.query import execute_hogql_query
-from posthog.hogql.timings import HogQLTimings
 from posthog.hogql_queries.query_runner import QueryRunner, get_query_runner
-from posthog.models import Team
 from posthog.schema import PersonsQuery, PersonsQueryResponse
 
 PERSON_FULL_TUPLE = ["id", "properties", "created_at", "is_identified"]
@@ -19,25 +17,13 @@ class PersonsQueryRunner(QueryRunner):
     query: PersonsQuery
     query_type = PersonsQuery
 
-    def __init__(
-        self,
-        query: PersonsQuery | Dict[str, Any],
-        team: Team,
-        timings: Optional[HogQLTimings] = None,
-        in_export_context: Optional[bool] = False,
-    ):
-        super().__init__(query=query, team=team, timings=timings, in_export_context=in_export_context)
-        if isinstance(query, PersonsQuery):
-            self.query = query
-        else:
-            self.query = PersonsQuery.model_validate(query)
-
     def calculate(self) -> PersonsQueryResponse:
         response = execute_hogql_query(
             query_type="PersonsQuery",
             query=self.to_query(),
             team=self.team,
             timings=self.timings,
+            modifiers=self.modifiers,
         )
         input_columns = self.input_columns()
         if "person" in input_columns:

posthog/hogql_queries/query_runner.py

@@ -28,6 +28,7 @@
     HogQLQuery,
     InsightPersonsQuery,
     DashboardFilter,
+    HogQLQueryModifiers,
 )
 from posthog.utils import generate_cache_key, get_safe_cache
 
@@ -86,6 +87,7 @@ def get_query_runner(
     team: Team,
     timings: Optional[HogQLTimings] = None,
     in_export_context: Optional[bool] = False,
+    modifiers: Optional[HogQLQueryModifiers] = None,
 ) -> "QueryRunner":
     kind = None
     if isinstance(query, dict):
@@ -103,6 +105,7 @@ def get_query_runner(
             team=team,
             timings=timings,
             in_export_context=in_export_context,
+            modifiers=modifiers,
         )
     if kind == "TrendsQuery":
         from .insights.trends.trends_query_runner import TrendsQueryRunner
@@ -112,6 +115,7 @@ def get_query_runner(
             team=team,
             timings=timings,
             in_export_context=in_export_context,
+            modifiers=modifiers,
         )
     if kind == "EventsQuery":
         from .events_query_runner import EventsQueryRunner
@@ -121,6 +125,7 @@ def get_query_runner(
             team=team,
             timings=timings,
             in_export_context=in_export_context,
+            modifiers=modifiers,
         )
     if kind == "PersonsQuery":
         from .persons_query_runner import PersonsQueryRunner
@@ -130,6 +135,7 @@ def get_query_runner(
             team=team,
             timings=timings,
             in_export_context=in_export_context,
+            modifiers=modifiers,
         )
     if kind == "InsightPersonsQuery":
         from .insights.insight_persons_query_runner import InsightPersonsQueryRunner
@@ -139,6 +145,7 @@ def get_query_runner(
             team=team,
             timings=timings,
             in_export_context=in_export_context,
+            modifiers=modifiers,
         )
     if kind == "HogQLQuery":
         from .hogql_query_runner import HogQLQueryRunner
@@ -148,6 +155,7 @@ def get_query_runner(
             team=team,
             timings=timings,
             in_export_context=in_export_context,
+            modifiers=modifiers,
         )
     if kind == "SessionsTimelineQuery":
         from .sessions_timeline_query_runner import SessionsTimelineQueryRunner
@@ -156,19 +164,20 @@ def get_query_runner(
             query=cast(SessionsTimelineQuery | Dict[str, Any], query),
             team=team,
             timings=timings,
+            modifiers=modifiers,
         )
     if kind == "WebOverviewQuery":
         from .web_analytics.web_overview import WebOverviewQueryRunner
 
-        return WebOverviewQueryRunner(query=query, team=team, timings=timings)
+        return WebOverviewQueryRunner(query=query, team=team, timings=timings, modifiers=modifiers)
     if kind == "WebTopClicksQuery":
         from .web_analytics.top_clicks import WebTopClicksQueryRunner
 
-        return WebTopClicksQueryRunner(query=query, team=team, timings=timings)
+        return WebTopClicksQueryRunner(query=query, team=team, timings=timings, modifiers=modifiers)
     if kind == "WebStatsTableQuery":
         from .web_analytics.stats_table import WebStatsTableQueryRunner
 
-        return WebStatsTableQueryRunner(query=query, team=team, timings=timings)
+        return WebStatsTableQueryRunner(query=query, team=team, timings=timings, modifiers=modifiers)
 
     raise ValueError(f"Can't get a runner for an unknown query kind: {kind}")
 
@@ -178,18 +187,21 @@ class QueryRunner(ABC):
     query_type: Type[RunnableQueryNode]
     team: Team
     timings: HogQLTimings
+    modifiers: HogQLQueryModifiers
     in_export_context: bool
 
     def __init__(
         self,
         query: RunnableQueryNode | BaseModel | Dict[str, Any],
         team: Team,
         timings: Optional[HogQLTimings] = None,
+        modifiers: Optional[HogQLQueryModifiers] = None,
         in_export_context: Optional[bool] = False,
     ):
         self.team = team
         self.timings = timings or HogQLTimings()
         self.in_export_context = in_export_context or False
+        self.modifiers = create_default_modifiers_for_team(team, modifiers)
         if isinstance(query, self.query_type):
             self.query = query  # type: ignore
         else:
@@ -244,7 +256,7 @@ def to_hogql(self) -> str:
                     team_id=self.team.pk,
                     enable_select_queries=True,
                     timings=self.timings,
-                    modifiers=create_default_modifiers_for_team(self.team),
+                    modifiers=self.modifiers,
                 ),
                 "hogql",
             )
@@ -253,8 +265,9 @@ def toJSON(self) -> str:
         return self.query.model_dump_json(exclude_defaults=True, exclude_none=True)
 
     def _cache_key(self) -> str:
+        modifiers = self.modifiers.model_dump_json(exclude_defaults=True, exclude_none=True)
         return generate_cache_key(
-            f"query_{self.toJSON()}_{self.__class__.__name__}_{self.team.pk}_{self.team.timezone}"
+            f"query_{self.toJSON()}_{self.__class__.__name__}_{self.team.pk}_{self.team.timezone}_{modifiers}"
         )
 
     @abstractmethod