From 92bfec6615a5869f093edfd04b37ce055c275e61 Mon Sep 17 00:00:00 2001
From: Michael Matloka <michal@matloka.com>
Date: Tue, 19 Mar 2024 20:11:17 +0100
Subject: [PATCH] Improve typing

---
 posthog/api/query.py                       | 2 +-
 posthog/clickhouse/client/execute_async.py | 9 +++++++--
 posthog/tasks/tasks.py                     | 9 +++++----
 3 files changed, 13 insertions(+), 7 deletions(-)

diff --git a/posthog/api/query.py b/posthog/api/query.py
index 19b70677397e3..e30853655c749 100644
--- a/posthog/api/query.py
+++ b/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)
 
diff --git a/posthog/clickhouse/client/execute_async.py b/posthog/clickhouse/client/execute_async.py
index fc7e2cc4c84c6..4671b0060299b 100644
--- a/posthog/clickhouse/client/execute_async.py
+++ b/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,
diff --git a/posthog/tasks/tasks.py b/posthog/tasks/tasks.py
index 721eb92ee31ac..78895c28d7b90 100644
--- a/posthog/tasks/tasks.py
+++ b/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:
     """