add changes + fix linting issues

astronautas · astronautas · commit 508bf0f66ea2 · 2025-12-22T12:44:16.000+02:00
Signed-off-by: lukas.valatka &lt;lukas.valatka@cast.ai&gt;
diff --git a/sdk/python/feast/infra/utils/clickhouse/clickhouse_config.py b/sdk/python/feast/infra/utils/clickhouse/clickhouse_config.py
@@ -11,4 +11,7 @@ class ClickhouseConfig(FeastConfigBaseModel):
     password: StrictStr
     use_temporary_tables_for_entity_df: bool = True
 
+    # Set this to higher than default, for larger scale offline store jobs
+    send_receive_timeout: int | None = None
+
     model_config = ConfigDict(frozen=True)
diff --git a/sdk/python/feast/infra/utils/clickhouse/connection_utils.py b/sdk/python/feast/infra/utils/clickhouse/connection_utils.py
@@ -11,12 +11,22 @@
 def get_client(config: ClickhouseConfig) -> Client:
     # Clickhouse client is not thread-safe, so we need to create a separate instance for each thread.
     if not hasattr(thread_local, "clickhouse_client"):
-        thread_local.clickhouse_client = clickhouse_connect.get_client(
-            host=config.host,
-            port=config.port,
-            user=config.user,
-            password=config.password,
-            database=config.database,
-        )
+        if config.send_receive_timeout is not None:
+            thread_local.clickhouse_client = clickhouse_connect.get_client(
+                host=config.host,
+                port=config.port,
+                user=config.user,
+                password=config.password,
+                database=config.database,
+                send_receive_timeout=config.send_receive_timeout,
+            )
+        else:
+            thread_local.clickhouse_client = clickhouse_connect.get_client(
+                host=config.host,
+                port=config.port,
+                user=config.user,
+                password=config.password,
+                database=config.database,
+            )
 
     return thread_local.clickhouse_client
diff --git a/sdk/python/tests/unit/infra/offline_stores/test_clickhouse.py b/sdk/python/tests/unit/infra/offline_stores/test_clickhouse.py
@@ -1,11 +1,16 @@
+import logging
 import threading
 from unittest.mock import MagicMock, patch
 
 import pytest
+from testcontainers.clickhouse import ClickHouseContainer
+from testcontainers.core.waiting_utils import wait_for_logs
 
 from feast.infra.utils.clickhouse.clickhouse_config import ClickhouseConfig
 from feast.infra.utils.clickhouse.connection_utils import get_client, thread_local
 
+logger = logging.getLogger(__name__)
+
 
 @pytest.fixture
 def clickhouse_config():
@@ -76,3 +81,52 @@ def thread_2_work():
     assert client_1a is not client_2, (
         "Different threads should get different client instances (not cached)"
     )
+
+
+@pytest.fixture(scope="module")
+def clickhouse_container():
+    """Start a ClickHouse container for integration testing."""
+    container = ClickHouseContainer(
+        username="default",
+        password="password",
+        dbname="default",
+    )
+    container.start()
+
+    log_string_to_wait_for = "Logging errors to"
+    waited = wait_for_logs(
+        container=container,
+        predicate=log_string_to_wait_for,
+        timeout=30,
+        interval=10,
+    )
+    logger.info("Waited for %s seconds until ClickHouse container was up", waited)
+
+    yield container
+    container.stop()
+
+
+def test_get_client_with_additional_params(clickhouse_container):
+    """
+    Test that get_client works with a real ClickHouse container and properly passes
+    additional settings like send_receive_timeout.
+    """
+    # Create config with custom send_receive_timeout
+    config = ClickhouseConfig(
+        host=clickhouse_container.get_container_host_ip(),
+        port=clickhouse_container.get_exposed_port(8123),
+        user="default",
+        password="password",
+        database="default",
+        send_receive_timeout=60,
+    )
+
+    # Get client and verify it works
+    client = get_client(config)
+
+    # Verify client is connected and functional by running a simple query
+    result = client.query("SELECT 1 AS test_value")
+    assert result.result_rows == [(1,)]
+
+    # Verify the send_receive_timeout was applied
+    assert client.timeout._read == 60
