feat(query-pipeline): Add configs to customize maximum size of a query (#7546)

Author: phacops

snuba/pipeline/stages/query_execution.py

@@ -11,6 +11,7 @@
 
 from snuba import environment
 from snuba import settings as snuba_settings
+from snuba import state
 from snuba.attribution.attribution_info import AttributionInfo
 from snuba.clickhouse.formatter.query import format_query
 from snuba.clickhouse.query import Query as ClickhouseQuery
@@ -45,6 +46,11 @@
 metrics = MetricsWrapper(environment.metrics, "api")
 logger = logging.getLogger("snuba.pipeline.stages.query_execution")
 
+DISABLE_MAX_QUERY_SIZE_CHECK_FOR_CLUSTERS_CONFIG = (
+    "ExecutionStage.disable_max_query_size_check_for_clusters"
+)
+MAX_QUERY_SIZE_BYTES_CONFIG = "ExecutionStage.max_query_size_bytes"
+
 
 class ExecutionStage(QueryPipelineStage[ClickhouseQuery | CompositeQuery[Table], QueryResult]):
     def __init__(
@@ -157,6 +163,21 @@ def _run_and_apply_column_names(
     return result
 
 
+def _max_query_size_bytes() -> int:
+    return (
+        state.get_int_config(MAX_QUERY_SIZE_BYTES_CONFIG, MAX_QUERY_SIZE_BYTES)
+        or MAX_QUERY_SIZE_BYTES
+    )
+
+
+def _disable_max_query_size_check_for_clusters() -> set[str]:
+    return set(
+        (state.get_str_config(DISABLE_MAX_QUERY_SIZE_CHECK_FOR_CLUSTERS_CONFIG, "") or "").split(
+            ","
+        )
+    )
+
+
 def _format_storage_query_and_run(
     timer: Timer,
     query_metadata: SnubaQueryMetadata,
@@ -207,11 +228,16 @@ def _format_storage_query_and_run(
         "cluster_name": cluster_name,
     }
 
-    if query_size_bytes > MAX_QUERY_SIZE_BYTES:
+    if (
+        not cluster_name
+        or
+        # This will force to fallback on the ClickHouse limit.
+        cluster_name not in _disable_max_query_size_check_for_clusters()
+    ) and query_size_bytes > _max_query_size_bytes():
         cause = QueryTooLongException(
             f"After processing, query is {query_size_bytes} bytes, "
             "which is too long for ClickHouse to process. "
-            f"Max size is {MAX_QUERY_SIZE_BYTES} bytes."
+            f"Max size is {_max_query_size_bytes()} bytes."
         )
         stats = update_query_metadata_and_stats(
             query=clickhouse_query,
@@ -273,10 +299,10 @@ def get_query_size_group(query_size_bytes: int) -> str:
     Eg. If the query size is equal to the max query size, this function
     returns "100%".
     """
-    if query_size_bytes == MAX_QUERY_SIZE_BYTES:
+    if query_size_bytes == _max_query_size_bytes():
         return "100%"
     else:
-        query_size_group = int(floor(query_size_bytes / MAX_QUERY_SIZE_BYTES * 10)) * 10
+        query_size_group = int(floor(query_size_bytes / _max_query_size_bytes() * 10)) * 10
         return f">={query_size_group}%"
 
 

snuba/settings/__init__.py

@@ -119,6 +119,7 @@
             "profile_chunks",
         },
         "single_node": True,
+        "cluster_name": "test_cluster",
     },
 ]
 

tests/pipeline/test_execution_stage.py

@@ -3,6 +3,7 @@
 import pytest
 
 from snuba import settings as snubasettings
+from snuba import state
 from snuba.attribution import get_app_id
 from snuba.attribution.attribution_info import AttributionInfo
 from snuba.clickhouse.columns import ColumnSet
@@ -11,7 +12,11 @@
 from snuba.datasets.entities.entity_key import EntityKey
 from snuba.datasets.storages.storage_key import StorageKey
 from snuba.pipeline.query_pipeline import QueryPipelineResult
-from snuba.pipeline.stages.query_execution import ExecutionStage
+from snuba.pipeline.stages.query_execution import (
+    DISABLE_MAX_QUERY_SIZE_CHECK_FOR_CLUSTERS_CONFIG,
+    MAX_QUERY_SIZE_BYTES_CONFIG,
+    ExecutionStage,
+)
 from snuba.query import SelectedExpression
 from snuba.query.allocation_policies import (
     MAX_THRESHOLD,
@@ -35,6 +40,7 @@
 from snuba.request import Request
 from snuba.utils.metrics.timer import Timer
 from snuba.utils.schemas import UUID, String, UInt
+from snuba.web import QueryException
 
 
 class MockAllocationPolicy(AllocationPolicy):
@@ -72,9 +78,7 @@ def get_fake_metadata() -> SnubaQueryMetadata:
         Request(
             uuid.uuid4(),
             {},
-            LogicalQuery(
-                from_clause=Entity(key=EntityKey.TRANSACTIONS, schema=ColumnSet([]))
-            ),
+            LogicalQuery(from_clause=Entity(key=EntityKey.TRANSACTIONS, schema=ColumnSet([]))),
             HTTPQuerySettings(),
             AttributionInfo(
                 get_app_id("blah"),
@@ -220,3 +224,61 @@ def test_turbo(ch_query: Query) -> None:
         and "avg(duration)" in res.data.result["data"][0]
     )
     assert ch_query.get_from_clause().sampling_rate == snubasettings.TURBO_SAMPLE_RATE
+
+
+@pytest.mark.clickhouse_db
+@pytest.mark.redis_db
+def test_max_query_size_bytes(ch_query: Query) -> None:
+    attinfo = AttributionInfo(
+        get_app_id("blah"), {"tenant_type": "tenant_id"}, "blah", None, None, None
+    )
+    settings = HTTPQuerySettings()
+    timer = Timer("test")
+    metadata = get_fake_metadata()
+
+    state.set_config(MAX_QUERY_SIZE_BYTES_CONFIG, 1)
+
+    res = ExecutionStage(attinfo, query_metadata=metadata).execute(
+        QueryPipelineResult(
+            data=ch_query,
+            query_settings=settings,
+            timer=timer,
+            error=None,
+        )
+    )
+
+    assert res.data is None
+    assert isinstance(res.error, QueryException)
+    assert "which is too long for ClickHouse to process" in res.error.message
+
+
+@pytest.mark.clickhouse_db
+@pytest.mark.redis_db
+def test_disable_max_query_size_check(ch_query: Query) -> None:
+    attinfo = AttributionInfo(
+        get_app_id("blah"), {"tenant_type": "tenant_id"}, "blah", None, None, None
+    )
+    settings = HTTPQuerySettings()
+    timer = Timer("test")
+    metadata = get_fake_metadata()
+    cluster_name = (
+        snubasettings.CLUSTERS[0]["cluster_name"]
+        if "cluster_name" in snubasettings.CLUSTERS[0]
+        else "test_cluster"
+    )
+
+    # Lowering this should make the query too big...
+    state.set_config(MAX_QUERY_SIZE_BYTES_CONFIG, 1)
+    # Unless we disable the check for this cluster.
+    state.set_config(DISABLE_MAX_QUERY_SIZE_CHECK_FOR_CLUSTERS_CONFIG, cluster_name)
+
+    res = ExecutionStage(attinfo, query_metadata=metadata).execute(
+        QueryPipelineResult(
+            data=ch_query,
+            query_settings=settings,
+            timer=timer,
+            error=None,
+        )
+    )
+
+    assert res.data