Skip to content

fix: Complete data source form support for all SDK-configurable parameters #6613

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
ntkathole wants to merge 1 commit into
base: master
Choose a base branch
from
Open

fix: Complete data source form support for all SDK-configurable parameters #6613

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
84 changes: 84 additions & 0 deletions sdk/python/feast/api/registry/rest/data_sources.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -31,6 +31,8 @@

class FileOptionsModel(BaseModel):
uri: str = ""
file_format: Optional[str] = "parquet"
s3_endpoint_override: Optional[str] = ""


class BigQueryOptionsModel(BaseModel):
Expand All @@ -42,22 +44,59 @@ class SnowflakeOptionsModel(BaseModel):
table: str = ""
database: str = ""
schema_: str = ""
query: Optional[str] = ""
warehouse: Optional[str] = ""


class RedshiftOptionsModel(BaseModel):
table: str = ""
database: str = ""
schema_: str = ""
query: Optional[str] = ""


class KafkaOptionsModel(BaseModel):
kafka_bootstrap_servers: str = ""
topic: str = ""
message_format: Optional[str] = "json"
watermark_delay_threshold: Optional[str] = ""


class SparkOptionsModel(BaseModel):
table: str = ""
path: str = ""
query: Optional[str] = ""
file_format: Optional[str] = ""
table_format: Optional[str] = ""
table_format_catalog: Optional[str] = ""
table_format_namespace: Optional[str] = ""
table_format_properties: Optional[str] = ""
date_partition_column: Optional[str] = ""
date_partition_column_format: Optional[str] = ""


class KinesisOptionsModel(BaseModel):
region: str = ""
stream_name: str = ""
record_format: Optional[str] = "json"


class TrinoOptionsModel(BaseModel):
table: str = ""
query: str = ""


class AthenaOptionsModel(BaseModel):
table: str = ""
query: str = ""
database: str = ""
data_source: Optional[str] = ""


class CustomOptionsModel(BaseModel):
configuration: Optional[str] = ""
class_name: Optional[str] = ""
config: Optional[str] = ""


class ApplyDataSourceRequestBody(BaseModel):
Expand All @@ -66,6 +105,7 @@ class ApplyDataSourceRequestBody(BaseModel):
type: Optional[int] = None
timestamp_field: Optional[str] = ""
created_timestamp_column: Optional[str] = ""
date_partition_column: Optional[str] = ""
description: Optional[str] = ""
tags: Optional[Dict[str, str]] = {}
owner: Optional[str] = ""
Expand All @@ -75,6 +115,11 @@ class ApplyDataSourceRequestBody(BaseModel):
redshift_options: Optional[RedshiftOptionsModel] = None
kafka_options: Optional[KafkaOptionsModel] = None
spark_options: Optional[SparkOptionsModel] = None
trino_options: Optional[TrinoOptionsModel] = None
athena_options: Optional[AthenaOptionsModel] = None
kinesis_options: Optional[KinesisOptionsModel] = None
custom_options: Optional[CustomOptionsModel] = None
data_source_class_type: Optional[str] = None


def get_data_source_router(grpc_handler) -> APIRouter:
Expand Down Expand Up @@ -221,19 +266,30 @@ def apply_data_source(body: ApplyDataSourceRequestBody):
if body.type is not None:
ds_proto.type = body.type # type: ignore[assignment]

if body.date_partition_column:
ds_proto.date_partition_column = body.date_partition_column

if body.file_options:
ds_proto.file_options.uri = body.file_options.uri
if body.file_options.s3_endpoint_override:
ds_proto.file_options.s3_endpoint_override = (
body.file_options.s3_endpoint_override
)
elif body.bigquery_options:
ds_proto.bigquery_options.table = body.bigquery_options.table
ds_proto.bigquery_options.query = body.bigquery_options.query
elif body.snowflake_options:
ds_proto.snowflake_options.table = body.snowflake_options.table
ds_proto.snowflake_options.database = body.snowflake_options.database
ds_proto.snowflake_options.schema = body.snowflake_options.schema_
if body.snowflake_options.query:
ds_proto.snowflake_options.query = body.snowflake_options.query
elif body.redshift_options:
ds_proto.redshift_options.table = body.redshift_options.table
ds_proto.redshift_options.database = body.redshift_options.database
ds_proto.redshift_options.schema = body.redshift_options.schema_
if body.redshift_options.query:
ds_proto.redshift_options.query = body.redshift_options.query
elif body.kafka_options:
ds_proto.kafka_options.kafka_bootstrap_servers = (
body.kafka_options.kafka_bootstrap_servers
Expand All @@ -242,6 +298,34 @@ def apply_data_source(body: ApplyDataSourceRequestBody):
elif body.spark_options:
ds_proto.spark_options.table = body.spark_options.table
ds_proto.spark_options.path = body.spark_options.path
if body.spark_options.query:
ds_proto.spark_options.query = body.spark_options.query
if body.spark_options.file_format:
ds_proto.spark_options.file_format = body.spark_options.file_format
if body.spark_options.date_partition_column_format:
ds_proto.spark_options.date_partition_column_format = (
body.spark_options.date_partition_column_format
)
elif body.trino_options:
ds_proto.trino_options.table = body.trino_options.table
ds_proto.trino_options.query = body.trino_options.query
elif body.athena_options:
ds_proto.athena_options.table = body.athena_options.table
ds_proto.athena_options.query = body.athena_options.query
ds_proto.athena_options.database = body.athena_options.database
if body.athena_options.data_source:
ds_proto.athena_options.data_source = body.athena_options.data_source
elif body.kinesis_options:
ds_proto.kinesis_options.region = body.kinesis_options.region
ds_proto.kinesis_options.stream_name = body.kinesis_options.stream_name
elif body.custom_options:
if body.custom_options.configuration:
ds_proto.custom_options.configuration = (
body.custom_options.configuration.encode("utf-8")
)

if body.data_source_class_type:
ds_proto.data_source_class_type = body.data_source_class_type

req = RegistryServer_pb2.ApplyDataSourceRequest(
data_source=ds_proto,
Expand Down
119 changes: 117 additions & 2 deletions sdk/python/feast/api/registry/rest/saved_datasets.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,3 +1,4 @@
import json
import logging
from typing import Dict, List, Optional

Expand Down Expand Up @@ -32,13 +33,58 @@
logger = logging.getLogger(__name__)


def _build_feature_view_to_datasource_map(
grpc_handler, project: str, allow_cache: bool = True
) -> Dict[str, str]:
"""Build a map of feature_view_name → data_source_name for a project."""
try:
fv_req = RegistryServer_pb2.ListAllFeatureViewsRequest(
project=project, allow_cache=allow_cache
)
fv_response = grpc_call(grpc_handler.ListAllFeatureViews, fv_req)
feature_views = fv_response.get("featureViews", [])
except Exception:
return {}

fv_ds_map: Dict[str, str] = {}
for any_fv in feature_views:
# AnyFeatureView wraps under "featureView", "onDemandFeatureView", etc.
fv_data = any_fv.get("featureView") or any_fv.get("streamFeatureView") or {}
spec = fv_data.get("spec", {})
fv_name = spec.get("name", "")
if not fv_name:
continue
batch_source = spec.get("batchSource", {})
ds_name = batch_source.get("name", "") if batch_source else ""
if ds_name:
fv_ds_map[fv_name] = ds_name
return fv_ds_map


def _get_dataset_data_sources(dataset: dict, fv_ds_map: Dict[str, str]) -> List[str]:
"""Get data source names for a dataset via its features' feature view lineage."""
spec = dataset.get("spec", dataset)
features = spec.get("features", [])
data_sources: List[str] = []
seen: set = set()
for feature_ref in features:
if isinstance(feature_ref, str) and ":" in feature_ref:
fv_name = feature_ref.split(":")[0]
ds_name = fv_ds_map.get(fv_name)
if ds_name and ds_name not in seen:
seen.add(ds_name)
data_sources.append(ds_name)
return data_sources


class RegisterDatasetRequest(BaseModel):
name: str
project: str
features: List[str] = []
join_keys: List[str] = []
storage_path: Optional[str] = None
storage_type: str = "file"
storage_file_format: Optional[str] = "parquet"
tags: Dict[str, str] = {}
full_feature_names: bool = False
feature_service_name: Optional[str] = None
Expand All @@ -61,6 +107,7 @@ class CreateDatasetRequest(BaseModel):
# Storage
storage_type: str = "file"
storage_path: str = ""
storage_file_format: Optional[str] = "parquet"
tags: Dict[str, str] = {}
allow_overwrite: bool = False

Expand All @@ -79,7 +126,7 @@ def list_saved_datasets_all(
False, description="Include relationships for each saved dataset"
),
):
return aggregate_across_projects(
result = aggregate_across_projects(
grpc_handler=grpc_handler,
list_method=grpc_handler.ListSavedDatasets,
request_cls=RegistryServer_pb2.ListSavedDatasetsRequest,
Expand All @@ -93,6 +140,21 @@ def list_saved_datasets_all(
include_relationships=include_relationships,
)

saved_datasets = result.get("savedDatasets", [])
projects_seen: Dict[str, Dict[str, str]] = {}
for sd in saved_datasets:
proj = sd.get("project", "")
if proj not in projects_seen:
projects_seen[proj] = _build_feature_view_to_datasource_map(
grpc_handler, proj, allow_cache
)
ds_names = _get_dataset_data_sources(sd, projects_seen[proj])
if ds_names:
spec = sd.get("spec", sd)
spec["dataSources"] = ds_names

return result

@router.get("/saved_datasets/data/{name}")
def get_dataset_data(
name: str,
Expand Down Expand Up @@ -193,6 +255,14 @@ def get_saved_dataset(

result = saved_dataset

fv_ds_map = _build_feature_view_to_datasource_map(
grpc_handler, project, allow_cache
)
ds_names = _get_dataset_data_sources(result, fv_ds_map)
if ds_names:
spec = result.get("spec", result)
spec["dataSources"] = ds_names

if include_relationships:
relationships = get_object_relationships(
grpc_handler, "savedDataset", name, project, allow_cache
Expand Down Expand Up @@ -246,6 +316,15 @@ def list_saved_datasets(
response = grpc_call(grpc_handler.ListSavedDatasets, req)
saved_datasets = response.get("savedDatasets", [])

fv_ds_map = _build_feature_view_to_datasource_map(
grpc_handler, project, allow_cache
)
for sd in saved_datasets:
ds_names = _get_dataset_data_sources(sd, fv_ds_map)
if ds_names:
spec = sd.get("spec", sd)
spec["dataSources"] = ds_names

result = {
"savedDatasets": saved_datasets,
"pagination": response.get("pagination", {}),
Expand Down Expand Up @@ -308,12 +387,48 @@ def register_saved_dataset(payload: RegisterDatasetRequest = Body(...)):
)
elif payload.storage_type == "spark":
storage_proto.spark_storage.CopyFrom(
DataSourceProto.SparkOptions(path=path)
DataSourceProto.SparkOptions(
path=path,
file_format=payload.storage_file_format or "parquet",
)
)
elif payload.storage_type == "athena":
storage_proto.athena_storage.CopyFrom(
DataSourceProto.AthenaOptions(table=path)
)
elif payload.storage_type == "postgres":
config = json.dumps({"name": None, "query": None, "table": path})
storage_proto.custom_storage.CopyFrom(
DataSourceProto.CustomSourceOptions(
configuration=config.encode("utf-8")
)
)
elif payload.storage_type == "clickhouse":
config = json.dumps({"name": None, "query": None, "table": path})
storage_proto.custom_storage.CopyFrom(
DataSourceProto.CustomSourceOptions(
configuration=config.encode("utf-8")
)
)
elif payload.storage_type == "couchbase":
parts = path.split(".", 2)
database = parts[0] if len(parts) > 0 else ""
scope = parts[1] if len(parts) > 1 else ""
collection = parts[2] if len(parts) > 2 else ""
config = json.dumps(
{
"name": None,
"query": None,
"database": database,
"scope": scope,
"collection": collection,
}
)
storage_proto.custom_storage.CopyFrom(
DataSourceProto.CustomSourceOptions(
configuration=config.encode("utf-8")
)
)
elif payload.storage_type == "custom":
storage_proto.custom_storage.CopyFrom(
DataSourceProto.CustomSourceOptions(configuration=path.encode("utf-8"))
Expand Down
Loading
Loading