feat(queries): Set ClickHouse max execution time higher for async queries

posthog/clickhouse/client/execute_async.py

@@ -81,7 +81,6 @@ def enqueue_process_query_task(
     query_json,
     query_id=None,
     refresh_requested=False,
-    in_export_context=False,
     bypass_celery=False,
     force=False,
 ):
@@ -115,12 +114,10 @@ def enqueue_process_query_task(
 
     if bypass_celery:
         # Call directly ( for testing )
-        process_query_task(
-            team_id, query_id, query_json, in_export_context=in_export_context, refresh_requested=refresh_requested
-        )
+        process_query_task(team_id, query_id, query_json, in_export_context=True, refresh_requested=refresh_requested)
     else:
         task = process_query_task.delay(
-            team_id, query_id, query_json, in_export_context=in_export_context, refresh_requested=refresh_requested
+            team_id, query_id, query_json, in_export_context=True, refresh_requested=refresh_requested
         )
         query_status.task_id = task.id
         redis_client.set(key, query_status.model_dump_json(), ex=REDIS_STATUS_TTL_SECONDS)

posthog/clickhouse/client/test/__snapshots__/test_execute_async.ambr

@@ -0,0 +1,8 @@
+# name: ClickhouseClientTestCase.test_async_query_client
+  '
+  SELECT plus(1, 1)
+  LIMIT 10000 SETTINGS readonly=2,
+                       max_execution_time=600,
+                       allow_experimental_object_type=1
+  '
+---

posthog/clickhouse/client/test/test_execute_async.py

@@ -7,7 +7,7 @@
 from posthog.client import sync_execute
 from posthog.hogql.errors import HogQLException
 from posthog.models import Organization, Team
-from posthog.test.base import ClickhouseTestMixin
+from posthog.test.base import ClickhouseTestMixin, snapshot_clickhouse_queries
 
 
 def build_query(sql):
@@ -23,6 +23,7 @@ def setUp(self):
         self.team = Team.objects.create(organization=self.organization)
         self.team_id = self.team.pk
 
+    @snapshot_clickhouse_queries
     def test_async_query_client(self):
         query = build_query("SELECT 1+1")
         team_id = self.team_id

posthog/hogql/printer.py

@@ -73,6 +73,7 @@ def print_ast(
     stack: Optional[List[ast.SelectQuery]] = None,
     settings: Optional[HogQLGlobalSettings] = None,
     pretty: bool = False,
+    timeout: Optional[int] = None,
 ) -> str:
     prepared_ast = prepare_ast_for_printing(node=node, context=context, dialect=dialect, stack=stack, settings=settings)
     return print_prepared_ast(
@@ -82,6 +83,7 @@ def print_ast(
         stack=stack,
         settings=settings,
         pretty=pretty,
+        timeout=timeout,
     )
 
 
@@ -125,6 +127,7 @@ def print_prepared_ast(
     stack: Optional[List[ast.SelectQuery]] = None,
     settings: Optional[HogQLGlobalSettings] = None,
     pretty: bool = False,
+    timeout: Optional[int] = None,
 ) -> str:
     with context.timings.measure("printer"):
         # _Printer also adds a team_id guard if printing clickhouse
@@ -134,6 +137,7 @@ def print_prepared_ast(
             stack=stack or [],
             settings=settings,
             pretty=pretty,
+            timeout=timeout,
         ).visit(node)
 
 
@@ -153,6 +157,7 @@ def __init__(
         stack: Optional[List[AST]] = None,
         settings: Optional[HogQLGlobalSettings] = None,
         pretty: bool = False,
+        timeout: Optional[int] = None,
     ):
         self.context = context
         self.dialect = dialect
@@ -161,6 +166,7 @@ def __init__(
         self.pretty = pretty
         self._indent = -1
         self.tab_size = 4
+        self.timeout = timeout
 
     def indent(self, extra: int = 0):
         return " " * self.tab_size * (self._indent + extra)
@@ -1102,6 +1108,9 @@ def _is_nullable(self, node: ast.Expr) -> bool:
         return True
 
     def _print_settings(self, settings):
+        if self.timeout:
+            settings.max_execution_time = self.timeout
+
         pairs = []
         for key, value in settings:
             if value is None:

posthog/hogql/query.py

@@ -132,6 +132,7 @@ def execute_hogql_query(
             context=clickhouse_context,
             dialect="clickhouse",
             settings=settings or HogQLGlobalSettings(),
+            timeout=600 if in_export_context else None,
         )
 
     timings_dict = timings.to_dict()