Refactor OfflineStoreConfig classes into their owning modules (#1657)

* Refactor OfflineStoreConfig classes into their owning modules Signed-off-by: Achal Shah <achals@gmail.com> * Fix error string Signed-off-by: Achal Shah <achals@gmail.com> * Generic error class Signed-off-by: Achal Shah <achals@gmail.com> * Merge conflicts Signed-off-by: Achal Shah <achals@gmail.com> * make the store type work, and add a test that uses the fully qualified name of the OnlineStore Signed-off-by: Achal Shah <achals@gmail.com> * Address comments from previous PR Signed-off-by: Achal Shah <achals@gmail.com> * CR updates Signed-off-by: Achal Shah <achals@gmail.com> Signed-off-by: Mwad22 <51929507+Mwad22@users.noreply.github.com>
feast-dev · feast-ci-bot · Jul 8, 2021 · Jun 12, 2021 · Jun 16, 2021 · Jun 16, 2021
commit 4ab4c607a5fddfc11dfc637836f903769ad50f28
@@ -50,7 +50,7 @@ def __init__(self, provider_name):
 
 
 class FeastModuleImportError(Exception):
-    def __init__(self, module_name, module_type="provider"):
+    def __init__(self, module_name: str, module_type: str):
         super().__init__(f"Could not import {module_type} module '{module_name}'")
 
 
@@ -92,10 +92,11 @@ def __init__(self, online_store_class_name: str):
         )
 
 
-class FeastOnlineStoreConfigInvalidName(Exception):
-    def __init__(self, online_store_config_class_name: str):
+class FeastClassInvalidName(Exception):
+    def __init__(self, class_name: str, class_type: str):
         super().__init__(
-            f"Online Store Config Class '{online_store_config_class_name}' should end with the string `OnlineStoreConfig`.'"
+            f"Config Class '{class_name}' "
+            f"should end with the string `{class_type}`.'"
         )
 
 

diff --git a/sdk/python/feast/importer.py b/sdk/python/feast/importer.py
@@ -0,0 +1,28 @@
+import importlib
+
+from feast import errors
+
+
+def get_class_from_type(module_name: str, class_name: str, class_type: str):
+    if not class_name.endswith(class_type):
+        raise errors.FeastClassInvalidName(class_name, class_type)
+
+    # Try importing the module that contains the custom provider
+    try:
+        module = importlib.import_module(module_name)
+    except Exception as e:
+        # The original exception can be anything - either module not found,
+        # or any other kind of error happening during the module import time.
+        # So we should include the original error as well in the stack trace.
+        raise errors.FeastModuleImportError(module_name, class_type) from e
+
+    # Try getting the provider class definition
+    try:
+        _class = getattr(module, class_name)
+    except AttributeError:
+        # This can only be one type of error, when class_name attribute does not exist in the module
+        # So we don't have to include the original exception here
+        raise errors.FeastClassImportError(
+            module_name, class_name, class_type=class_type
+        ) from None
+    return _class
@@ -7,6 +7,8 @@
 import pandas
 import pyarrow
 from jinja2 import BaseLoader, Environment
+from pydantic import StrictStr
+from pydantic.typing import Literal
 from tenacity import retry, stop_after_delay, wait_fixed
 
 from feast import errors
@@ -20,7 +22,7 @@
     _get_requested_feature_views_to_features_dict,
 )
 from feast.registry import Registry
-from feast.repo_config import BigQueryOfflineStoreConfig, RepoConfig
+from feast.repo_config import FeastConfigBaseModel, RepoConfig
 
 try:
     from google.api_core.exceptions import NotFound
@@ -34,6 +36,19 @@
     raise FeastExtrasDependencyImportError("gcp", str(e))
 
 
+class BigQueryOfflineStoreConfig(FeastConfigBaseModel):
+    """ Offline store config for GCP BigQuery """
+
+    type: Literal["bigquery"] = "bigquery"
+    """ Offline store type selector"""
+
+    dataset: StrictStr = "feast"
+    """ (optional) BigQuery Dataset name for temporary tables """
+
+    project_id: Optional[StrictStr] = None
+    """ (optional) GCP project name used for the BigQuery offline store """
+
+
 class BigQueryOfflineStore(OfflineStore):
     @staticmethod
     def pull_latest_from_table_or_query(

@@ -4,6 +4,7 @@
 import pandas as pd
 import pyarrow
 import pytz
+from pydantic.typing import Literal
 
 from feast.data_source import DataSource, FileSource
 from feast.errors import FeastJoinKeysDuringMaterialization
@@ -15,7 +16,14 @@
     _run_field_mapping,
 )
 from feast.registry import Registry
-from feast.repo_config import RepoConfig
+from feast.repo_config import FeastConfigBaseModel, RepoConfig
+
+
+class FileOfflineStoreConfig(FeastConfigBaseModel):
+    """ Offline store config for local (file-based) store """
+
+    type: Literal["file"] = "file"
+    """ Offline store type selector"""
 
 
 class FileRetrievalJob(RetrievalJob):

diff --git a/sdk/python/feast/infra/offline_stores/helpers.py b/sdk/python/feast/infra/offline_stores/helpers.py
@@ -1,41 +1,31 @@
-from feast.data_source import BigQuerySource, DataSource, FileSource
-from feast.errors import FeastOfflineStoreUnsupportedDataSource
+import importlib
+from typing import Any
+
+from feast import errors
 from feast.infra.offline_stores.offline_store import OfflineStore
-from feast.repo_config import (
-    BigQueryOfflineStoreConfig,
-    FileOfflineStoreConfig,
-    OfflineStoreConfig,
-)
 
 
-def get_offline_store_from_config(
-    offline_store_config: OfflineStoreConfig,
-) -> OfflineStore:
+def get_offline_store_from_config(offline_store_config: Any,) -> OfflineStore:
     """Get the offline store from offline store config"""
 
-    if isinstance(offline_store_config, FileOfflineStoreConfig):
-        from feast.infra.offline_stores.file import FileOfflineStore
-
-        return FileOfflineStore()
-    elif isinstance(offline_store_config, BigQueryOfflineStoreConfig):
-        from feast.infra.offline_stores.bigquery import BigQueryOfflineStore
-
-        return BigQueryOfflineStore()
-
-    raise ValueError(f"Unsupported offline store config '{offline_store_config}'")
-
-
-def assert_offline_store_supports_data_source(
-    offline_store_config: OfflineStoreConfig, data_source: DataSource
-):
-    if (
-        isinstance(offline_store_config, FileOfflineStoreConfig)
-        and isinstance(data_source, FileSource)
-    ) or (
-        isinstance(offline_store_config, BigQueryOfflineStoreConfig)
-        and isinstance(data_source, BigQuerySource)
-    ):
-        return
-    raise FeastOfflineStoreUnsupportedDataSource(
-        offline_store_config.type, data_source.__class__.__name__
-    )
+    module_name = offline_store_config.__module__
+    qualified_name = type(offline_store_config).__name__
+    store_class_name = qualified_name.replace("Config", "")
+    try:
+        module = importlib.import_module(module_name)
+    except Exception as e:
+        # The original exception can be anything - either module not found,
+        # or any other kind of error happening during the module import time.
+        # So we should include the original error as well in the stack trace.
+        raise errors.FeastModuleImportError(module_name, "OfflineStore") from e
+
+    # Try getting the provider class definition
+    try:
+        offline_store_class = getattr(module, store_class_name)
+    except AttributeError:
+        # This can only be one type of error, when class_name attribute does not exist in the module
+        # So we don't have to include the original exception here
+        raise errors.FeastClassImportError(
+            module_name, store_class_name, class_type="OfflineStore"
+        ) from None
+    return offline_store_class()
diff --git a/sdk/python/feast/infra/online_stores/helpers.py b/sdk/python/feast/infra/online_stores/helpers.py
@@ -22,9 +22,7 @@ def get_online_store_from_config(online_store_config: Any,) -> OnlineStore:
         # The original exception can be anything - either module not found,
         # or any other kind of error happening during the module import time.
         # So we should include the original error as well in the stack trace.
-        raise errors.FeastModuleImportError(
-            module_name, module_type="OnlineStore"
-        ) from e
+        raise errors.FeastModuleImportError(module_name, "OnlineStore") from e
 
     # Try getting the provider class definition
     try:

diff --git a/sdk/python/feast/infra/online_stores/sqlite.py b/sdk/python/feast/infra/online_stores/sqlite.py
@@ -34,7 +34,9 @@
 class SqliteOnlineStoreConfig(FeastConfigBaseModel):
     """ Online store config for local (SQLite-based) store """
 
-    type: Literal["sqlite"] = "sqlite"
+    type: Literal[
+        "sqlite", "feast.infra.online_stores.sqlite.SqliteOnlineStore"
+    ] = "sqlite"
     """ Online store type selector"""
 
     path: StrictStr = "data/online.db"
@@ -51,7 +53,10 @@ class SqliteOnlineStore(OnlineStore):
 
     @staticmethod
     def _get_db_path(config: RepoConfig) -> str:
-        assert config.online_store.type == "sqlite"
+        assert (
+            config.online_store.type == "sqlite"
+            or config.online_store.type.endswith("SqliteOnlineStore")
+        )
 
         if config.repo_path and not Path(config.online_store.path).is_absolute():
             db_path = str(config.repo_path / config.online_store.path)

@@ -1,5 +1,4 @@
 import abc
-import importlib
 from datetime import datetime
 from pathlib import Path
 from typing import Any, Callable, Dict, List, Optional, Sequence, Tuple, Union
@@ -8,7 +7,7 @@
 import pyarrow
 from tqdm import tqdm
 
-from feast import errors
+from feast import errors, importer
 from feast.entity import Entity
 from feast.feature_table import FeatureTable
 from feast.feature_view import FeatureView
@@ -157,24 +156,9 @@ def get_provider(config: RepoConfig, repo_path: Path) -> Provider:
         # For example, provider 'foo.bar.MyProvider' will be parsed into 'foo.bar' and 'MyProvider'
         module_name, class_name = config.provider.rsplit(".", 1)
 
-        # Try importing the module that contains the custom provider
-        try:
-            module = importlib.import_module(module_name)
-        except Exception as e:
-            # The original exception can be anything - either module not found,
-            # or any other kind of error happening during the module import time.
-            # So we should include the original error as well in the stack trace.
-            raise errors.FeastModuleImportError(module_name) from e
-
-        # Try getting the provider class definition
-        try:
-            ProviderCls = getattr(module, class_name)
-        except AttributeError:
-            # This can only be one type of error, when class_name attribute does not exist in the module
-            # So we don't have to include the original exception here
-            raise errors.FeastClassImportError(module_name, class_name) from None
-
-        return ProviderCls(config, repo_path)
+        cls = importer.get_class_from_type(module_name, class_name, "Provider")
+
+        return cls(config, repo_path)
 
 
 def _get_requested_feature_views_to_features_dict(