feat: Expose thrift connection configuration to hbase online store config

sudohainguyen · sudohainguyen · commit 9a6ecd392242 · 2023-10-21T23:41:09.000+07:00
Signed-off-by: Hai Nguyen &lt;quanghai.ng1512@gmail.com&gt;
diff --git a/sdk/python/feast/infra/online_stores/contrib/hbase_online_store/hbase.py b/sdk/python/feast/infra/online_stores/contrib/hbase_online_store/hbase.py
@@ -4,6 +4,7 @@
 from typing import Any, Callable, Dict, List, Optional, Sequence, Tuple
 
 from happybase import ConnectionPool
+from happybase.connection import DEFAULT_TRANSPORT, DEFAULT_PROTOCOL
 from pydantic.typing import Literal
 
 from feast import Entity
@@ -32,6 +33,12 @@ class HbaseOnlineStoreConfig(FeastConfigBaseModel):
     connection_pool_size: int = 4
     """Number of connections to Hbase Thrift server to keep in the connection pool"""
 
+    protocol: str = DEFAULT_PROTOCOL
+    """Protocol used to communicate with Hbase Thrift server"""
+
+    transport: str = DEFAULT_TRANSPORT
+    """Transport used to communicate with Hbase Thrift server"""
+
 
 class HbaseOnlineStore(OnlineStore):
     """
@@ -59,6 +66,8 @@ def _get_conn(self, config: RepoConfig):
                 host=store_config.host,
                 port=int(store_config.port),
                 size=int(store_config.connection_pool_size),
+                protocol=store_config.protocol,
+                transport=store_config.transport,
             )
         return self._conn
 
