fix: Adopt connection pooling for HBase

sudohainguyen · sudohainguyen · commit 5390d7b7ea28 · 2023-10-12T07:44:00.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
@@ -3,7 +3,7 @@
 from datetime import datetime
 from typing import Any, Callable, Dict, List, Optional, Sequence, Tuple
 
-from happybase import Connection
+from happybase import Connection, ConnectionPool
 from pydantic.typing import Literal
 
 from feast import Entity
@@ -29,6 +29,9 @@ class HbaseOnlineStoreConfig(FeastConfigBaseModel):
     port: str
     """Port in which Hbase Thrift server is running"""
 
+    connection_pool_size: int = 4
+    """Number of connections to Hbase Thrift server to keep in the connection pool"""
+
 
 class HbaseConnection:
     """
@@ -62,7 +65,7 @@ class HbaseOnlineStore(OnlineStore):
         _conn: Happybase Connection to connect to hbase thrift server.
     """
 
-    _conn: Connection = None
+    _conn: ConnectionPool = None
 
     def _get_conn(self, config: RepoConfig):
         """
@@ -76,7 +79,11 @@ def _get_conn(self, config: RepoConfig):
         assert isinstance(store_config, HbaseOnlineStoreConfig)
 
         if not self._conn:
-            self._conn = Connection(host=store_config.host, port=int(store_config.port))
+            self._conn = ConnectionPool(
+                host=store_config.host,
+                port=int(store_config.port),
+                size=int(store_config.connection_pool_size),
+            )
         return self._conn
 
     @log_exceptions_and_usage(online_store="hbase")
diff --git a/sdk/python/feast/infra/utils/hbase_utils.py b/sdk/python/feast/infra/utils/hbase_utils.py
@@ -1,9 +1,6 @@
 from typing import List
 
-from happybase import Connection
-
-from feast.infra.key_encoding_utils import serialize_entity_key
-from feast.protos.feast.types.EntityKey_pb2 import EntityKey
+from happybase import ConnectionPool
 
 
 class HbaseConstants:
@@ -40,14 +37,22 @@ class HbaseUtils:
     """
 
     def __init__(
-        self, conn: Connection = None, host: str = None, port: int = None, timeout=None
+        self,
+        pool: ConnectionPool = None,
+        host: str = None,
+        port: int = None,
+        connection_pool_size: int = 4,
     ):
-        if conn is None:
+        if pool is None:
             self.host = host
             self.port = port
-            self.conn = Connection(host=host, port=port, timeout=timeout)
+            self.pool = ConnectionPool(
+                host=host,
+                port=port,
+                size=connection_pool_size,
+            )
         else:
-            self.conn = conn
+            self.pool = pool
 
     def create_table(self, table_name: str, colm_family: List[str]):
         """
@@ -60,7 +65,9 @@ def create_table(self, table_name: str, colm_family: List[str]):
         cf_dict: dict = {}
         for cf in colm_family:
             cf_dict[cf] = dict()
-        return self.conn.create_table(table_name, cf_dict)
+
+        with self.pool.connection() as conn:
+            return conn.create_table(table_name, cf_dict)
 
     def create_table_with_default_cf(self, table_name: str):
         """
@@ -69,7 +76,8 @@ def create_table_with_default_cf(self, table_name: str):
         Arguments:
             table_name: Name of the Hbase table.
         """
-        return self.conn.create_table(table_name, {"default": dict()})
+        with self.pool.connection() as conn:
+            return conn.create_table(table_name, {"default": dict()})
 
     def check_if_table_exist(self, table_name: str):
         """
@@ -78,16 +86,18 @@ def check_if_table_exist(self, table_name: str):
         Arguments:
             table_name: Name of the Hbase table.
         """
-        return bytes(table_name, "utf-8") in self.conn.tables()
+        with self.pool.connection() as conn:
+            return bytes(table_name, "utf-8") in conn.tables()
 
     def batch(self, table_name: str):
         """
-        Returns a 'Batch' instance that can be used for mass data manipulation in the hbase table.
+        Returns a "Batch" instance that can be used for mass data manipulation in the hbase table.
 
         Arguments:
             table_name: Name of the Hbase table.
         """
-        return self.conn.table(table_name).batch()
+        with self.pool.connection() as conn:
+            return conn.table(table_name).batch()
 
     def put(self, table_name: str, row_key: str, data: dict):
         """
@@ -98,8 +108,9 @@ def put(self, table_name: str, row_key: str, data: dict):
             row_key: Row key of the row to be inserted to hbase table.
             data: Mapping of column family name:column name to column values
         """
-        table = self.conn.table(table_name)
-        table.put(row_key, data)
+        with self.pool.connection() as conn:
+            table = conn.table(table_name)
+            table.put(row_key, data)
 
     def row(
         self,
@@ -119,8 +130,9 @@ def row(
             timestamp: timestamp specifies the maximum version the cells can have.
             include_timestamp: specifies if (column, timestamp) to be return instead of only column.
         """
-        table = self.conn.table(table_name)
-        return table.row(row_key, columns, timestamp, include_timestamp)
+        with self.pool.connection() as conn:
+            table = conn.table(table_name)
+            return table.row(row_key, columns, timestamp, include_timestamp)
 
     def rows(
         self,
@@ -140,52 +152,69 @@ def rows(
             timestamp: timestamp specifies the maximum version the cells can have.
             include_timestamp: specifies if (column, timestamp) to be return instead of only column.
         """
-        table = self.conn.table(table_name)
-        return table.rows(row_keys, columns, timestamp, include_timestamp)
+        with self.pool.connection() as conn:
+            table = conn.table(table_name)
+            return table.rows(row_keys, columns, timestamp, include_timestamp)
 
     def print_table(self, table_name):
         """Prints the table scanning all the rows of the hbase table."""
-        table = self.conn.table(table_name)
-        scan_data = table.scan()
-        for row_key, cols in scan_data:
-            print(row_key.decode("utf-8"), cols)
+        with self.pool.connection() as conn:
+            table = conn.table(table_name)
+            scan_data = table.scan()
+            for row_key, cols in scan_data:
+                print(row_key.decode("utf-8"), cols)
 
     def delete_table(self, table: str):
         """Deletes the hbase table given the table name."""
         if self.check_if_table_exist(table):
-            self.conn.delete_table(table, disable=True)
+            with self.pool.connection() as conn:
+                conn.delete_table(table, disable=True)
 
     def close_conn(self):
         """Closes the happybase connection."""
-        self.conn.close()
+        with self.pool.connection() as conn:
+            conn.close()
 
 
 def main():
+    from feast.infra.key_encoding_utils import serialize_entity_key
+    from feast.protos.feast.types.EntityKey_pb2 import EntityKey
     from feast.protos.feast.types.Value_pb2 import Value
 
-    connection = Connection(host="localhost", port=9090)
-    table = connection.table("test_hbase_driver_hourly_stats")
-    row_keys = [
-        serialize_entity_key(
-            EntityKey(join_keys=["driver_id"], entity_values=[Value(int64_val=1004)]),
-            entity_key_serialization_version=2,
-        ).hex(),
-        serialize_entity_key(
-            EntityKey(join_keys=["driver_id"], entity_values=[Value(int64_val=1005)]),
-            entity_key_serialization_version=2,
-        ).hex(),
-        serialize_entity_key(
-            EntityKey(join_keys=["driver_id"], entity_values=[Value(int64_val=1024)]),
-            entity_key_serialization_version=2,
-        ).hex(),
-    ]
-    rows = table.rows(row_keys)
-
-    for row_key, row in rows:
-        for key, value in row.items():
-            col_name = bytes.decode(key, "utf-8").split(":")[1]
-            print(col_name, value)
-        print()
+    pool = ConnectionPool(
+        host="localhost",
+        port=9090,
+        size=2,
+    )
+    with pool.connection() as connection:
+        table = connection.table("test_hbase_driver_hourly_stats")
+        row_keys = [
+            serialize_entity_key(
+                EntityKey(
+                    join_keys=["driver_id"], entity_values=[Value(int64_val=1004)]
+                ),
+                entity_key_serialization_version=2,
+            ).hex(),
+            serialize_entity_key(
+                EntityKey(
+                    join_keys=["driver_id"], entity_values=[Value(int64_val=1005)]
+                ),
+                entity_key_serialization_version=2,
+            ).hex(),
+            serialize_entity_key(
+                EntityKey(
+                    join_keys=["driver_id"], entity_values=[Value(int64_val=1024)]
+                ),
+                entity_key_serialization_version=2,
+            ).hex(),
+        ]
+        rows = table.rows(row_keys)
+
+        for _, row in rows:
+            for key, value in row.items():
+                col_name = bytes.decode(key, "utf-8").split(":")[1]
+                print(col_name, value)
+            print()
 
 
 if __name__ == "__main__":
