feat: Add request_timeout setting for cassandra online store adapter (feast-dev#3256)

arcabucero · web-flow · commit da20757aa1d5 · 2022-09-28T07:35:34.000-07:00
Add request_timeout setting for cassandra online store adapter

Signed-off-by: arcabucero &lt;arcabucero@gmail.com&gt;

Signed-off-by: arcabucero &lt;arcabucero@gmail.com&gt;
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
@@ -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
