Improve query runner type names

PostHog · Sep 29, 2023 · 9698c04 · 9698c04
1 parent dbf41d8
commit 9698c04
Show file tree

Hide file tree

Showing 3 changed files with 31 additions and 33 deletions.
diff --git a/posthog/hogql_queries/query_runner.py b/posthog/hogql_queries/query_runner.py
@@ -1,6 +1,6 @@
 from abc import ABC, abstractmethod
 from datetime import datetime
-from typing import Any, Generic, List, Optional, Type, Dict, TypeVar
+from typing import Any, Generic, List, Optional, Type, Dict, TypeVar, Union
 
 from django.conf import settings
 from django.core.cache import cache
@@ -14,9 +14,18 @@
 from posthog.hogql.timings import HogQLTimings
 from posthog.metrics import LABEL_TEAM_ID
 from posthog.models import Team
-from posthog.schema import QueryTiming
-
-from posthog.types import Node
+from posthog.schema import (
+    QueryTiming,
+    TrendsQuery,
+    FunnelsQuery,
+    RetentionQuery,
+    PathsQuery,
+    StickinessQuery,
+    LifecycleQuery,
+    WebTopSourcesQuery,
+    WebTopClicksQuery,
+    WebTopPagesQuery,
+)
 from posthog.utils import generate_cache_key, get_safe_cache
 
 QUERY_CACHE_WRITE_COUNTER = Counter(
@@ -53,13 +62,26 @@ class CachedQueryResponse(QueryResponse):
     next_allowed_client_refresh: str
 
 
+RunnableQueryNode = Union[
+    TrendsQuery,
+    FunnelsQuery,
+    RetentionQuery,
+    PathsQuery,
+    StickinessQuery,
+    LifecycleQuery,
+    WebTopSourcesQuery,
+    WebTopClicksQuery,
+    WebTopPagesQuery,
+]
+
+
 class QueryRunner(ABC):
-    query: Node
-    query_type: Type[Node]
+    query: RunnableQueryNode
+    query_type: Type[RunnableQueryNode]
     team: Team
     timings: HogQLTimings
 
-    def __init__(self, query: Node | Dict[str, Any], team: Team, timings: Optional[HogQLTimings] = None):
+    def __init__(self, query: RunnableQueryNode | Dict[str, Any], team: Team, timings: Optional[HogQLTimings] = None):
         self.team = team
         self.timings = timings or HogQLTimings()
         if isinstance(query, self.query_type):

diff --git a/posthog/hogql_queries/test/test_query_runner.py b/posthog/hogql_queries/test/test_query_runner.py
@@ -6,10 +6,9 @@
 from freezegun import freeze_time
 from pydantic import BaseModel
 
-from posthog.hogql_queries.query_runner import QueryResponse, QueryRunner
+from posthog.hogql_queries.query_runner import QueryResponse, QueryRunner, RunnableQueryNode
 from posthog.models.team.team import Team
 from posthog.test.base import BaseTest
-from posthog.types import Node
 
 
 class TestQuery(BaseModel):
@@ -19,7 +18,7 @@ class TestQuery(BaseModel):
 
 
 class QueryRunnerTest(BaseTest):
-    def setup_test_query_runner_class(self, query_class: Type[Node] = TestQuery):  # type: ignore
+    def setup_test_query_runner_class(self, query_class: Type[RunnableQueryNode] = TestQuery):  # type: ignore
         """Setup required methods and attributes of the abstract base class."""
 
         class TestQueryRunner(QueryRunner):

diff --git a/posthog/types.py b/posthog/types.py
@@ -4,28 +4,5 @@
 from posthog.models.filters.path_filter import PathFilter
 from posthog.models.filters.retention_filter import RetentionFilter
 from posthog.models.filters.stickiness_filter import StickinessFilter
-from posthog.schema import (
-    FunnelsQuery,
-    LifecycleQuery,
-    PathsQuery,
-    RetentionQuery,
-    StickinessQuery,
-    TrendsQuery,
-    WebTopSourcesQuery,
-    WebTopClicksQuery,
-    WebTopPagesQuery,
-)
 
 FilterType = Union[Filter, PathFilter, RetentionFilter, StickinessFilter]
-
-Node = Union[
-    TrendsQuery,
-    FunnelsQuery,
-    RetentionQuery,
-    PathsQuery,
-    StickinessQuery,
-    LifecycleQuery,
-    WebTopSourcesQuery,
-    WebTopClicksQuery,
-    WebTopPagesQuery,
-]