From de565299d21f2cf04edf1aaff23f0923fd65dfca Mon Sep 17 00:00:00 2001 From: arcabucero Date: Wed, 28 Sep 2022 00:49:26 +0300 Subject: [PATCH] Add request_timeout setting for cassandra online store adapter Signed-off-by: arcabucero --- .../cassandra_online_store/cassandra_online_store.py | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/sdk/python/feast/infra/online_stores/contrib/cassandra_online_store/cassandra_online_store.py b/sdk/python/feast/infra/online_stores/contrib/cassandra_online_store/cassandra_online_store.py index f89517c41eb..71d3167c46f 100644 --- a/sdk/python/feast/infra/online_stores/contrib/cassandra_online_store/cassandra_online_store.py +++ b/sdk/python/feast/infra/online_stores/contrib/cassandra_online_store/cassandra_online_store.py @@ -32,7 +32,7 @@ ) from cassandra.policies import DCAwareRoundRobinPolicy, TokenAwarePolicy from cassandra.query import PreparedStatement -from pydantic import StrictInt, StrictStr +from pydantic import StrictFloat, StrictInt, StrictStr from pydantic.typing import Literal from feast import Entity, FeatureView, RepoConfig @@ -141,6 +141,9 @@ class CassandraOnlineStoreConfig(FeastConfigBaseModel): protocol_version: Optional[StrictInt] = None """Explicit specification of the CQL protocol version used.""" + request_timeout: Optional[StrictFloat] = None + """Request timeout in seconds.""" + class CassandraLoadBalancingPolicy(FeastConfigBaseModel): """ Configuration block related to the Cluster's load-balancing policy. @@ -240,7 +243,10 @@ def _get_session(self, config: RepoConfig): raise CassandraInvalidConfig(E_CASSANDRA_UNKNOWN_LB_POLICY) # wrap it up in a map of ex.profiles with a default - exe_profile = ExecutionProfile(load_balancing_policy=lb_policy) + exe_profile = ExecutionProfile( + request_timeout=online_store_config.request_timeout, + load_balancing_policy=lb_policy, + ) execution_profiles = {EXEC_PROFILE_DEFAULT: exe_profile} else: execution_profiles = None