Skip to content

Commit ec1e61d

Browse files
feat: Add Snowflake Registry (feast-dev#3363)
* feat: Add Snowflake Registry Signed-off-by: Miles Adkins <miles.adkins@snowflake.com> * Fix custom registry store parsing Signed-off-by: Danny Chiao <danny@tecton.ai> * fix test Signed-off-by: Danny Chiao <danny@tecton.ai> --------- Signed-off-by: Miles Adkins <miles.adkins@snowflake.com> Signed-off-by: Danny Chiao <danny@tecton.ai> Co-authored-by: Danny Chiao <danny@tecton.ai>
1 parent f0eea99 commit ec1e61d

17 files changed

Lines changed: 1377 additions & 51 deletions

File tree

docs/tutorials/azure/notebooks/src/score.py

Lines changed: 5 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -6,9 +6,11 @@
66
import json
77
import joblib
88
from feast import FeatureStore, RepoConfig
9-
from feast.infra.registry.registry import RegistryConfig
9+
from feast.repo_config import RegistryConfig
1010

11-
from feast.infra.offline_stores.contrib.mssql_offline_store.mssql import MsSqlServerOfflineStoreConfig
11+
from feast.infra.offline_stores.contrib.mssql_offline_store.mssql import (
12+
MsSqlServerOfflineStoreConfig,
13+
)
1214
from feast.infra.online_stores.redis import RedisOnlineStoreConfig, RedisOnlineStore
1315

1416

@@ -73,4 +75,4 @@ def run(raw_data):
7375
y_hat = model.predict(data)
7476
return y_hat.tolist()
7577
else:
76-
return 0.0
78+
return 0.0

sdk/python/feast/errors.py

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -132,6 +132,11 @@ def __init__(self):
132132
super().__init__("Provider is not set, but is required")
133133

134134

135+
class FeastRegistryNotSetError(Exception):
136+
def __init__(self):
137+
super().__init__("Registry is not set, but is required")
138+
139+
135140
class FeastFeatureServerTypeSetError(Exception):
136141
def __init__(self, feature_server_type: str):
137142
super().__init__(
@@ -146,6 +151,13 @@ def __init__(self, feature_server_type: str):
146151
)
147152

148153

154+
class FeastRegistryTypeInvalidError(Exception):
155+
def __init__(self, registry_type: str):
156+
super().__init__(
157+
f"Feature server type was set to {registry_type}, but this type is invalid"
158+
)
159+
160+
149161
class FeastModuleImportError(Exception):
150162
def __init__(self, module_name: str, class_name: str):
151163
super().__init__(

sdk/python/feast/feature_store.py

Lines changed: 8 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -164,9 +164,15 @@ def __init__(
164164
self.repo_path, utils.get_default_yaml_file_path(self.repo_path)
165165
)
166166

167-
registry_config = self.config.get_registry_config()
167+
registry_config = self.config.registry
168168
if registry_config.registry_type == "sql":
169169
self._registry = SqlRegistry(registry_config, self.config.project, None)
170+
elif registry_config.registry_type == "snowflake.registry":
171+
from feast.infra.registry.snowflake import SnowflakeRegistry
172+
173+
self._registry = SnowflakeRegistry(
174+
registry_config, self.config.project, None
175+
)
170176
else:
171177
r = Registry(self.config.project, registry_config, repo_path=self.repo_path)
172178
r._initialize_registry(self.config.project)
@@ -209,7 +215,7 @@ def refresh_registry(self):
209215
greater than 0, then once the cache becomes stale (more time than the TTL has passed), a new cache will be
210216
downloaded synchronously, which may increase latencies if the triggering method is get_online_features().
211217
"""
212-
registry_config = self.config.get_registry_config()
218+
registry_config = self.config.registry
213219
registry = Registry(
214220
self.config.project, registry_config, repo_path=self.repo_path
215221
)

sdk/python/feast/infra/materialization/snowflake_engine.py

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -132,6 +132,10 @@ def update(
132132
# if the stage already exists,
133133
# assumes that the materialization functions have been deployed
134134
if f"feast_{project}" in stage_list["name"].tolist():
135+
click.echo(
136+
f"Materialization functions for {Style.BRIGHT + Fore.GREEN}{project}{Style.RESET_ALL} already detected."
137+
)
138+
click.echo()
135139
return None
136140

137141
click.echo(

sdk/python/feast/infra/registry/proto_registry_utils.py

Lines changed: 19 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -9,7 +9,6 @@
99
EntityNotFoundException,
1010
FeatureServiceNotFoundException,
1111
FeatureViewNotFoundException,
12-
OnDemandFeatureViewNotFoundException,
1312
SavedDatasetNotFound,
1413
ValidationReferenceNotFound,
1514
)
@@ -98,7 +97,7 @@ def get_on_demand_feature_view(
9897
and on_demand_feature_view.spec.name == name
9998
):
10099
return OnDemandFeatureView.from_proto(on_demand_feature_view)
101-
raise OnDemandFeatureViewNotFoundException(name, project=project)
100+
raise FeatureViewNotFoundException(name, project=project)
102101

103102

104103
def get_data_source(
@@ -138,10 +137,6 @@ def get_validation_reference(
138137
raise ValidationReferenceNotFound(name, project=project)
139138

140139

141-
def list_validation_references(registry_proto: RegistryProto):
142-
return registry_proto.validation_references
143-
144-
145140
def list_feature_services(
146141
registry_proto: RegistryProto, project: str, allow_cache: bool = False
147142
) -> List[FeatureService]:
@@ -215,13 +210,25 @@ def list_data_sources(registry_proto: RegistryProto, project: str) -> List[DataS
215210

216211

217212
def list_saved_datasets(
218-
registry_proto: RegistryProto, project: str, allow_cache: bool = False
213+
registry_proto: RegistryProto, project: str
219214
) -> List[SavedDataset]:
220-
return [
221-
SavedDataset.from_proto(saved_dataset)
222-
for saved_dataset in registry_proto.saved_datasets
223-
if saved_dataset.spec.project == project
224-
]
215+
saved_datasets = []
216+
for saved_dataset in registry_proto.saved_datasets:
217+
if saved_dataset.project == project:
218+
saved_datasets.append(SavedDataset.from_proto(saved_dataset))
219+
return saved_datasets
220+
221+
222+
def list_validation_references(
223+
registry_proto: RegistryProto, project: str
224+
) -> List[ValidationReference]:
225+
validation_references = []
226+
for validation_reference in registry_proto.validation_references:
227+
if validation_reference.project == project:
228+
validation_references.append(
229+
ValidationReference.from_proto(validation_reference)
230+
)
231+
return validation_references
225232

226233

227234
def list_project_metadata(

sdk/python/feast/infra/registry/registry.py

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -174,6 +174,10 @@ def __new__(
174174
from feast.infra.registry.sql import SqlRegistry
175175

176176
return SqlRegistry(registry_config, project, repo_path)
177+
elif registry_config and registry_config.registry_type == "snowflake.registry":
178+
from feast.infra.registry.snowflake import SnowflakeRegistry
179+
180+
return SnowflakeRegistry(registry_config, project, repo_path)
177181
else:
178182
return super(Registry, cls).__new__(cls)
179183

@@ -731,7 +735,7 @@ def list_validation_references(
731735
registry_proto = self._get_registry_proto(
732736
project=project, allow_cache=allow_cache
733737
)
734-
return proto_registry_utils.list_validation_references(registry_proto)
738+
return proto_registry_utils.list_validation_references(registry_proto, project)
735739

736740
def delete_validation_reference(self, name: str, project: str, commit: bool = True):
737741
registry_proto = self._prepare_registry_for_changes(project)

0 commit comments

Comments
 (0)