dymanic load of RegistryStore class

Signed-off-by: DvirDukhan <dvir@redislabs.com>

sdk/python/feast/infra/aws.py

@@ -1,6 +1,7 @@
 import os
 import uuid
 from datetime import datetime
+from pathlib import Path
 from tempfile import TemporaryFile
 from typing import Any, Callable, Dict, List, Optional, Sequence, Tuple, Union
 from urllib.parse import urlparse
@@ -26,7 +27,7 @@
 from feast.protos.feast.types.Value_pb2 import Value as ValueProto
 from feast.registry import Registry
 from feast.registry_store import RegistryStore
-from feast.repo_config import RepoConfig
+from feast.repo_config import RegistryConfig, RepoConfig
 
 
 class AwsProvider(Provider):
@@ -151,7 +152,8 @@ def get_historical_features(
 
 
 class S3RegistryStore(RegistryStore):
-    def __init__(self, uri: str):
+    def __init__(self, registry_config: RegistryConfig, repo_path: Path):
+        uri = registry_config.path
         try:
             import boto3
         except ImportError as e:

sdk/python/feast/infra/gcp.py

@@ -1,5 +1,6 @@
 import uuid
 from datetime import datetime
+from pathlib import Path
 from tempfile import TemporaryFile
 from typing import Any, Callable, Dict, List, Optional, Sequence, Tuple, Union
 from urllib.parse import urlparse
@@ -24,7 +25,7 @@
 from feast.protos.feast.types.Value_pb2 import Value as ValueProto
 from feast.registry import Registry
 from feast.registry_store import RegistryStore
-from feast.repo_config import RepoConfig
+from feast.repo_config import RegistryConfig, RepoConfig
 
 
 class GcpProvider(Provider):
@@ -151,7 +152,8 @@ def get_historical_features(
 
 
 class GCSRegistryStore(RegistryStore):
-    def __init__(self, uri: str):
+    def __init__(self, registry_config: RegistryConfig, repo_path: Path):
+        uri = registry_config.path
         try:
             from google.cloud import storage
         except ImportError as e:

sdk/python/feast/infra/local.py

@@ -24,7 +24,7 @@
 from feast.protos.feast.types.Value_pb2 import Value as ValueProto
 from feast.registry import Registry
 from feast.registry_store import RegistryStore
-from feast.repo_config import RepoConfig
+from feast.repo_config import RegistryConfig, RepoConfig
 
 
 class LocalProvider(Provider):
@@ -159,8 +159,8 @@ def _to_naive_utc(ts: datetime):
 
 
 class LocalRegistryStore(RegistryStore):
-    def __init__(self, repo_path: Path, registry_path_string: str):
-        registry_path = Path(registry_path_string)
+    def __init__(self, registry_config: RegistryConfig, repo_path: Path):
+        registry_path = Path(registry_config.path)
         if registry_path.is_absolute():
             self._filepath = registry_path
         else:

sdk/python/feast/registry.py

@@ -32,12 +32,49 @@
 from feast.feature_view import FeatureView
 from feast.on_demand_feature_view import OnDemandFeatureView
 from feast.protos.feast.core.Registry_pb2 import Registry as RegistryProto
-from feast.registry_store import RegistryStore
 from feast.repo_config import RegistryConfig
 
 REGISTRY_SCHEMA_VERSION = "1"
 
 
+REGISTRY_STORE_CLASS_FOR_TYPE = {
+    "GCSRegistryStore": "feast.infra.gcp.GCSRegistryStore",
+    "S3RegistryStore": "feast.infra.aws.S3RegistryStore",
+    "LocalRegistryStore": "feast.infra.local.LocalRegistryStore",
+}
+
+REGISTRY_STORE_CLASS_FOR_SCHEME = {
+    "gs": "GCSRegistryStore",
+    "s3": "S3RegistryStore",
+    "file": "LocalRegistryStore",
+    "": "LocalRegistryStore",
+}
+
+
+def get_registry_store_class_from_type(registry_store_type: str):
+    if not registry_store_type.endswith("RegistryStore"):
+        raise Exception('Registry store class name should end with "RegistryStore"')
+    if registry_store_type in REGISTRY_STORE_CLASS_FOR_TYPE:
+        registry_store_type = REGISTRY_STORE_CLASS_FOR_TYPE[registry_store_type]
+    module_name, registry_store_class_name = registry_store_type.rsplit(".", 1)
+
+    return importer.get_class_from_type(
+        module_name, registry_store_class_name, "RegistryStore"
+    )
+
+
+def get_registry_store_class_from_scheme(registry_path: str):
+    uri = urlparse(registry_path)
+    if uri.scheme not in REGISTRY_STORE_CLASS_FOR_SCHEME:
+        raise Exception(
+            f"Registry path {registry_path} has unsupported scheme {uri.scheme}. "
+            f"Supported schemes are file, s3 and gs."
+        )
+    else:
+        registry_store_type = REGISTRY_STORE_CLASS_FOR_SCHEME[uri.scheme]
+        return get_registry_store_class_from_type(registry_store_type)
+
+
 class Registry:
     """
     Registry: A registry allows for the management and persistence of feature definitions and related metadata.
@@ -62,61 +99,12 @@ def __init__(self, registry_config: RegistryConfig, repo_path: Path):
         """
         registry_store_type = registry_config.registry_store_type
         registry_path = registry_config.path
-        uri = urlparse(registry_path)
         if registry_store_type is None:
-            if uri.scheme == "gs":
-                from feast.infra.gcp import GCSRegistryStore
-
-                self._registry_store: RegistryStore = GCSRegistryStore(registry_path)
-            elif uri.scheme == "s3":
-                from feast.infra.aws import S3RegistryStore
-
-                self._registry_store = S3RegistryStore(registry_path)
-            elif uri.scheme == "file" or uri.scheme == "":
-                from feast.infra.local import LocalRegistryStore
-
-                self._registry_store = LocalRegistryStore(
-                    repo_path=repo_path, registry_path_string=registry_path
-                )
-            else:
-                raise Exception(
-                    f"Registry path {registry_path} has unsupported scheme {uri.scheme}. "
-                    f"Supported schemes are file, s3 and gs."
-                )
+            cls = get_registry_store_class_from_scheme(registry_path)
         else:
-            registry_store_type = str(registry_store_type)
-            if "." not in registry_store_type:
-                if uri.scheme == "gs":
-                    from feast.infra.gcp import GCSRegistryStore
-
-                    self._registry_store = GCSRegistryStore(registry_path)
-                elif uri.scheme == "s3":
-                    from feast.infra.aws import S3RegistryStore
-
-                    self._registry_store = S3RegistryStore(registry_path)
-                elif uri.scheme == "file" or uri.scheme == "":
-                    from feast.infra.local import LocalRegistryStore
-
-                    self._registry_store = LocalRegistryStore(
-                        repo_path=repo_path, registry_path_string=registry_path
-                    )
-                else:
-                    raise Exception(
-                        f"Registry path {registry_path} has unsupported scheme {uri.scheme} or does not match {registry_store_type}. "
-                        f"Supported schemes are file, s3 and gs."
-                    )
-            else:
-                # Split provider into module and class names by finding the right-most dot.
-                # For example, provider 'foo.bar.MyRegistryStore' will be parsed into 'foo.bar' and 'MyRegistryStore'
-                module_name, class_name = str(
-                    registry_config.registry_store_type
-                ).rsplit(".", 1)
-
-                cls = importer.get_class_from_type(
-                    module_name, class_name, "RegistryStore"
-                )
+            cls = get_registry_store_class_from_type(str(registry_store_type))
 
-                self._registry_store = cls(registry_config, repo_path)
+        self._registry_store = cls(registry_config, repo_path)
         self.cached_registry_proto_ttl = timedelta(
             seconds=registry_config.cache_ttl_seconds
             if registry_config.cache_ttl_seconds is not None

sdk/python/feast/registry_store.py

@@ -9,7 +9,7 @@ class RegistryStore(ABC):
     """
 
     @abstractmethod
-    def get_registry_proto(self )->RegistryProto:
+    def get_registry_proto(self) -> RegistryProto:
         """
         Retrieves the registry proto from the registry path. If there is no file at that path,
         raises a FileNotFoundError.

sdk/python/tests/integration/registration/test_cli.py

@@ -215,7 +215,7 @@ def test_3rd_party_registry_store() -> None:
         return_code, output = runner.run_with_output(["apply"], cwd=repo_path)
         assertpy.assert_that(return_code).is_equal_to(1)
         assertpy.assert_that(output).contains(
-            b"Registry path foobar://foo.bar has unsupported scheme foobar or does not match feast123. Supported schemes are file, s3 and gs"
+            b'Registry store class name should end with "RegistryStore"'
         )
     # Check with incorrect third-party registry store name (with dots)
     with setup_third_party_registry_store_repo("feast_foo.RegistryStore") as repo_path: