Skip to content
Merged
Show file tree
Hide file tree
Changes from 1 commit
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Prev Previous commit
Next Next commit
generalize based on pr suggestion
Signed-off-by: lukas.valatka <lukas.valatka@cast.ai>
  • Loading branch information
astronautas authored and ntkathole committed Dec 23, 2025
commit 24f6f40fc8f65f10360603eef7e015c85196f4fa
Original file line number Diff line number Diff line change
@@ -1,3 +1,5 @@
from typing import Any

from pydantic import ConfigDict, StrictStr

from feast.repo_config import FeastConfigBaseModel
Expand All @@ -11,7 +13,8 @@ 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
# See https://github.com/ClickHouse/clickhouse-connect/blob/main/clickhouse_connect/driver/__init__.py#L51
# Some typical ones e.g. send_receive_timeout (read_timeout), etc
additional_client_args: dict[str, Any] | None = None

model_config = ConfigDict(frozen=True)
6 changes: 4 additions & 2 deletions sdk/python/feast/infra/utils/clickhouse/connection_utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -11,14 +11,16 @@
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"):
if config.send_receive_timeout is not None:
additional_client_args = config.additional_client_args

if additional_client_args:
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,
**additional_client_args,
)
else:
thread_local.clickhouse_client = clickhouse_connect.get_client(
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -118,7 +118,7 @@ def test_get_client_with_additional_params(clickhouse_container):
user="default",
password="password",
database="default",
send_receive_timeout=60,
additional_client_args={"send_receive_timeout": 60},
)

# Get client and verify it works
Expand Down