Improve typing

posthog/api/query.py

@@ -70,7 +70,7 @@ def create(self, request, *args, **kwargs) -> Response:
                 user_id=self.request.user.pk,
                 query_json=request.data["query"],
                 query_id=client_query_id,
-                refresh_requested=data.refresh,
+                refresh_requested=data.refresh or False,
             )
             return Response(query_status.model_dump(), status=status.HTTP_202_ACCEPTED)
 

posthog/clickhouse/client/execute_async.py

@@ -70,7 +70,12 @@ def delete_query_status(self):
 
 
 def execute_process_query(
-    team_id: int, user_id: int, query_id: str, query_json: str, limit_context: LimitContext, refresh_requested: bool
+    team_id: int,
+    user_id: int,
+    query_id: str,
+    query_json: dict,
+    limit_context: Optional[LimitContext],
+    refresh_requested: bool,
 ):
     manager = QueryStatusManager(query_id, team_id)
 
@@ -112,7 +117,7 @@ def execute_process_query(
 def enqueue_process_query_task(
     team_id: int,
     user_id: int,
-    query_json: str,
+    query_json: dict,
     query_id: Optional[str] = None,
     refresh_requested: bool = False,
     force: bool = False,

posthog/tasks/tasks.py

@@ -1,5 +1,5 @@
 import time
-from typing import Any, Optional
+from typing import Optional
 from uuid import UUID
 
 from celery import shared_task
@@ -9,6 +9,7 @@
 from prometheus_client import Gauge
 
 from posthog.cloud_utils import is_cloud
+from posthog.hogql.constants import LimitContext
 from posthog.metrics import pushed_metrics_registry
 from posthog.ph_client import get_ph_client
 from posthog.redis import get_client
@@ -33,11 +34,11 @@ def redis_heartbeat() -> None:
 
 @shared_task(ignore_result=True, queue=CeleryQueue.ANALYTICS_QUERIES.value)
 def process_query_task(
-    team_id: str,
+    team_id: int,
     user_id: str,
     query_id: str,
-    query_json: Any,
-    limit_context: Any = None,
+    query_json: dict,
+    limit_context: Optional[LimitContext] = None,
     refresh_requested: bool = False,
 ) -> None:
     """