get cluster after checking storage availability

getsentry · Jun 20, 2023 · 323a215 · 323a215
1 parent 206f4af
commit 323a215
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 1 deletion.
diff --git 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
@@ -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(