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/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/query.py b/posthog/hogql/query.py index c7e8c82713b15..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,7 +137,7 @@ def execute_hogql_query( select_query, context=clickhouse_context, dialect="clickhouse", - settings=settings or HogQLGlobalSettings(), + settings=settings, ) timings_dict = timings.to_dict()