init and close methods in interface

sdk/python/feast/feature_store.py

@@ -2078,6 +2078,14 @@ def list_saved_datasets(
             self.project, allow_cache=allow_cache, tags=tags
         )
 
+    async def initialize(self) -> None:
+        """Initialize long-lived clients and/or resources needed for accessing datastores"""
+        await self._get_provider().initialize(self.config)
+
+    async def close(self) -> None:
+        """Cleanup any long-lived clients and/or resources"""
+        await self._get_provider().close()
+
 
 def _print_materialization_log(
     start_date, end_date, num_feature_views: int, online_store: str

sdk/python/feast/infra/offline_stores/offline_store.py

@@ -388,3 +388,9 @@ def get_table_column_names_and_types_from_data_source(
             data_source: DataSource object
         """
         return data_source.get_table_column_names_and_types(config=config)
+
+    async def initialize(self, config: RepoConfig) -> None:
+        pass
+
+    async def close(self) -> None:
+        pass

sdk/python/feast/infra/online_stores/dynamodb.py

@@ -91,25 +91,14 @@ class DynamoDBOnlineStore(OnlineStore):
 
     _dynamodb_client = None
     _dynamodb_resource = None
-    _aioboto_client = None
 
-    @staticmethod
-    async def initialize(config: RepoConfig):
-        if not DynamoDBOnlineStore._aioboto_client:
-            context_stack = contextlib.AsyncExitStack()
-            DynamoDBOnlineStore._aioboto_client = (
-                await context_stack.enter_async_context(
-                    _get_aiodynamodb_client(
-                        region=config.online_store.region,
-                        max_pool_connections=config.online_store.max_pool_connections,
-                    )
-                )
-            )
+    async def initialize(self, config: RepoConfig):
+        await _get_aiodynamodb_client(
+            config.online_store.region, config.online_store.max_pool_connections
+        )
 
-    @staticmethod
-    async def close():
-        if DynamoDBOnlineStore._aioboto_client:
-            DynamoDBOnlineStore._aioboto_client.close()
+    async def close(self):
+        await _aiodynamodb_close()
 
     def update(
         self,
@@ -325,8 +314,6 @@ async def online_read_async(
 
         deserialize = TypeDeserializer().deserialize
 
-        await self.initialize(config)
-
         def to_tbl_resp(raw_client_response):
             return {
                 "entity_id": deserialize(raw_client_response["entity_id"]),
@@ -346,18 +333,17 @@ def to_tbl_resp(raw_client_response):
             batches.append(batch)
             entity_id_batches.append(entity_id_batch)
 
-        client_context = _get_aiodynamodb_client(
+        client = _get_aiodynamodb_client(
             online_config.region, online_config.max_pool_connections
         )
-        async with client_context as client:
-            response_batches = await asyncio.gather(
-                *[
-                    client.batch_get_item(
-                        RequestItems=entity_id_batch,
-                    )
-                    for entity_id_batch in entity_id_batches
-                ]
-            )
+        response_batches = await asyncio.gather(
+            *[
+                client.batch_get_item(
+                    RequestItems=entity_id_batch,
+                )
+                for entity_id_batch in entity_id_batches
+            ]
+        )
 
         result_batches = []
         for batch, response in zip(batches, response_batches):
@@ -505,6 +491,7 @@ def _to_client_batch_get_payload(online_config, table_name, batch):
 
 
 _aioboto_session = None
+_aioboto_client = None
 
 
 def _get_aioboto_session():
@@ -514,12 +501,23 @@ def _get_aioboto_session():
     return _aioboto_session
 
 
-def _get_aiodynamodb_client(region: str, max_pool_connections: int):
-    return _get_aioboto_session().create_client(
-        "dynamodb",
-        region_name=region,
-        config=AioConfig(max_pool_connections=max_pool_connections),
-    )
+async def _get_aiodynamodb_client(region: str, max_pool_connections: int):
+    global _aioboto_client
+    if _aioboto_client is None:
+        client_context = _get_aioboto_session().create_client(
+            "dynamodb",
+            region_name=region,
+            config=AioConfig(max_pool_connections=max_pool_connections),
+        )
+        context_stack = contextlib.AsyncExitStack()
+        _aioboto_client = await context_stack.enter_async_context(client_context)
+    return _aioboto_client
+
+
+async def _aiodynamodb_close():
+    global _aioboto_client
+    if _aioboto_client:
+        await _aioboto_client.close()
 
 
 def _initialize_dynamodb_client(

sdk/python/feast/infra/online_stores/online_store.py

@@ -390,3 +390,9 @@ def retrieve_online_documents(
         raise NotImplementedError(
             f"Online store {self.__class__.__name__} does not support online retrieval"
         )
+
+    async def initialize(self, config: RepoConfig) -> None:
+        pass
+
+    async def close(self) -> None:
+        pass

sdk/python/feast/infra/passthrough_provider.py

@@ -473,3 +473,11 @@ def get_table_column_names_and_types_from_data_source(
         return self.offline_store.get_table_column_names_and_types_from_data_source(
             config=config, data_source=data_source
         )
+
+    async def initialize(self, config: RepoConfig) -> None:
+        await self.online_store.initialize(config)
+        await self.offline_store.initialize(config)
+
+    async def close(self) -> None:
+        await self.online_store.close()
+        await self.offline_store.close()

sdk/python/feast/infra/provider.py

@@ -429,6 +429,14 @@ def get_table_column_names_and_types_from_data_source(
         """
         pass
 
+    @abstractmethod
+    async def initialize(self, config: RepoConfig) -> None:
+        pass
+
+    @abstractmethod
+    async def close(self) -> None:
+        pass
+
 
 def get_provider(config: RepoConfig) -> Provider:
     if "." not in config.provider: