Skip to content

Commit f77bd1d

Browse files
rohit-bharmalcursoragent
rohit-bharmal
and
cursoragent
authored
feat: Adds registry REST API endpoints for managing entities, data sources, and feature views (#6413)
* feat: Implement REST API for entity management with apply and delete endpoints Co-authored-by: Cursor <[email protected]> Signed-off-by: Rohit Bharmal <[email protected]> * feat: removed proto and added REST endpoints Co-authored-by: Cursor <[email protected]> Signed-off-by: Rohit Bharmal <[email protected]> * chore: fix pre-commit lint issues in REST API files Signed-off-by: Rohit Bharmal <[email protected]> Co-authored-by: Cursor <[email protected]> * fix: handle default data source type in apply_data_source function Signed-off-by: Rohit Bharmal <[email protected]> * feat: Add REST API tests for entity, data source, and feature view management Signed-off-by: Rohit Bharmal <[email protected]> * refactor: fixed lint issues Signed-off-by: Rohit Bharmal <[email protected]> * fix: update default values for ApplyDataSourceRequestBody and ApplyFeatureViewRequestBody Signed-off-by: Rohit Bharmal <[email protected]> --------- Signed-off-by: Rohit Bharmal <[email protected]> Co-authored-by: Cursor <[email protected]>
1 parent f77981c commit f77bd1d

4 files changed

Lines changed: 382 additions & 2 deletions

File tree

‎sdk/python/feast/api/registry/rest/data_sources.py‎

Lines changed: 117 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,9 @@
11
import logging
2-
from typing import Dict
2+
from typing import Dict, Optional
33

44
from fastapi import APIRouter, Depends, Query
5+
from fastapi.responses import JSONResponse
6+
from pydantic import BaseModel
57

68
from feast.api.registry.rest.codegen_utils import (
79
render_data_source_code,
@@ -19,13 +21,62 @@
1921
grpc_call,
2022
parse_tags,
2123
)
24+
from feast.protos.feast.core.DataSource_pb2 import DataSource as DataSourceProto
2225
from feast.protos.feast.registry import RegistryServer_pb2
2326
from feast.type_map import _convert_value_type_str_to_value_type
2427
from feast.types import from_value_type
2528

2629
logger = logging.getLogger(__name__)
2730

2831

32+
class FileOptionsModel(BaseModel):
33+
uri: str = ""
34+
35+
36+
class BigQueryOptionsModel(BaseModel):
37+
table: str = ""
38+
query: str = ""
39+
40+
41+
class SnowflakeOptionsModel(BaseModel):
42+
table: str = ""
43+
database: str = ""
44+
schema_: str = ""
45+
46+
47+
class RedshiftOptionsModel(BaseModel):
48+
table: str = ""
49+
database: str = ""
50+
schema_: str = ""
51+
52+
53+
class KafkaOptionsModel(BaseModel):
54+
kafka_bootstrap_servers: str = ""
55+
topic: str = ""
56+
57+
58+
class SparkOptionsModel(BaseModel):
59+
table: str = ""
60+
path: str = ""
61+
62+
63+
class ApplyDataSourceRequestBody(BaseModel):
64+
name: str
65+
project: str
66+
type: Optional[int] = None
67+
timestamp_field: Optional[str] = ""
68+
created_timestamp_column: Optional[str] = ""
69+
description: Optional[str] = ""
70+
tags: Optional[Dict[str, str]] = {}
71+
owner: Optional[str] = ""
72+
file_options: Optional[FileOptionsModel] = None
73+
bigquery_options: Optional[BigQueryOptionsModel] = None
74+
snowflake_options: Optional[SnowflakeOptionsModel] = None
75+
redshift_options: Optional[RedshiftOptionsModel] = None
76+
kafka_options: Optional[KafkaOptionsModel] = None
77+
spark_options: Optional[SparkOptionsModel] = None
78+
79+
2980
def get_data_source_router(grpc_handler) -> APIRouter:
3081
router = APIRouter()
3182

@@ -157,4 +208,69 @@ def get_data_source(
157208
result["featureDefinition"] = render_data_source_code(context)
158209
return result
159210

211+
@router.post("/data_sources", status_code=201)
212+
def apply_data_source(body: ApplyDataSourceRequestBody):
213+
ds_proto = DataSourceProto(
214+
name=body.name,
215+
timestamp_field=body.timestamp_field or "",
216+
created_timestamp_column=body.created_timestamp_column or "",
217+
description=body.description or "",
218+
tags=body.tags or {},
219+
owner=body.owner or "",
220+
)
221+
if body.type is not None:
222+
ds_proto.type = body.type # type: ignore[assignment]
223+
224+
if body.file_options:
225+
ds_proto.file_options.uri = body.file_options.uri
226+
elif body.bigquery_options:
227+
ds_proto.bigquery_options.table = body.bigquery_options.table
228+
ds_proto.bigquery_options.query = body.bigquery_options.query
229+
elif body.snowflake_options:
230+
ds_proto.snowflake_options.table = body.snowflake_options.table
231+
ds_proto.snowflake_options.database = body.snowflake_options.database
232+
ds_proto.snowflake_options.schema = body.snowflake_options.schema_
233+
elif body.redshift_options:
234+
ds_proto.redshift_options.table = body.redshift_options.table
235+
ds_proto.redshift_options.database = body.redshift_options.database
236+
ds_proto.redshift_options.schema = body.redshift_options.schema_
237+
elif body.kafka_options:
238+
ds_proto.kafka_options.kafka_bootstrap_servers = (
239+
body.kafka_options.kafka_bootstrap_servers
240+
)
241+
ds_proto.kafka_options.topic = body.kafka_options.topic
242+
elif body.spark_options:
243+
ds_proto.spark_options.table = body.spark_options.table
244+
ds_proto.spark_options.path = body.spark_options.path
245+
246+
req = RegistryServer_pb2.ApplyDataSourceRequest(
247+
data_source=ds_proto,
248+
project=body.project,
249+
commit=True,
250+
)
251+
grpc_call(grpc_handler.ApplyDataSource, req)
252+
253+
return JSONResponse(
254+
status_code=201,
255+
content={
256+
"name": body.name,
257+
"project": body.project,
258+
"status": "applied",
259+
},
260+
)
261+
262+
@router.delete("/data_sources/{name}")
263+
def delete_data_source(
264+
name: str,
265+
project: str = Query(...),
266+
):
267+
req = RegistryServer_pb2.DeleteDataSourceRequest(
268+
name=name,
269+
project=project,
270+
commit=True,
271+
)
272+
grpc_call(grpc_handler.DeleteDataSource, req)
273+
274+
return {"name": name, "project": project, "status": "deleted"}
275+
160276
return router

‎sdk/python/feast/api/registry/rest/entities.py‎

Lines changed: 55 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,9 @@
11
import logging
2+
from typing import Dict, Optional
23

34
from fastapi import APIRouter, Depends, Query
5+
from fastapi.responses import JSONResponse
6+
from pydantic import BaseModel
47

58
from feast.api.registry.rest.codegen_utils import render_entity_code
69
from feast.api.registry.rest.rest_utils import (
@@ -13,11 +16,23 @@
1316
get_sorting_params,
1417
grpc_call,
1518
)
19+
from feast.protos.feast.core.Entity_pb2 import Entity as EntityProto
20+
from feast.protos.feast.core.Entity_pb2 import EntitySpecV2 as EntitySpecProto
1621
from feast.protos.feast.registry import RegistryServer_pb2
1722

1823
logger = logging.getLogger(__name__)
1924

2025

26+
class ApplyEntityRequestBody(BaseModel):
27+
name: str
28+
project: str
29+
join_key: Optional[str] = None
30+
value_type: Optional[int] = 2
31+
description: Optional[str] = ""
32+
tags: Optional[Dict[str, str]] = {}
33+
owner: Optional[str] = ""
34+
35+
2136
def get_entity_router(grpc_handler) -> APIRouter:
2237
router = APIRouter()
2338

@@ -136,4 +151,44 @@ def get_entity(
136151
result["featureDefinition"] = render_entity_code(context)
137152
return result
138153

154+
@router.post("/entities", status_code=201)
155+
def apply_entity(body: ApplyEntityRequestBody):
156+
join_key = body.join_key if body.join_key else body.name
157+
158+
spec = EntitySpecProto(
159+
name=body.name,
160+
value_type=body.value_type,
161+
join_key=join_key,
162+
description=body.description or "",
163+
tags=body.tags or {},
164+
owner=body.owner or "",
165+
)
166+
entity_proto = EntityProto(spec=spec)
167+
168+
req = RegistryServer_pb2.ApplyEntityRequest(
169+
entity=entity_proto,
170+
project=body.project,
171+
commit=True,
172+
)
173+
grpc_call(grpc_handler.ApplyEntity, req)
174+
175+
return JSONResponse(
176+
status_code=201,
177+
content={"name": body.name, "project": body.project, "status": "applied"},
178+
)
179+
180+
@router.delete("/entities/{name}")
181+
def delete_entity(
182+
name: str,
183+
project: str = Query(...),
184+
):
185+
req = RegistryServer_pb2.DeleteEntityRequest(
186+
name=name,
187+
project=project,
188+
commit=True,
189+
)
190+
grpc_call(grpc_handler.DeleteEntity, req)
191+
192+
return {"name": name, "project": project, "status": "deleted"}
193+
139194
return router

‎sdk/python/feast/api/registry/rest/feature_views.py‎

Lines changed: 89 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,10 @@
1-
from typing import Dict
1+
import logging
2+
from typing import Dict, List, Optional
23

34
from fastapi import APIRouter, Depends, Query
5+
from fastapi.responses import JSONResponse
6+
from google.protobuf.duration_pb2 import Duration
7+
from pydantic import BaseModel
48

59
from feast.api.registry.rest.codegen_utils import render_feature_view_code
610
from feast.api.registry.rest.rest_utils import (
@@ -14,10 +18,34 @@
1418
paginate_and_sort,
1519
parse_tags,
1620
)
21+
from feast.protos.feast.core.DataSource_pb2 import DataSource as DataSourceProto
22+
from feast.protos.feast.core.Feature_pb2 import FeatureSpecV2
23+
from feast.protos.feast.core.FeatureView_pb2 import FeatureView as FeatureViewProto
24+
from feast.protos.feast.core.FeatureView_pb2 import FeatureViewSpec
1725
from feast.registry_server import RegistryServer_pb2
1826
from feast.type_map import _convert_value_type_str_to_value_type
1927
from feast.types import from_value_type
2028

29+
logger = logging.getLogger(__name__)
30+
31+
32+
class FeatureModel(BaseModel):
33+
name: str
34+
value_type: int = 2
35+
36+
37+
class ApplyFeatureViewRequestBody(BaseModel):
38+
name: str
39+
project: str
40+
entities: Optional[List[str]] = []
41+
features: Optional[List[FeatureModel]] = []
42+
batch_source: Optional[str] = ""
43+
ttl_seconds: Optional[int] = None
44+
online: Optional[bool] = True
45+
description: Optional[str] = ""
46+
tags: Optional[Dict[str, str]] = {}
47+
owner: Optional[str] = ""
48+
2149

2250
def _extract_feature_view_from_any(any_feature_view: dict) -> dict:
2351
"""Extract the specific feature view type and data from an AnyFeatureView object.
@@ -275,4 +303,64 @@ def list_all_feature_views(
275303

276304
return result
277305

306+
@router.post("/feature_views", status_code=201)
307+
def apply_feature_view(body: ApplyFeatureViewRequestBody):
308+
feature_specs = []
309+
for f in body.features or []:
310+
feature_specs.append(FeatureSpecV2(name=f.name, value_type=f.value_type))
311+
312+
batch_source_proto = (
313+
DataSourceProto(name=body.batch_source) if body.batch_source else None
314+
)
315+
316+
ttl = (
317+
Duration(seconds=body.ttl_seconds) if body.ttl_seconds is not None else None
318+
)
319+
320+
spec = FeatureViewSpec(
321+
name=body.name,
322+
entities=body.entities or [],
323+
features=feature_specs,
324+
tags=body.tags or {},
325+
online=body.online if body.online is not None else True,
326+
description=body.description or "",
327+
owner=body.owner or "",
328+
)
329+
if ttl is not None:
330+
spec.ttl.CopyFrom(ttl)
331+
if batch_source_proto:
332+
spec.batch_source.CopyFrom(batch_source_proto)
333+
334+
fv_proto = FeatureViewProto(spec=spec)
335+
336+
req = RegistryServer_pb2.ApplyFeatureViewRequest(
337+
feature_view=fv_proto,
338+
project=body.project,
339+
commit=True,
340+
)
341+
grpc_call(grpc_handler.ApplyFeatureView, req)
342+
343+
return JSONResponse(
344+
status_code=201,
345+
content={
346+
"name": body.name,
347+
"project": body.project,
348+
"status": "applied",
349+
},
350+
)
351+
352+
@router.delete("/feature_views/{name}")
353+
def delete_feature_view(
354+
name: str,
355+
project: str = Query(...),
356+
):
357+
req = RegistryServer_pb2.DeleteFeatureViewRequest(
358+
name=name,
359+
project=project,
360+
commit=True,
361+
)
362+
grpc_call(grpc_handler.DeleteFeatureView, req)
363+
364+
return {"name": name, "project": project, "status": "deleted"}
365+
278366
return router

0 commit comments

Comments
 (0)