Skip to content

Commit 9add91d

Browse files
committed
postgres tests
Signed-off-by: Achal Shah <achals@gmail.com>
1 parent dcd88bd commit 9add91d

2 files changed

Lines changed: 357 additions & 14 deletions

File tree

sdk/python/feast/infra/registry_stores/sql.py

Lines changed: 87 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -21,7 +21,12 @@
2121
from feast.base_feature_view import BaseFeatureView
2222
from feast.data_source import DataSource
2323
from feast.entity import Entity
24-
from feast.errors import DataSourceObjectNotFoundException, EntityNotFoundException
24+
from feast.errors import (
25+
DataSourceObjectNotFoundException,
26+
EntityNotFoundException,
27+
FeatureServiceNotFoundException,
28+
FeatureViewNotFoundException,
29+
)
2530
from feast.feature_service import FeatureService
2631
from feast.feature_view import FeatureView
2732
from feast.on_demand_feature_view import OnDemandFeatureView
@@ -161,6 +166,57 @@ def get_entity(self, name: str, project: str, allow_cache: bool = False) -> Enti
161166
return Entity.from_proto(entity_proto)
162167
raise EntityNotFoundException(name, project=project)
163168

169+
def get_feature_view(
170+
self, name: str, project: str, allow_cache: bool = False
171+
) -> FeatureView:
172+
with self.engine.connect() as conn:
173+
stmt = select(feature_views).where(
174+
feature_views.c.feature_view_name == name
175+
)
176+
row = conn.execute(stmt).first()
177+
if row:
178+
fv_proto = FeatureViewProto.FromString(row["feature_view_proto"])
179+
return FeatureView.from_proto(fv_proto)
180+
raise FeatureViewNotFoundException(name, project=project)
181+
182+
def get_on_demand_feature_view(
183+
self, name: str, project: str, allow_cache: bool = False
184+
) -> OnDemandFeatureView:
185+
with self.engine.connect() as conn:
186+
stmt = select(on_demand_feature_views).where(
187+
on_demand_feature_views.c.feature_view_name == name
188+
)
189+
row = conn.execute(stmt).first()
190+
if row:
191+
fv_proto = OnDemandFeatureViewProto.FromString(
192+
row["feature_view_proto"]
193+
)
194+
return OnDemandFeatureView.from_proto(fv_proto)
195+
raise FeatureViewNotFoundException(name, project=project)
196+
197+
def get_feature_service(
198+
self, name: str, project: str, allow_cache: bool = False
199+
) -> FeatureService:
200+
with self.engine.connect() as conn:
201+
stmt = select(feature_services).where(
202+
feature_services.c.feature_service_name == name
203+
)
204+
row = conn.execute(stmt).first()
205+
if row:
206+
fv_proto = FeatureServiceProto.FromString(row["feature_service_proto"])
207+
return FeatureService.from_proto(fv_proto)
208+
raise FeatureServiceNotFoundException(name, project=project)
209+
210+
def get_saved_dataset(
211+
self, name: str, project: str, allow_cache: bool = False
212+
) -> SavedDataset:
213+
pass
214+
215+
def get_validation_reference(
216+
self, name: str, project: str, allow_cache: bool = False
217+
) -> ValidationReference:
218+
pass
219+
164220
def list_entities(self, project: str, allow_cache: bool = False) -> List[Entity]:
165221
return self._list_objects(entities, EntityProto, Entity, "entity_proto")
166222

@@ -171,6 +227,25 @@ def delete_entity(self, name: str, project: str, commit: bool = True):
171227
if rows.rowcount < 1:
172228
raise EntityNotFoundException(name, project)
173229

230+
def delete_feature_view(self, name: str, project: str, commit: bool = True):
231+
deleted_count = 0
232+
for table in {feature_views, request_feature_views, on_demand_feature_views}:
233+
with self.engine.connect() as conn:
234+
stmt = delete(table).where(table.c.feature_view_name == name)
235+
rows = conn.execute(stmt)
236+
deleted_count += rows.rowcount
237+
if deleted_count == 0:
238+
raise FeatureViewNotFoundException(name, project)
239+
240+
def delete_feature_service(self, name: str, project: str, commit: bool = True):
241+
with self.engine.connect() as conn:
242+
stmt = delete(feature_services).where(
243+
feature_services.c.feature_service_name == name
244+
)
245+
rows = conn.execute(stmt)
246+
if rows.rowcount < 1:
247+
raise FeatureServiceNotFoundException(name, project)
248+
174249
def get_data_source(
175250
self, name: str, project: str, allow_cache: bool = False
176251
) -> DataSource:
@@ -305,25 +380,23 @@ def _apply_object(
305380
if hasattr(obj, "last_updated_timestamp"):
306381
obj.last_updated_timestamp = update_datetime
307382
if row:
383+
values = {
384+
proto_field_name: obj.to_proto().SerializeToString(),
385+
"last_updated_timestamp": update_time,
386+
}
308387
update_stmt = (
309388
update(table)
310389
.where(getattr(table.c, id_field_name) == name)
311-
.values(
312-
**{
313-
proto_field_name: obj.to_proto().SerializeToString(),
314-
"last_updated_timestamp": update_time,
315-
},
316-
)
390+
.values(values,)
317391
)
318392
conn.execute(update_stmt)
319393
else:
320-
insert_stmt = insert(feature_services).values(
321-
**{
322-
id_field_name: name,
323-
proto_field_name: obj.to_proto().SerializeToString(),
324-
"last_updated_timestamp": update_time,
325-
},
326-
)
394+
values = {
395+
id_field_name: name,
396+
proto_field_name: obj.to_proto().SerializeToString(),
397+
"last_updated_timestamp": update_time,
398+
}
399+
insert_stmt = insert(table).values(values,)
327400
conn.execute(insert_stmt)
328401

329402
def _list_objects(self, table, proto_class, python_class, proto_field_name):

0 commit comments

Comments
 (0)