lint fix on base function create_online_store

Signed-off-by: Chester Ong <chester.ong.ch@gmail.com>
feast-dev · bushwhackr · Feb 11, 2024 · Feb 11, 2024 · Feb 11, 2024 · Feb 11, 2024
commit 39fb45b4c47c63df6a69e07b46c8a130fc80f35c
@@ -32,7 +32,7 @@
 from feast.infra.provider import RetrievalJob
 from feast.infra.registry.base_registry import BaseRegistry
 from feast.on_demand_feature_view import OnDemandFeatureView
-from feast.repo_config import FeastBaseModel, RepoConfig
+from feast.repo_config import FeastConfigBaseModel, RepoConfig
 from feast.saved_dataset import SavedDatasetStorage
 from feast.type_map import pa_to_mssql_type
 from feast.usage import log_exceptions_and_usage
@@ -43,7 +43,7 @@
 EntitySchema = Dict[str, np.dtype]
 
 
-class MsSqlServerOfflineStoreConfig(FeastBaseModel):
+class MsSqlServerOfflineStoreConfig(FeastConfigBaseModel):
     """Offline store config for SQL Server"""
 
     type: Literal["mssql"] = "mssql"

@@ -14,6 +14,7 @@
 from feast.infra.offline_stores.contrib.mssql_offline_store.mssqlserver_source import (
     MsSqlServerSource,
 )
+from feast.repo_config import FeastConfigBaseModel
 from feast.saved_dataset import SavedDatasetStorage
 from tests.integration.feature_repos.universal.data_source_creator import (
     DataSourceCreator,

@@ -20,6 +20,7 @@
 from testcontainers.core.container import DockerContainer
 from testcontainers.core.waiting_utils import wait_for_logs
 
+from feast.repo_config import FeastConfigBaseModel
 from tests.integration.feature_repos.universal.online_store_creator import (
     OnlineStoreCreator,
 )
@@ -32,7 +33,7 @@ def __init__(self, project_name: str, **kwargs):
             "9042"
         )
 
-    def create_online_store(self) -> Dict[str, object]:
+    def create_online_store(self) -> FeastConfigBaseModel:
         self.container.start()
         log_string_to_wait_for = "Startup complete"
         # on a modern machine it takes about 45-60 seconds for the container
@@ -45,12 +46,12 @@ def create_online_store(self) -> Dict[str, object]:
         self.container.exec(f'cqlsh -e "{keyspace_creation_command}"')
         time.sleep(2)
         exposed_port = int(self.container.get_exposed_port("9042"))
-        return {
-            "type": "cassandra",
-            "hosts": ["127.0.0.1"],
-            "port": exposed_port,
-            "keyspace": keyspace_name,
-        }
+        return FeastConfigBaseModel(
+            type="cassandra",
+            hosts=["127.0.0.1"],
+            port=exposed_port,
+            keyspace=keyspace_name,
+        )
 
     def teardown(self):
         self.container.stop()
@@ -6,6 +6,7 @@
 from testcontainers.core.container import DockerContainer
 from testcontainers.core.waiting_utils import wait_for_logs
 
+from feast.repo_config import FeastConfigBaseModel
 from tests.integration.feature_repos.universal.online_store_creator import (
     OnlineStoreCreator,
 )
@@ -29,7 +30,7 @@ def __init__(self, project_name: str, **kwargs):
             .with_exposed_ports(5701)
         )
 
-    def create_online_store(self) -> Dict[str, Any]:
+    def create_online_store(self) -> FeastConfigBaseModel:
         self.container.start()
         cluster_member = (
             self.container.get_container_host_ip()
@@ -38,11 +39,11 @@ def create_online_store(self) -> Dict[str, Any]:
         )
         log_string_to_wait_for = r"Cluster name: " + self.cluster_name
         wait_for_logs(self.container, predicate=log_string_to_wait_for, timeout=10)
-        return {
-            "type": "hazelcast",
-            "cluster_name": self.cluster_name,
-            "cluster_members": [cluster_member],
-        }
+        return FeastConfigBaseModel(
+            type="hazelcast",
+            cluster_name=self.cluster_name,
+            cluster_members=[cluster_member],
+        )
 
     def teardown(self):
         self.container.stop()
@@ -3,6 +3,7 @@
 from testcontainers.core.container import DockerContainer
 from testcontainers.core.waiting_utils import wait_for_logs
 
+from feast.repo_config import FeastConfigBaseModel
 from tests.integration.feature_repos.universal.online_store_creator import (
     OnlineStoreCreator,
 )
@@ -13,7 +14,7 @@ def __init__(self, project_name: str, **kwargs):
         super().__init__(project_name)
         self.container = DockerContainer("harisekhon/hbase").with_exposed_ports("9090")
 
-    def create_online_store(self) -> Dict[str, str]:
+    def create_online_store(self) -> FeastConfigBaseModel:
         self.container.start()
         log_string_to_wait_for = (
             "Initializing Hbase Local with the following configuration:"
@@ -22,7 +23,11 @@ def create_online_store(self) -> Dict[str, str]:
             container=self.container, predicate=log_string_to_wait_for, timeout=10
         )
         exposed_port = self.container.get_exposed_port("9090")
-        return {"type": "hbase", "host": "127.0.0.1", "port": exposed_port}
+        return FeastConfigBaseModel(
+            type="hbase",
+            host="127.0.0.1",
+            port=exposed_port,
+        )
 
     def teardown(self):
         self.container.stop()
@@ -2,6 +2,7 @@
 
 from testcontainers.mysql import MySqlContainer
 
+from feast.repo_config import FeastConfigBaseModel
 from tests.integration.feature_repos.universal.online_store_creator import (
     OnlineStoreCreator,
 )
@@ -18,16 +19,16 @@ def __init__(self, project_name: str, **kwargs):
             .with_env("MYSQL_DATABASE", "test")
         )
 
-    def create_online_store(self) -> Dict[str, str]:
+    def create_online_store(self) -> FeastConfigBaseModel:
         self.container.start()
         exposed_port = self.container.get_exposed_port(3306)
-        return {
-            "type": "mysql",
-            "user": "root",
-            "password": "test",
-            "database": "test",
-            "port": exposed_port,
-        }
+        return FeastConfigBaseModel(
+            type="mysql",
+            user="root",
+            password="test",
+            database="test",
+            port=exposed_port,
+        )
 
     def teardown(self):
         self.container.stop()
@@ -3,6 +3,7 @@
 from testcontainers.core.container import DockerContainer
 from testcontainers.core.waiting_utils import wait_for_logs
 
+from feast.repo_config import FeastConfigBaseModel
 from tests.integration.feature_repos.universal.online_store_creator import (
     OnlineStoreCreator,
 )
@@ -13,14 +14,16 @@ def __init__(self, project_name: str, **kwargs):
         super().__init__(project_name)
         self.container = DockerContainer("redis").with_exposed_ports("6379")
 
-    def create_online_store(self) -> Dict[str, str]:
+    def create_online_store(self) -> FeastConfigBaseModel:
         self.container.start()
         log_string_to_wait_for = "Ready to accept connections"
         wait_for_logs(
             container=self.container, predicate=log_string_to_wait_for, timeout=10
         )
         exposed_port = self.container.get_exposed_port("6379")
-        return {"type": "redis", "connection_string": f"localhost:{exposed_port},db=0"}
+        return FeastConfigBaseModel(
+            type="redis", connection_string=f"localhost:{exposed_port},db=0"
+        )
 
     def teardown(self):
         self.container.stop()