Skip to content

Commit 5b4b07f

Browse files
authored
chore: Make loading online and offline classes lazy (#2646)
* chore: Make loading online and offline classes lazy Signed-off-by: Achal Shah <[email protected]> * fixes Signed-off-by: Achal Shah <[email protected]> * fixes Signed-off-by: Achal Shah <[email protected]> * nullability fixes Signed-off-by: Achal Shah <[email protected]>
1 parent 41a1da4 commit 5b4b07f

File tree

4 files changed

+79
-23
lines changed

4 files changed

+79
-23
lines changed

sdk/python/feast/infra/passthrough_provider.py

Lines changed: 18 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -39,12 +39,24 @@ def __init__(self, config: RepoConfig):
3939
super().__init__(config)
4040

4141
self.repo_config = config
42-
self.offline_store = get_offline_store_from_config(config.offline_store)
43-
self.online_store = (
44-
get_online_store_from_config(config.online_store)
45-
if config.online_store
46-
else None
47-
)
42+
self._offline_store = None
43+
self._online_store = None
44+
45+
@property
46+
def online_store(self):
47+
if not self._online_store and self.repo_config.online_store:
48+
self._online_store = get_online_store_from_config(
49+
self.repo_config.online_store
50+
)
51+
return self._online_store
52+
53+
@property
54+
def offline_store(self):
55+
if not self._offline_store:
56+
self._offline_store = get_offline_store_from_config(
57+
self.repo_config.offline_store
58+
)
59+
return self._offline_store
4860

4961
def update_infra(
5062
self,

sdk/python/feast/repo_config.py

Lines changed: 56 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@
66
import yaml
77
from pydantic import (
88
BaseModel,
9+
Field,
910
StrictInt,
1011
StrictStr,
1112
ValidationError,
@@ -107,10 +108,10 @@ class RepoConfig(FeastBaseModel):
107108
provider: StrictStr
108109
""" str: local or gcp or aws """
109110

110-
online_store: Any
111+
_online_config: Any = Field(alias="online_store")
111112
""" OnlineStoreConfig: Online store configuration (optional depending on provider) """
112113

113-
offline_store: Any
114+
_offline_config: Any = Field(alias="offline_store")
114115
""" OfflineStoreConfig: Offline store configuration (optional depending on provider) """
115116

116117
feature_server: Optional[Any]
@@ -126,19 +127,27 @@ class RepoConfig(FeastBaseModel):
126127
def __init__(self, **data: Any):
127128
super().__init__(**data)
128129

129-
if isinstance(self.online_store, Dict):
130-
self.online_store = get_online_config_from_type(self.online_store["type"])(
131-
**self.online_store
132-
)
133-
elif isinstance(self.online_store, str):
134-
self.online_store = get_online_config_from_type(self.online_store)()
135-
136-
if isinstance(self.offline_store, Dict):
137-
self.offline_store = get_offline_config_from_type(
138-
self.offline_store["type"]
139-
)(**self.offline_store)
140-
elif isinstance(self.offline_store, str):
141-
self.offline_store = get_offline_config_from_type(self.offline_store)()
130+
self._offline_store = None
131+
if "offline_store" in data:
132+
self._offline_config = data["offline_store"]
133+
else:
134+
if data["provider"] == "local":
135+
self._offline_config = "file"
136+
elif data["provider"] == "gcp":
137+
self._offline_config = "bigquery"
138+
elif data["provider"] == "aws":
139+
self._offline_config = "redshift"
140+
141+
self._online_store = None
142+
if "online_store" in data:
143+
self._online_config = data["online_store"]
144+
else:
145+
if data["provider"] == "local":
146+
self._online_config = "sqlite"
147+
elif data["provider"] == "gcp":
148+
self._online_config = "datastore"
149+
elif data["provider"] == "aws":
150+
self._online_config = "dynamodb"
142151

143152
if isinstance(self.feature_server, Dict):
144153
self.feature_server = get_feature_server_config_from_type(
@@ -151,6 +160,35 @@ def get_registry_config(self):
151160
else:
152161
return self.registry
153162

163+
@property
164+
def offline_store(self):
165+
if not self._offline_store:
166+
if isinstance(self._offline_config, Dict):
167+
self._offline_store = get_offline_config_from_type(
168+
self._offline_config["type"]
169+
)(**self._offline_config)
170+
elif isinstance(self._offline_config, str):
171+
self._offline_store = get_offline_config_from_type(
172+
self._offline_config
173+
)()
174+
elif self._offline_config:
175+
self._offline_store = self._offline_config
176+
return self._offline_store
177+
178+
@property
179+
def online_store(self):
180+
if not self._online_store:
181+
if isinstance(self._online_config, Dict):
182+
self._online_store = get_online_config_from_type(
183+
self._online_config["type"]
184+
)(**self._online_config)
185+
elif isinstance(self._online_config, str):
186+
self._online_store = get_online_config_from_type(self._online_config)()
187+
elif self._online_config:
188+
self._online_store = self._online_config
189+
190+
return self._online_store
191+
154192
@root_validator(pre=True)
155193
@log_exceptions
156194
def _validate_online_store_config(cls, values):
@@ -304,6 +342,9 @@ def write_to_path(self, repo_path: Path):
304342
sort_keys=False,
305343
)
306344

345+
class Config:
346+
allow_population_by_field_name = True
347+
307348

308349
class FeastConfigError(Exception):
309350
def __init__(self, error_message, config_path):

sdk/python/tests/integration/feature_repos/integration_test_repo_config.py

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -19,7 +19,7 @@ class IntegrationTestRepoConfig:
1919
"""
2020

2121
provider: str = "local"
22-
online_store: Union[str, Dict] = "sqlite"
22+
online_store: Optional[Union[str, Dict]] = "sqlite"
2323

2424
offline_store_creator: Type[DataSourceCreator] = FileDataSourceCreator
2525
online_store_creator: Optional[Type[OnlineStoreCreator]] = None
@@ -38,8 +38,10 @@ def __repr__(self) -> str:
3838
online_store_type = self.online_store.get("redis_type", "redis")
3939
else:
4040
online_store_type = self.online_store["type"]
41-
else:
41+
elif self.online_store:
4242
online_store_type = self.online_store.__name__
43+
else:
44+
online_store_type = "none"
4345
else:
4446
online_store_type = self.online_store_creator.__name__
4547

sdk/python/tests/unit/infra/online_store/test_dynamodb_online_store.py

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -39,6 +39,7 @@ def repo_config():
3939
project=PROJECT,
4040
provider=PROVIDER,
4141
online_store=DynamoDBOnlineStoreConfig(region=REGION),
42+
# online_store={"type": "dynamodb", "region": REGION},
4243
offline_store=FileOfflineStoreConfig(),
4344
)
4445

0 commit comments

Comments
 (0)