Skip to content

Commit 8ce51b4

Browse files
achalsachals
authored
Refactor OfflineStoreConfig classes into their owning modules (feast-dev#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>
1 parent e5bea11 commit 8ce51b4

13 files changed

Lines changed: 152 additions & 145 deletions

File tree

‎sdk/python/feast/errors.py‎

Lines changed: 5 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -50,7 +50,7 @@ def __init__(self, provider_name):
5050

5151

5252
class FeastModuleImportError(Exception):
53-
def __init__(self, module_name, module_type="provider"):
53+
def __init__(self, module_name: str, module_type: str):
5454
super().__init__(f"Could not import {module_type} module '{module_name}'")
5555

5656

@@ -85,10 +85,11 @@ def __init__(self, online_store_class_name: str):
8585
)
8686

8787

88-
class FeastOnlineStoreConfigInvalidName(Exception):
89-
def __init__(self, online_store_config_class_name: str):
88+
class FeastClassInvalidName(Exception):
89+
def __init__(self, class_name: str, class_type: str):
9090
super().__init__(
91-
f"Online Store Config Class '{online_store_config_class_name}' should end with the string `OnlineStoreConfig`.'"
91+
f"Config Class '{class_name}' "
92+
f"should end with the string `{class_type}`.'"
9293
)
9394

9495

‎sdk/python/feast/importer.py‎

Lines changed: 28 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,28 @@
1+
import importlib
2+
3+
from feast import errors
4+
5+
6+
def get_class_from_type(module_name: str, class_name: str, class_type: str):
7+
if not class_name.endswith(class_type):
8+
raise errors.FeastClassInvalidName(class_name, class_type)
9+
10+
# Try importing the module that contains the custom provider
11+
try:
12+
module = importlib.import_module(module_name)
13+
except Exception as e:
14+
# The original exception can be anything - either module not found,
15+
# or any other kind of error happening during the module import time.
16+
# So we should include the original error as well in the stack trace.
17+
raise errors.FeastModuleImportError(module_name, class_type) from e
18+
19+
# Try getting the provider class definition
20+
try:
21+
_class = getattr(module, class_name)
22+
except AttributeError:
23+
# This can only be one type of error, when class_name attribute does not exist in the module
24+
# So we don't have to include the original exception here
25+
raise errors.FeastClassImportError(
26+
module_name, class_name, class_type=class_type
27+
) from None
28+
return _class

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

Lines changed: 16 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,8 @@
77
import pandas
88
import pyarrow
99
from jinja2 import BaseLoader, Environment
10+
from pydantic import StrictStr
11+
from pydantic.typing import Literal
1012
from tenacity import retry, stop_after_delay, wait_fixed
1113

1214
from feast import errors
@@ -20,7 +22,7 @@
2022
_get_requested_feature_views_to_features_dict,
2123
)
2224
from feast.registry import Registry
23-
from feast.repo_config import BigQueryOfflineStoreConfig, RepoConfig
25+
from feast.repo_config import FeastConfigBaseModel, RepoConfig
2426

2527
try:
2628
from google.api_core.exceptions import NotFound
@@ -34,6 +36,19 @@
3436
raise FeastExtrasDependencyImportError("gcp", str(e))
3537

3638

39+
class BigQueryOfflineStoreConfig(FeastConfigBaseModel):
40+
""" Offline store config for GCP BigQuery """
41+
42+
type: Literal["bigquery"] = "bigquery"
43+
""" Offline store type selector"""
44+
45+
dataset: StrictStr = "feast"
46+
""" (optional) BigQuery Dataset name for temporary tables """
47+
48+
project_id: Optional[StrictStr] = None
49+
""" (optional) GCP project name used for the BigQuery offline store """
50+
51+
3752
class BigQueryOfflineStore(OfflineStore):
3853
@staticmethod
3954
def pull_latest_from_table_or_query(

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

Lines changed: 9 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@
44
import pandas as pd
55
import pyarrow
66
import pytz
7+
from pydantic.typing import Literal
78

89
from feast.data_source import DataSource, FileSource
910
from feast.errors import FeastJoinKeysDuringMaterialization
@@ -15,7 +16,14 @@
1516
_run_field_mapping,
1617
)
1718
from feast.registry import Registry
18-
from feast.repo_config import RepoConfig
19+
from feast.repo_config import FeastConfigBaseModel, RepoConfig
20+
21+
22+
class FileOfflineStoreConfig(FeastConfigBaseModel):
23+
""" Offline store config for local (file-based) store """
24+
25+
type: Literal["file"] = "file"
26+
""" Offline store type selector"""
1927

2028

2129
class FileRetrievalJob(RetrievalJob):

‎sdk/python/feast/infra/offline_stores/helpers.py‎

Lines changed: 26 additions & 36 deletions
Original file line numberDiff line numberDiff line change
@@ -1,41 +1,31 @@
1-
from feast.data_source import BigQuerySource, DataSource, FileSource
2-
from feast.errors import FeastOfflineStoreUnsupportedDataSource
1+
import importlib
2+
from typing import Any
3+
4+
from feast import errors
35
from feast.infra.offline_stores.offline_store import OfflineStore
4-
from feast.repo_config import (
5-
BigQueryOfflineStoreConfig,
6-
FileOfflineStoreConfig,
7-
OfflineStoreConfig,
8-
)
96

107

11-
def get_offline_store_from_config(
12-
offline_store_config: OfflineStoreConfig,
13-
) -> OfflineStore:
8+
def get_offline_store_from_config(offline_store_config: Any,) -> OfflineStore:
149
"""Get the offline store from offline store config"""
1510

16-
if isinstance(offline_store_config, FileOfflineStoreConfig):
17-
from feast.infra.offline_stores.file import FileOfflineStore
18-
19-
return FileOfflineStore()
20-
elif isinstance(offline_store_config, BigQueryOfflineStoreConfig):
21-
from feast.infra.offline_stores.bigquery import BigQueryOfflineStore
22-
23-
return BigQueryOfflineStore()
24-
25-
raise ValueError(f"Unsupported offline store config '{offline_store_config}'")
26-
27-
28-
def assert_offline_store_supports_data_source(
29-
offline_store_config: OfflineStoreConfig, data_source: DataSource
30-
):
31-
if (
32-
isinstance(offline_store_config, FileOfflineStoreConfig)
33-
and isinstance(data_source, FileSource)
34-
) or (
35-
isinstance(offline_store_config, BigQueryOfflineStoreConfig)
36-
and isinstance(data_source, BigQuerySource)
37-
):
38-
return
39-
raise FeastOfflineStoreUnsupportedDataSource(
40-
offline_store_config.type, data_source.__class__.__name__
41-
)
11+
module_name = offline_store_config.__module__
12+
qualified_name = type(offline_store_config).__name__
13+
store_class_name = qualified_name.replace("Config", "")
14+
try:
15+
module = importlib.import_module(module_name)
16+
except Exception as e:
17+
# The original exception can be anything - either module not found,
18+
# or any other kind of error happening during the module import time.
19+
# So we should include the original error as well in the stack trace.
20+
raise errors.FeastModuleImportError(module_name, "OfflineStore") from e
21+
22+
# Try getting the provider class definition
23+
try:
24+
offline_store_class = getattr(module, store_class_name)
25+
except AttributeError:
26+
# This can only be one type of error, when class_name attribute does not exist in the module
27+
# So we don't have to include the original exception here
28+
raise errors.FeastClassImportError(
29+
module_name, store_class_name, class_type="OfflineStore"
30+
) from None
31+
return offline_store_class()

‎sdk/python/feast/infra/online_stores/helpers.py‎

Lines changed: 1 addition & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -22,9 +22,7 @@ def get_online_store_from_config(online_store_config: Any,) -> OnlineStore:
2222
# The original exception can be anything - either module not found,
2323
# or any other kind of error happening during the module import time.
2424
# So we should include the original error as well in the stack trace.
25-
raise errors.FeastModuleImportError(
26-
module_name, module_type="OnlineStore"
27-
) from e
25+
raise errors.FeastModuleImportError(module_name, "OnlineStore") from e
2826

2927
# Try getting the provider class definition
3028
try:

‎sdk/python/feast/infra/online_stores/sqlite.py‎

Lines changed: 7 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -34,7 +34,9 @@
3434
class SqliteOnlineStoreConfig(FeastConfigBaseModel):
3535
""" Online store config for local (SQLite-based) store """
3636

37-
type: Literal["sqlite"] = "sqlite"
37+
type: Literal[
38+
"sqlite", "feast.infra.online_stores.sqlite.SqliteOnlineStore"
39+
] = "sqlite"
3840
""" Online store type selector"""
3941

4042
path: StrictStr = "data/online.db"
@@ -51,7 +53,10 @@ class SqliteOnlineStore(OnlineStore):
5153

5254
@staticmethod
5355
def _get_db_path(config: RepoConfig) -> str:
54-
assert config.online_store.type == "sqlite"
56+
assert (
57+
config.online_store.type == "sqlite"
58+
or config.online_store.type.endswith("SqliteOnlineStore")
59+
)
5560

5661
if config.repo_path and not Path(config.online_store.path).is_absolute():
5762
db_path = str(config.repo_path / config.online_store.path)

‎sdk/python/feast/infra/provider.py‎

Lines changed: 4 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,4 @@
11
import abc
2-
import importlib
32
from datetime import datetime
43
from pathlib import Path
54
from typing import Any, Callable, Dict, List, Optional, Sequence, Tuple, Union
@@ -8,7 +7,7 @@
87
import pyarrow
98
from tqdm import tqdm
109

11-
from feast import errors
10+
from feast import errors, importer
1211
from feast.entity import Entity
1312
from feast.feature_table import FeatureTable
1413
from feast.feature_view import FeatureView
@@ -156,24 +155,9 @@ def get_provider(config: RepoConfig, repo_path: Path) -> Provider:
156155
# For example, provider 'foo.bar.MyProvider' will be parsed into 'foo.bar' and 'MyProvider'
157156
module_name, class_name = config.provider.rsplit(".", 1)
158157

159-
# Try importing the module that contains the custom provider
160-
try:
161-
module = importlib.import_module(module_name)
162-
except Exception as e:
163-
# The original exception can be anything - either module not found,
164-
# or any other kind of error happening during the module import time.
165-
# So we should include the original error as well in the stack trace.
166-
raise errors.FeastModuleImportError(module_name) from e
167-
168-
# Try getting the provider class definition
169-
try:
170-
ProviderCls = getattr(module, class_name)
171-
except AttributeError:
172-
# This can only be one type of error, when class_name attribute does not exist in the module
173-
# So we don't have to include the original exception here
174-
raise errors.FeastClassImportError(module_name, class_name) from None
175-
176-
return ProviderCls(config, repo_path)
158+
cls = importer.get_class_from_type(module_name, class_name, "Provider")
159+
160+
return cls(config, repo_path)
177161

178162

179163
def _get_requested_feature_views_to_features_dict(

0 commit comments

Comments
 (0)