From 681d703c64d9dc47a049a05332f42d24516f3ec9 Mon Sep 17 00:00:00 2001 From: Julian Bez Date: Mon, 20 Nov 2023 12:55:49 +0100 Subject: [PATCH 1/3] Set ClickHouse max execution time higher for async queries --- posthog/clickhouse/client/execute_async.py | 7 ++----- posthog/hogql/printer.py | 9 +++++++++ posthog/hogql/query.py | 1 + 3 files changed, 12 insertions(+), 5 deletions(-) diff --git a/posthog/clickhouse/client/execute_async.py b/posthog/clickhouse/client/execute_async.py index fc9e292b08ee4..7e42d52d4836c 100644 --- a/posthog/clickhouse/client/execute_async.py +++ b/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) diff --git a/posthog/hogql/printer.py b/posthog/hogql/printer.py index f89614d0dc95a..14a3db835a816 100644 --- a/posthog/hogql/printer.py +++ b/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: diff --git a/posthog/hogql/query.py b/posthog/hogql/query.py index c7e8c82713b15..37994f529dc6c 100644 --- a/posthog/hogql/query.py +++ b/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() From 3e4c1ee87917c115ec5ed061fd882e85d18b344e Mon Sep 17 00:00:00 2001 From: Julian Bez Date: Mon, 20 Nov 2023 14:31:08 +0100 Subject: [PATCH 2/3] Fix setting settings, add snapshot for test --- .../client/test/__snapshots__/test_execute_async.ambr | 8 ++++++++ posthog/clickhouse/client/test/test_execute_async.py | 3 ++- posthog/hogql/printer.py | 2 +- 3 files changed, 11 insertions(+), 2 deletions(-) create mode 100644 posthog/clickhouse/client/test/__snapshots__/test_execute_async.ambr diff --git a/posthog/clickhouse/client/test/__snapshots__/test_execute_async.ambr b/posthog/clickhouse/client/test/__snapshots__/test_execute_async.ambr new file mode 100644 index 0000000000000..282191d2015c7 --- /dev/null +++ b/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 + ' +--- diff --git a/posthog/clickhouse/client/test/test_execute_async.py b/posthog/clickhouse/client/test/test_execute_async.py index 1ab4bf49e03d3..4958c23b3f0a0 100644 --- a/posthog/clickhouse/client/test/test_execute_async.py +++ b/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 diff --git a/posthog/hogql/printer.py b/posthog/hogql/printer.py index 14a3db835a816..9f16643437493 100644 --- a/posthog/hogql/printer.py +++ b/posthog/hogql/printer.py @@ -1109,7 +1109,7 @@ def _is_nullable(self, node: ast.Expr) -> bool: def _print_settings(self, settings): if self.timeout: - settings["max_execution_time"] = self.timeout + settings.max_execution_time = self.timeout pairs = [] for key, value in settings: From 41f07a3276b5d629fb5f269425fbd200cd6304a8 Mon Sep 17 00:00:00 2001 From: Julian Bez Date: Tue, 21 Nov 2023 09:17:40 +0100 Subject: [PATCH 3/3] Use export context to directly modify max_execution_time via settings --- posthog/hogql/printer.py | 9 --------- posthog/hogql/query.py | 9 +++++++-- 2 files changed, 7 insertions(+), 11 deletions(-) diff --git a/posthog/hogql/printer.py b/posthog/hogql/printer.py index 9f16643437493..f89614d0dc95a 100644 --- a/posthog/hogql/printer.py +++ b/posthog/hogql/printer.py @@ -73,7 +73,6 @@ 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( @@ -83,7 +82,6 @@ def print_ast( stack=stack, settings=settings, pretty=pretty, - timeout=timeout, ) @@ -127,7 +125,6 @@ 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 @@ -137,7 +134,6 @@ def print_prepared_ast( stack=stack or [], settings=settings, pretty=pretty, - timeout=timeout, ).visit(node) @@ -157,7 +153,6 @@ def __init__( stack: Optional[List[AST]] = None, settings: Optional[HogQLGlobalSettings] = None, pretty: bool = False, - timeout: Optional[int] = None, ): self.context = context self.dialect = dialect @@ -166,7 +161,6 @@ 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) @@ -1108,9 +1102,6 @@ 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: diff --git a/posthog/hogql/query.py b/posthog/hogql/query.py index 37994f529dc6c..751b9fb46b860 100644 --- a/posthog/hogql/query.py +++ b/posthog/hogql/query.py @@ -22,6 +22,8 @@ from posthog.client import sync_execute from posthog.schema import HogQLQueryResponse, HogQLFilters, HogQLQueryModifiers +EXPORT_CONTEXT_MAX_EXECUTION_TIME = 600 + def execute_hogql_query( query: Union[str, ast.SelectQuery, ast.SelectUnionQuery], @@ -119,6 +121,10 @@ def execute_hogql_query( ) ) + settings = settings or HogQLGlobalSettings() + if in_export_context: + settings.max_execution_time = EXPORT_CONTEXT_MAX_EXECUTION_TIME + # Print the ClickHouse SQL query with timings.measure("print_ast"): clickhouse_context = HogQLContext( @@ -131,8 +137,7 @@ def execute_hogql_query( select_query, context=clickhouse_context, dialect="clickhouse", - settings=settings or HogQLGlobalSettings(), - timeout=600 if in_export_context else None, + settings=settings, ) timings_dict = timings.to_dict()