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>

sdk/python/feast/errors.py

@@ -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}'")
 
 
@@ -85,10 +85,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}`.'"
         )
 
 

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

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

@@ -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(

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

@@ -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):

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()

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:

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)

sdk/python/feast/infra/provider.py

@@ -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
@@ -156,24 +155,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(