feat: Add SSL/TLS support for ClickHouse connections

snuba/cli/cleanup.py

@@ -18,6 +18,24 @@
     type=int,
     help="Clickhouse native port to write to.",
 )
+@click.option(
+    "--clickhouse-secure",
+    type=bool,
+    default=False,
+    help="If true, an encrypted connection will be used",
+)
+@click.option(
+    "--clickhouse-ca-certs",
+    type=str,
+    default=None,
+    help="An optional path to certificates directory.",
+)
+@click.option(
+    "--clickhouse-verify",
+    type=bool,
+    default=False,
+    help="Verify ClickHouse SSL cert.",
+)
 @click.option(
     "--dry-run",
     type=bool,
@@ -36,6 +54,9 @@ def cleanup(
     *,
     clickhouse_host: Optional[str],
     clickhouse_port: Optional[int],
+    clickhouse_secure: bool,
+    clickhouse_ca_certs: Optional[str],
+    clickhouse_verify: Optional[bool],
     dry_run: bool,
     storage_name: str,
     log_level: Optional[str] = None,
@@ -67,6 +88,9 @@ def cleanup(
             clickhouse_user,
             clickhouse_password,
             database,
+            clickhouse_secure,
+            clickhouse_ca_certs,
+            clickhouse_verify,
         )
     elif not cluster.is_single_node():
         raise click.ClickException("Provide ClickHouse host and port for cleanup")

snuba/cli/migrations.py

@@ -318,12 +318,33 @@ def reverse_in_progress(
     required=True,
     default=os.environ.get("CLICKHOUSE_DATABASE", "default"),
 )
+@click.option(
+    "--secure",
+    type=bool,
+    default=False,
+    help="If true, an encrypted connection will be used",
+)
+@click.option(
+    "--ca-certs",
+    type=str,
+    default=None,
+    help="An optional path to certificates directory.",
+)
+@click.option(
+    "--verify",
+    type=bool,
+    default=False,
+    help="Verify ClickHouse SSL cert.",
+)
 def add_node(
     node_type: str,
     storage_set_names: Sequence[str],
     host_name: str,
     port: int,
     database: str,
+    secure: bool,
+    ca_certs: Optional[str],
+    verify: Optional[bool],
 ) -> None:
     """
     Runs all migrations on a brand new ClickHouse node. This should be performed
@@ -364,6 +385,9 @@ def add_node(
         user=user,
         password=password,
         database=database,
+        secure=secure,
+        ca_certs=ca_certs,
+        verify=verify,
     )
 
 

snuba/cli/optimize.py

@@ -25,6 +25,24 @@
     type=int,
     help="Clickhouse native port to write to.",
 )
+@click.option(
+    "--clickhouse-secure",
+    type=bool,
+    default=False,
+    help="If true, an encrypted connection will be used",
+)
+@click.option(
+    "--clickhouse-ca-certs",
+    type=str,
+    default=None,
+    help="An optional path to certificates directory.",
+)
+@click.option(
+    "--clickhouse-verify",
+    type=bool,
+    default=False,
+    help="Verify ClickHouse SSL cert.",
+)
 @click.option(
     "--storage",
     "storage_name",
@@ -44,6 +62,9 @@ def optimize(
     *,
     clickhouse_host: Optional[str],
     clickhouse_port: Optional[int],
+    clickhouse_secure: bool,
+    clickhouse_ca_certs: Optional[str],
+    clickhouse_verify: Optional[bool],
     storage_name: str,
     default_parallel_threads: int,
     log_level: Optional[str] = None,
@@ -79,6 +100,9 @@ def optimize(
             clickhouse_user,
             clickhouse_password,
             database,
+            clickhouse_secure,
+            clickhouse_ca_certs,
+            clickhouse_verify,
             send_receive_timeout=ClickhouseClientSettings.OPTIMIZE.value.timeout,
         )
     elif not storage.get_cluster().is_single_node():

snuba/clickhouse/http.py

@@ -20,7 +20,7 @@
 
 import rapidjson
 import sentry_sdk
-from urllib3.connectionpool import HTTPConnectionPool
+from urllib3.connectionpool import HTTPConnectionPool, HTTPSConnectionPool
 from urllib3.exceptions import HTTPError
 
 from snuba import settings, state
@@ -293,6 +293,9 @@ def __init__(
         port: int,
         user: str,
         password: str,
+        secure: bool,
+        ca_certs: Optional[str],
+        verify: Optional[bool],
         metrics: MetricsBackend,
         statement: InsertStatement,
         encoding: Optional[str],
@@ -302,9 +305,10 @@ def __init__(
         max_connections: int = 1,
         block_connections: bool = False,
     ):
-        self.__pool = HTTPConnectionPool(
-            host, port, maxsize=max_connections, block=block_connections
-        )
+        if secure:
+            self.__pool = HTTPSConnectionPool(host, port, ca_certs=ca_certs, verify=verify)
+        else:
+            self.__pool = HTTPConnectionPool(host, port)
         self.__executor = ThreadPoolExecutor()
         self.__metrics = metrics
 

snuba/clickhouse/native.py

@@ -81,6 +81,9 @@ def __init__(
         user: str,
         password: str,
         database: str,
+        secure: bool = False,
+        ca_certs: Optional[str] = None,
+        verify: Optional[bool] = False,
         connect_timeout: int = 1,
         send_receive_timeout: Optional[int] = 300,
         max_pool_size: int = settings.CLICKHOUSE_MAX_POOL_SIZE,
@@ -91,6 +94,9 @@ def __init__(
         self.user = user
         self.password = password
         self.database = database
+        self.secure = secure
+        self.ca_certs = ca_certs
+        self.verify = verify
         self.connect_timeout = connect_timeout
         self.send_receive_timeout = send_receive_timeout
         self.client_settings = client_settings
@@ -380,6 +386,9 @@ def _create_conn(self, use_fallback_host: bool = False) -> Client:
             user=self.user,
             password=self.password,
             database=self.database,
+            secure=self.secure,
+            ca_certs=self.ca_certs,
+            verify=self.verify,
             connect_timeout=self.connect_timeout,
             send_receive_timeout=self.send_receive_timeout,
             settings=self.client_settings,

snuba/clusters/cluster.py

@@ -226,6 +226,9 @@ def __init__(
         password: str,
         database: str,
         http_port: int,
+        secure: bool,
+        ca_certs: Optional[str],
+        verify: Optional[bool],
         storage_sets: Set[str],
         single_node: bool,
         # The cluster name and distributed cluster name only apply if single_node is set to False
@@ -246,6 +249,9 @@ def __init__(
         self.__password = password
         self.__database = database
         self.__http_port = http_port
+        self.__secure = secure
+        self.__ca_certs = ca_certs
+        self.__verify = verify
         self.__single_node = single_node
         self.__cluster_name = cluster_name
         self.__distributed_cluster_name = distributed_cluster_name
@@ -290,6 +296,9 @@ def get_node_connection(
             self.__user,
             self.__password,
             self.__database,
+            self.__secure,
+            self.__ca_certs,
+            self.__verify,
         )
 
     def get_deleter(self) -> Reader:
@@ -331,6 +340,9 @@ def get_batch_writer(
             block_connections=self.__block_connections,
             user=self.__user,
             password=self.__password,
+            secure=self.__secure,
+            ca_certs=self.__ca_certs,
+            verify=self.__verify,
             metrics=metrics,
             statement=insert_statement.with_database(self.__database),
             encoding=encoding,
@@ -413,6 +425,9 @@ def get_http_port(self) -> int:
         password=cluster.get("password", ""),
         database=cluster.get("database", "default"),
         http_port=cluster["http_port"],
+        secure=cluster["secure"],
+        ca_certs=cluster["ca_certs"],
+        verify=cluster["verify"],
         storage_sets=cluster["storage_sets"],
         single_node=cluster["single_node"],
         cluster_name=cluster["cluster_name"] if "cluster_name" in cluster else None,
@@ -459,6 +474,9 @@ def _build_sliced_cluster(cluster: Mapping[str, Any]) -> ClickhouseCluster:
         password=cluster.get("password", ""),
         database=cluster.get("database", "default"),
         http_port=cluster["http_port"],
+        secure=cluster["secure"],
+        ca_certs=cluster["ca_certs"],
+        verify=cluster["verify"],
         storage_sets={
             storage_tuple[0] for storage_tuple in cluster["storage_set_slices"]
         },

snuba/migrations/runner.py

@@ -572,6 +572,9 @@ def add_node(
         user: str,
         password: str,
         database: str,
+        secure: bool = False,
+        ca_certs: Optional[str] = None,
+        verify: Optional[bool] = False,
     ) -> None:
         client_settings = ClickhouseClientSettings.MIGRATE.value
         clickhouse = ClickhousePool(
@@ -580,6 +583,9 @@ def add_node(
             user,
             password,
             database,
+            secure,
+            ca_certs,
+            verify,
             client_settings=client_settings.settings,
             send_receive_timeout=client_settings.timeout,
         )

snuba/settings/__init__.py

@@ -93,6 +93,9 @@
         "password": os.environ.get("CLICKHOUSE_PASSWORD", ""),
         "database": os.environ.get("CLICKHOUSE_DATABASE", "default"),
         "http_port": int(os.environ.get("CLICKHOUSE_HTTP_PORT", 8123)),
+        "secure": os.environ.get("CLICKHOUSE_SECURE", "0") == "1",
+        "ca_certs": os.environ.get("CLICKHOUSE_CA_CERTS"),
+        "verify": os.environ.get("CLICKHOUSE_VERIFY"),
         "storage_sets": {
             "cdc",
             "discover",

snuba/settings/settings_distributed.py

@@ -12,6 +12,9 @@
         "password": os.environ.get("CLICKHOUSE_PASSWORD", ""),
         "database": os.environ.get("CLICKHOUSE_DATABASE", "default"),
         "http_port": int(os.environ.get("CLICKHOUSE_HTTP_PORT", 8123)),
+        "secure": os.environ.get("CLICKHOUSE_SECURE", "0") == "1",
+        "ca_certs": os.environ.get("CLICKHOUSE_CA_CERTS"),
+        "verify": os.environ.get("CLICKHOUSE_VERIFY"),
         "storage_sets": {
             "cdc",
             "discover",

tests/clusters/fake_cluster.py

@@ -69,6 +69,9 @@ def __init__(
         password: str,
         database: str,
         http_port: int,
+        secure: bool,
+        ca_certs: Optional[str],
+        verify: Optional[bool],
         storage_sets: Set[str],
         single_node: bool,
         # The cluster name and distributed cluster name only apply if single_node is set to False

tests/clusters/test_cluster.py

@@ -60,6 +60,9 @@
         "password": "",
         "database": "default",
         "http_port": 8123,
+        "secure": False,
+        "ca_certs": None,
+        "verify": False,
         "storage_sets": ALL_STORAGE_SETS,
         "single_node": True,
     },
@@ -70,6 +73,9 @@
         "password": "",
         "database": "default",
         "http_port": 8123,
+        "secure": False,
+        "ca_certs": None,
+        "verify": False,
         "storage_sets": {"transactions"},
         "single_node": False,
         "cluster_name": "clickhouse_hosts",
@@ -195,16 +201,16 @@ def test_get_local_nodes() -> None:
 @pytest.mark.clickhouse_db
 def test_cache_connections() -> None:
     cluster_1 = cluster.ClickhouseCluster(
-        "127.0.0.1", 8000, "default", "", "default", 8001, {"events"}, True
+        "127.0.0.1", 8000, "default", "", "default", 8001, False, None, False, {"events"}, True
     )
 
     cluster_2 = cluster.ClickhouseCluster(
-        "127.0.0.1", 8000, "default", "", "default", 8001, {"transactions"}, True
+        "127.0.0.1", 8000, "default", "", "default", 8001, False, None, False, {"transactions"}, True
     )
 
     # Same node but different user
     cluster_3 = cluster.ClickhouseCluster(
-        "127.0.0.1", 8000, "readonly", "", "default", 8001, {"metrics"}, True
+        "127.0.0.1", 8000, "readonly", "", "default", 8001, False, None, False, {"metrics"}, True
     )
 
     assert cluster_1.get_query_connection(

tests/conftest.py

@@ -45,6 +45,9 @@ def create_databases() -> None:
             password="",
             database="default",
             http_port=cluster["http_port"],
+            secure=cluster["secure"],
+            ca_certs=cluster["ca_certs"],
+            verify=cluster["verify"],
             storage_sets=cluster["storage_sets"],
             single_node=cluster["single_node"],
             cluster_name=cluster["cluster_name"] if "cluster_name" in cluster else None,

tests/migrations/test_connect.py

@@ -24,6 +24,9 @@
     password="",
     database="default",
     http_port=420,
+    secure=False,
+    ca_certs=None,
+    verify=False,
     storage_sets={
         "querylog",
     },
@@ -37,6 +40,9 @@
     password="",
     database="default",
     http_port=420,
+    secure=False,
+    ca_certs=None,
+    verify=False,
     storage_sets={
         "events",
     },
@@ -50,6 +56,9 @@
     password="",
     database="default",
     http_port=420,
+    secure=False,
+    ca_certs=None,
+    verify=False,
     storage_sets=_REMAINING_STORAGE_SET_KEYS,
     single_node=True,
 )