Fix LimitContext causing too high default limit with pagination

PostHog · Dec 21, 2023 · cc08544 · cc08544
1 parent 8865be4
commit cc08544
Show file tree

Hide file tree

Showing 3 changed files with 8 additions and 7 deletions.
diff --git a/posthog/clickhouse/client/execute_async.py b/posthog/clickhouse/client/execute_async.py
@@ -137,11 +137,11 @@ def enqueue_process_query_task(
     if bypass_celery:
         # Call directly ( for testing )
         process_query_task(
-            team_id, query_id, query_json, limit_context=LimitContext.EXPORT, refresh_requested=refresh_requested
+            team_id, query_id, query_json, limit_context=LimitContext.QUERY_ASYNC, refresh_requested=refresh_requested
         )
     else:
         task = process_query_task.delay(
-            team_id, query_id, query_json, limit_context=LimitContext.EXPORT, refresh_requested=refresh_requested
+            team_id, query_id, query_json, limit_context=LimitContext.QUERY_ASYNC, refresh_requested=refresh_requested
         )
         query_status.task_id = task.id
         manager.store_query_status(query_status)

diff --git a/posthog/hogql/constants.py b/posthog/hogql/constants.py
@@ -39,14 +39,15 @@
 
 class LimitContext(str, Enum):
     QUERY = "query"
+    QUERY_ASYNC = "query_async"
     EXPORT = "export"
     COHORT_CALCULATION = "cohort_calculation"
 
 
 def get_max_limit_for_context(limit_context: LimitContext) -> int:
     if limit_context == LimitContext.EXPORT:
         return MAX_SELECT_RETURNED_ROWS  # 10k
-    elif limit_context == LimitContext.QUERY:
+    elif limit_context in (LimitContext.QUERY, LimitContext.QUERY_ASYNC):
         return MAX_SELECT_RETURNED_ROWS  # 10k
     elif limit_context == LimitContext.COHORT_CALCULATION:
         return MAX_SELECT_COHORT_CALCULATION_LIMIT  # 1b
@@ -58,7 +59,7 @@ def get_default_limit_for_context(limit_context: LimitContext) -> int:
     """Limit used if no limit is provided"""
     if limit_context == LimitContext.EXPORT:
         return MAX_SELECT_RETURNED_ROWS  # 10k
-    elif limit_context == LimitContext.QUERY:
+    elif limit_context in (LimitContext.QUERY, LimitContext.QUERY_ASYNC):
         return DEFAULT_RETURNED_ROWS  # 100
     elif limit_context == LimitContext.COHORT_CALCULATION:
         return MAX_SELECT_COHORT_CALCULATION_LIMIT  # 1b

diff --git a/posthog/hogql/query.py b/posthog/hogql/query.py
@@ -22,7 +22,7 @@
 from posthog.client import sync_execute
 from posthog.schema import HogQLQueryResponse, HogQLFilters, HogQLQueryModifiers
 
-EXPORT_CONTEXT_MAX_EXECUTION_TIME = 600
+INCREASED_MAX_EXECUTION_TIME = 600
 
 
 def execute_hogql_query(
@@ -114,8 +114,8 @@ def execute_hogql_query(
                     )
 
     settings = settings or HogQLGlobalSettings()
-    if limit_context == LimitContext.EXPORT or limit_context == LimitContext.COHORT_CALCULATION:
-        settings.max_execution_time = EXPORT_CONTEXT_MAX_EXECUTION_TIME
+    if limit_context in (LimitContext.EXPORT, LimitContext.COHORT_CALCULATION, LimitContext.QUERY_ASYNC):
+        settings.max_execution_time = INCREASED_MAX_EXECUTION_TIME
 
     # Print the ClickHouse SQL query
     with timings.measure("print_ast"):