add changes + fix linting issues

astronautas · astronautas · commit a9c9902d2485 · 2025-12-22T12:43:20.000+02:00
Signed-off-by: lukas.valatka &lt;lukas.valatka@cast.ai&gt;
diff --git a/infra/website/src/pages/index.astro b/infra/website/src/pages/index.astro
@@ -114,9 +114,6 @@ features = store.retrieve_online_documents(
           <div class="logo-item">
             <img src="/images/logos/seatgeek.svg" alt="SeatGeek" class="company-logo">
           </div>
-          <div class="logo-item">
-            <img src="/images/logos/castai.png" alt="Cast.ai" class="company-logo">
-          </div>
         </div>
       </div>
     </section>
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/feast/templates/postgres/feature_repo/feature_store.yaml b/sdk/python/feast/templates/postgres/feature_repo/feature_store.yaml
@@ -2,7 +2,7 @@ project: my_project
 provider: local
 registry:
     registry_type: sql
-    path: postgresql://postgres:mysecretpassword@127.0.0.1:55001/feast
+    path: postgresql://DB_USERNAME:DB_PASSWORD@DB_HOST:DB_PORT/DB_NAME
     cache_ttl_seconds: 60
     sqlalchemy_config_kwargs:
         echo: false
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
diff --git a/sdk/python/tests/unit/local_feast_tests/test_init.py b/sdk/python/tests/unit/local_feast_tests/test_init.py
@@ -67,3 +67,17 @@ def test_repo_init_with_underscore_in_project_name() -> None:
         )
         result = runner.run(["apply"], cwd=repo_dir)
         assert result.returncode != 0
+
+
+def test_postgres_template_registry_path_is_parameterized() -> None:
+    template_fs_yaml = (
+        Path(__file__).resolve().parents[3]
+        / "feast"
+        / "templates"
+        / "postgres"
+        / "feature_repo"
+        / "feature_store.yaml"
+    )
+    contents = template_fs_yaml.read_text(encoding="utf-8")
+    expected = "path: postgresql://DB_USERNAME:DB_PASSWORD@DB_HOST:DB_PORT/DB_NAME"
+    assert expected in contents
