check connections on consumer startup (#4497)

* check connections on consumer startup

* make retries a setting

* allow only checking a set of clusters

* remove retries setting

snuba/cli/consumer.py

@@ -13,9 +13,10 @@
     ProcessingParameters,
 )
 from snuba.consumers.consumer_config import resolve_consumer_config
-from snuba.datasets.storages.factory import get_writable_storage_keys
+from snuba.datasets.storages.factory import get_storage, get_writable_storage_keys
 from snuba.datasets.storages.storage_key import StorageKey
 from snuba.environment import setup_logging, setup_sentry
+from snuba.migrations.connect import check_clickhouse_connections
 from snuba.utils.metrics.wrapper import MetricsWrapper
 from snuba.utils.streams.metrics_adapter import StreamMetricsAdapter
 
@@ -179,11 +180,17 @@ def consumer(
 
     setup_logging(log_level)
     setup_sentry()
+
     logger.info("Consumer Starting")
 
     storage_key = StorageKey(storage_name)
     sentry_sdk.set_tag("storage", storage_name)
 
+    logger.info("Checking Clickhouse connections")
+    storage = get_storage(storage_key)
+    cluster = storage.get_cluster()
+    check_clickhouse_connections([cluster])
+
     metrics_tags = {
         "consumer_group": consumer_group,
         "storage": storage_key.value,

snuba/migrations/connect.py

@@ -9,6 +9,7 @@
 from snuba.clusters.cluster import (
     CLUSTERS,
     ClickhouseClientSettings,
+    ClickhouseCluster,
     ClickhouseNode,
     UndefinedClickhouseCluster,
 )
@@ -25,13 +26,15 @@
 logger = structlog.get_logger().bind(module=__name__)
 
 
-def check_clickhouse_connections() -> None:
+def check_clickhouse_connections(
+    clusters: Sequence[ClickhouseCluster] = CLUSTERS,
+) -> None:
     """
     Ensure that we can establish a connection with every cluster.
     """
     attempts = 0
 
-    for cluster in CLUSTERS:
+    for cluster in clusters:
         clickhouse = cluster.get_query_connection(ClickhouseClientSettings.MIGRATE)
 
         while True: