From 323a2151dcb83ac98bdf4689cf84dff4f510850d Mon Sep 17 00:00:00 2001 From: Enoch Tang Date: Tue, 20 Jun 2023 11:54:43 -0400 Subject: [PATCH] get cluster after checking storage availability --- snuba/datasets/plans/storage_plan_builder.py | 4 +++- tests/datasets/plans/test_storage_plan_builder.py | 1 + 2 files changed, 4 insertions(+), 1 deletion(-) diff --git a/snuba/datasets/plans/storage_plan_builder.py b/snuba/datasets/plans/storage_plan_builder.py index d8b8ba5cce..389f5a7ebe 100644 --- a/snuba/datasets/plans/storage_plan_builder.py +++ b/snuba/datasets/plans/storage_plan_builder.py @@ -178,7 +178,6 @@ def build_and_rank_plans( storage_connection = self.get_storage(query, settings) storage = storage_connection.storage mappers = storage_connection.translation_mappers - cluster = self.get_cluster(storage, query, settings) # Return failure if storage readiness state is not supported in current environment if snuba_settings.READINESS_STATE_FAIL_QUERIES: @@ -190,6 +189,9 @@ def build_and_rank_plans( f"The selected storage={storage.get_storage_key().value} is not available in this environment yet. To enable it, consider bumping the storage's readiness_state.", ) + # If storage is supported in this environment, we can safely check get the existing cluster. + cluster = self.get_cluster(storage, query, settings) + with sentry_sdk.start_span( op="build_plan.storage_query_plan_builder", description="translate" ): diff --git a/tests/datasets/plans/test_storage_plan_builder.py b/tests/datasets/plans/test_storage_plan_builder.py index bac0e987cb..4c58a5d7ea 100644 --- a/tests/datasets/plans/test_storage_plan_builder.py +++ b/tests/datasets/plans/test_storage_plan_builder.py @@ -165,6 +165,7 @@ def test_storage_unavailable_error_in_plan_builder(temp_settings: Any) -> None: query, _ = parse_snql_query(str(snql_query), dataset) temp_settings.SUPPORTED_STATES = {} # remove all supported states + temp_settings.READINESS_STATE_FAIL_QUERIES = True assert isinstance(query, Query) with pytest.raises(