Skip to content

Commit 4a2c84c

Browse files
achalsachals
committed
Fix CR comments
Signed-off-by: Achal Shah <achals@gmail.com>
1 parent 629083c commit 4a2c84c

2 files changed

Lines changed: 91 additions & 166 deletions

File tree

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

Lines changed: 91 additions & 52 deletions
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,8 @@
2626
EntityNotFoundException,
2727
FeatureServiceNotFoundException,
2828
FeatureViewNotFoundException,
29+
SavedDatasetNotFound,
30+
ValidationReferenceNotFound,
2931
)
3032
from feast.feature_service import FeatureService
3133
from feast.feature_view import FeatureView
@@ -43,6 +45,9 @@
4345
RequestFeatureView as RequestFeatureViewProto,
4446
)
4547
from feast.protos.feast.core.SavedDataset_pb2 import SavedDataset as SavedDatasetProto
48+
from feast.protos.feast.core.ValidationProfile_pb2 import (
49+
ValidationReference as ValidationReferenceProto,
50+
)
4651
from feast.registry import Registry
4752
from feast.repo_config import RegistryConfig
4853
from feast.request_feature_view import RequestFeatureView
@@ -53,7 +58,7 @@
5358
entities = Table(
5459
"entities",
5560
metadata,
56-
Column("entity_id", String(50), primary_key=True),
61+
Column("entity_name", String(50), primary_key=True),
5762
Column("last_updated_timestamp", BigInteger, nullable=False),
5863
Column("entity_proto", LargeBinary, nullable=False),
5964
)
@@ -91,14 +96,6 @@
9196
Column("feature_view_proto", LargeBinary, nullable=False),
9297
)
9398

94-
feature_user_metadata = Table(
95-
"feature_metadata",
96-
metadata,
97-
Column("feature_name", String(50), primary_key=True),
98-
Column("last_updated_timestamp", BigInteger, nullable=False),
99-
Column("feature_metadata_binary", LargeBinary, nullable=False),
100-
)
101-
10299
feature_services = Table(
103100
"feature_services",
104101
metadata,
@@ -139,9 +136,10 @@ def __init__(
139136

140137
def teardown(self):
141138
for t in {
139+
entities,
140+
data_sources,
142141
feature_views,
143142
feature_services,
144-
data_sources,
145143
on_demand_feature_views,
146144
request_feature_views,
147145
saved_datasets,
@@ -155,74 +153,96 @@ def refresh(self):
155153
pass
156154

157155
def apply_entity(self, entity: Entity, project: str, commit: bool = True):
158-
return self._apply_object(entities, "entity_id", entity, "entity_proto")
156+
return self._apply_object(entities, "entity_name", entity, "entity_proto")
159157

160158
def get_entity(self, name: str, project: str, allow_cache: bool = False) -> Entity:
161-
with self.engine.connect() as conn:
162-
stmt = select(entities).where(entities.c.entity_id == name)
163-
row = conn.execute(stmt).first()
164-
if row:
165-
entity_proto = EntityProto.FromString(row["entity_proto"])
166-
return Entity.from_proto(entity_proto)
167-
raise EntityNotFoundException(name, project=project)
159+
return self._get_object(
160+
entities,
161+
name,
162+
project,
163+
EntityProto,
164+
Entity,
165+
"entity_name",
166+
"entity_proto",
167+
EntityNotFoundException,
168+
)
168169

169170
def get_feature_view(
170171
self, name: str, project: str, allow_cache: bool = False
171172
) -> 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)
173+
return self._get_object(
174+
feature_views,
175+
name,
176+
project,
177+
FeatureViewProto,
178+
FeatureView,
179+
"feature_view_name",
180+
"feature_view_proto",
181+
FeatureViewNotFoundException,
182+
)
181183

182184
def get_on_demand_feature_view(
183185
self, name: str, project: str, allow_cache: bool = False
184186
) -> 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)
187+
return self._get_object(
188+
on_demand_feature_views,
189+
name,
190+
project,
191+
OnDemandFeatureViewProto,
192+
OnDemandFeatureView,
193+
"feature_view_name",
194+
"feature_view_proto",
195+
FeatureViewNotFoundException,
196+
)
196197

197198
def get_feature_service(
198199
self, name: str, project: str, allow_cache: bool = False
199200
) -> 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)
201+
return self._get_object(
202+
feature_services,
203+
name,
204+
project,
205+
FeatureServiceProto,
206+
FeatureService,
207+
"feature_service_name",
208+
"feature_service_proto",
209+
FeatureServiceNotFoundException,
210+
)
209211

210212
def get_saved_dataset(
211213
self, name: str, project: str, allow_cache: bool = False
212214
) -> SavedDataset:
213-
pass
215+
return self._get_object(
216+
saved_datasets,
217+
name,
218+
project,
219+
SavedDatasetProto,
220+
SavedDataset,
221+
"saved_dataset_name",
222+
"saved_dataset_proto",
223+
SavedDatasetNotFound,
224+
)
214225

215226
def get_validation_reference(
216227
self, name: str, project: str, allow_cache: bool = False
217228
) -> ValidationReference:
218-
pass
229+
return self._get_object(
230+
validation_references,
231+
name,
232+
project,
233+
ValidationReferenceProto,
234+
ValidationReference,
235+
"validation_reference_name",
236+
"validation_reference_proto",
237+
ValidationReferenceNotFound,
238+
)
219239

220240
def list_entities(self, project: str, allow_cache: bool = False) -> List[Entity]:
221241
return self._list_objects(entities, EntityProto, Entity, "entity_proto")
222242

223243
def delete_entity(self, name: str, project: str, commit: bool = True):
224244
with self.engine.connect() as conn:
225-
stmt = delete(entities).where(entities.c.entity_id == name)
245+
stmt = delete(entities).where(entities.c.entity_name == name)
226246
rows = conn.execute(stmt)
227247
if rows.rowcount < 1:
228248
raise EntityNotFoundException(name, project)
@@ -250,7 +270,7 @@ def get_data_source(
250270
self, name: str, project: str, allow_cache: bool = False
251271
) -> DataSource:
252272
with self.engine.connect() as conn:
253-
stmt = select(data_sources).where(data_sources.c.entity_id == name)
273+
stmt = select(data_sources).where(data_sources.c.entity_name == name)
254274
row = conn.execute(stmt).first()
255275
if row:
256276
ds_proto = DataSourceProto.FromString(row["data_source_proto"])
@@ -299,7 +319,7 @@ def apply_feature_service(
299319

300320
def delete_data_source(self, name: str, project: str, commit: bool = True):
301321
with self.engine.connect() as conn:
302-
stmt = delete(data_sources).where(data_sources.c.entity_id == name)
322+
stmt = delete(data_sources).where(data_sources.c.data_source_name == name)
303323
rows = conn.execute(stmt)
304324
if rows.rowcount < 1:
305325
raise DataSourceObjectNotFoundException(name, project)
@@ -411,3 +431,22 @@ def _list_objects(self, table, proto_class, python_class, proto_field_name):
411431
for row in rows
412432
]
413433
return []
434+
435+
def _get_object(
436+
self,
437+
table,
438+
name,
439+
project,
440+
proto_class,
441+
python_class,
442+
id_field_name,
443+
proto_field_name,
444+
not_found_exception,
445+
):
446+
with self.engine.connect() as conn:
447+
stmt = select(table).where(getattr(table.c, id_field_name) == name)
448+
row = conn.execute(stmt).first()
449+
if row:
450+
_proto = proto_class.FromString(row[proto_field_name])
451+
return python_class.from_proto(_proto)
452+
raise not_found_exception(name, project)

‎sdk/python/tests/integration/registration/test_sql_registry.py‎

Lines changed: 0 additions & 114 deletions
Original file line numberDiff line numberDiff line change
@@ -141,45 +141,6 @@ def test_apply_entity_success(sql_registry):
141141
sql_registry.teardown()
142142

143143

144-
@pytest.mark.skipif(
145-
sys.platform == "darwin", reason="does not run on mac github actions"
146-
)
147-
@pytest.mark.integration
148-
@pytest.mark.parametrize(
149-
"sql_registry", [lazy_fixture("mysql_registry"), lazy_fixture("pg_registry")],
150-
)
151-
def test_apply_entity_integration(sql_registry):
152-
entity = Entity(
153-
name="driver_car_id", description="Car driver id", tags={"team": "matchmaking"},
154-
)
155-
156-
project = "project"
157-
158-
# Register Entity
159-
sql_registry.apply_entity(entity, project)
160-
161-
entities = sql_registry.list_entities(project)
162-
163-
entity = entities[0]
164-
assert (
165-
len(entities) == 1
166-
and entity.name == "driver_car_id"
167-
and entity.description == "Car driver id"
168-
and "team" in entity.tags
169-
and entity.tags["team"] == "matchmaking"
170-
)
171-
172-
entity = sql_registry.get_entity("driver_car_id", project)
173-
assert (
174-
entity.name == "driver_car_id"
175-
and entity.description == "Car driver id"
176-
and "team" in entity.tags
177-
and entity.tags["team"] == "matchmaking"
178-
)
179-
180-
sql_registry.teardown()
181-
182-
183144
@pytest.mark.skipif(
184145
sys.platform == "darwin", reason="does not run on mac github actions"
185146
)
@@ -449,81 +410,6 @@ def odfv1(feature_df: pd.DataFrame) -> pd.DataFrame:
449410
sql_registry.teardown()
450411

451412

452-
@pytest.mark.skipif(
453-
sys.platform == "darwin", reason="does not run on mac github actions"
454-
)
455-
@pytest.mark.integration
456-
@pytest.mark.parametrize(
457-
"sql_registry", [lazy_fixture("mysql_registry"), lazy_fixture("pg_registry")],
458-
)
459-
def test_apply_feature_view_integration(sql_registry):
460-
# Create Feature Views
461-
batch_source = FileSource(
462-
file_format=ParquetFormat(),
463-
path="file://feast/*",
464-
timestamp_field="ts_col",
465-
created_timestamp_column="timestamp",
466-
)
467-
468-
entity = Entity(name="fs1_my_entity_1", join_keys=["test"])
469-
470-
fv1 = FeatureView(
471-
name="my_feature_view_1",
472-
schema=[
473-
Field(name="fs1_my_feature_1", dtype=Int64),
474-
Field(name="fs1_my_feature_2", dtype=String),
475-
Field(name="fs1_my_feature_3", dtype=Array(String)),
476-
Field(name="fs1_my_feature_4", dtype=Array(Bytes)),
477-
],
478-
entities=[entity],
479-
tags={"team": "matchmaking"},
480-
batch_source=batch_source,
481-
ttl=timedelta(minutes=5),
482-
)
483-
484-
project = "project"
485-
486-
# Register Feature View
487-
sql_registry.apply_feature_view(fv1, project)
488-
489-
feature_views = sql_registry.list_feature_views(project)
490-
491-
# List Feature Views
492-
assert (
493-
len(feature_views) == 1
494-
and feature_views[0].name == "my_feature_view_1"
495-
and feature_views[0].features[0].name == "fs1_my_feature_1"
496-
and feature_views[0].features[0].dtype == Int64
497-
and feature_views[0].features[1].name == "fs1_my_feature_2"
498-
and feature_views[0].features[1].dtype == String
499-
and feature_views[0].features[2].name == "fs1_my_feature_3"
500-
and feature_views[0].features[2].dtype == Array(String)
501-
and feature_views[0].features[3].name == "fs1_my_feature_4"
502-
and feature_views[0].features[3].dtype == Array(Bytes)
503-
and feature_views[0].entities[0] == "fs1_my_entity_1"
504-
)
505-
506-
feature_view = sql_registry.get_feature_view("my_feature_view_1", project)
507-
assert (
508-
feature_view.name == "my_feature_view_1"
509-
and feature_view.features[0].name == "fs1_my_feature_1"
510-
and feature_view.features[0].dtype == Int64
511-
and feature_view.features[1].name == "fs1_my_feature_2"
512-
and feature_view.features[1].dtype == String
513-
and feature_view.features[2].name == "fs1_my_feature_3"
514-
and feature_view.features[2].dtype == Array(String)
515-
and feature_view.features[3].name == "fs1_my_feature_4"
516-
and feature_view.features[3].dtype == Array(Bytes)
517-
and feature_view.entities[0] == "fs1_my_entity_1"
518-
)
519-
520-
sql_registry.delete_feature_view("my_feature_view_1", project)
521-
feature_views = sql_registry.list_feature_views(project)
522-
assert len(feature_views) == 0
523-
524-
sql_registry.teardown()
525-
526-
527413
@pytest.mark.skipif(
528414
sys.platform == "darwin", reason="does not run on mac github actions"
529415
)

0 commit comments

Comments
 (0)