diff --git a/sdk/python/feast/infra/aws.py b/sdk/python/feast/infra/aws.py index e58a407370..f93553bf9f 100644 --- a/sdk/python/feast/infra/aws.py +++ b/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: diff --git a/sdk/python/feast/infra/gcp.py b/sdk/python/feast/infra/gcp.py index bad39fe456..c57450c876 100644 --- a/sdk/python/feast/infra/gcp.py +++ b/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: diff --git a/sdk/python/feast/infra/local.py b/sdk/python/feast/infra/local.py index e5851d9ad5..7304e262b1 100644 --- a/sdk/python/feast/infra/local.py +++ b/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: diff --git a/sdk/python/feast/registry.py b/sdk/python/feast/registry.py index bdf9a82d2b..0681ecb679 100644 --- a/sdk/python/feast/registry.py +++ b/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 diff --git a/sdk/python/feast/registry_store.py b/sdk/python/feast/registry_store.py index e538e8fa1f..22c8b0f5ed 100644 --- a/sdk/python/feast/registry_store.py +++ b/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. diff --git a/sdk/python/tests/integration/registration/test_cli.py b/sdk/python/tests/integration/registration/test_cli.py index c3d0f4b6b7..38663685e2 100644 --- a/sdk/python/tests/integration/registration/test_cli.py +++ b/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: