Skip to content

Commit ff298df

Browse files
ntkatholentkathole
committed
feat: Add saved datasets to lineage
Signed-off-by: ntkathole <nikhilkathole2683@gmail.com>
1 parent c42db95 commit ff298df

17 files changed

Lines changed: 944 additions & 39 deletions

File tree

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

Lines changed: 58 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -25,6 +25,11 @@
2525
_ol_processor: Optional[Any] = None
2626

2727

28+
def get_ol_processor() -> Optional[Any]:
29+
"""Return the global OpenLineage processor, if initialized."""
30+
return _ol_processor
31+
32+
2833
def register_all_routes(app: FastAPI, grpc_handler, server=None, store=None):
2934
app.include_router(get_entity_router(grpc_handler))
3035
app.include_router(get_data_source_router(grpc_handler))
@@ -133,19 +138,67 @@ def _register_openlineage_consumer(app: FastAPI, feast_store):
133138

134139
# Wire the local processor into Feast's own OL emitter so Feast events
135140
# are also stored in the consumer DB automatically.
141+
# The emitter is lazy-initialized, so it may be None at startup.
142+
# Two-pronged approach:
143+
# 1. If already initialized, wire now.
144+
# 2. Store processor globally so _init_openlineage_emitter() can
145+
# pick it up when the emitter is lazily created later.
136146
try:
137-
if feast_store and hasattr(feast_store, "_openlineage_emitter"):
138-
emitter = feast_store._openlineage_emitter
139-
if emitter and hasattr(emitter, "_client") and emitter._client:
140-
emitter._client.set_local_processor(processor)
141-
logger.info("Feast OL emitter wired to local consumer processor")
147+
emitter = getattr(feast_store, "_openlineage_emitter", None)
148+
if emitter and hasattr(emitter, "_client") and emitter._client:
149+
emitter._client.set_local_processor(processor)
150+
logger.info(
151+
"Feast OL emitter wired to local consumer processor (eager)"
152+
)
142153
except Exception as wire_err:
143154
logger.debug(f"Could not wire emitter to local processor: {wire_err}")
144155

156+
def _build_allowed_namespaces_fn(fs):
157+
"""Build a callback that derives OL namespace access from Feast RBAC.
158+
159+
Uses ``permitted_resources`` with ``DESCRIBE`` on all known projects.
160+
The projects the current user may describe become the allowed OL
161+
namespaces (mapped through ``resolve_namespace``).
162+
"""
163+
164+
def _get_allowed():
165+
try:
166+
from feast.openlineage.identity import resolve_namespace
167+
from feast.permissions.action import AuthzedAction
168+
from feast.permissions.security_manager import (
169+
get_security_manager,
170+
permitted_resources,
171+
)
172+
173+
sm = get_security_manager()
174+
if sm is None:
175+
return None
176+
177+
all_projects = fs.registry.list_projects(allow_cache=True)
178+
if not all_projects:
179+
return None
180+
181+
allowed_projects = permitted_resources(
182+
all_projects, AuthzedAction.DESCRIBE
183+
)
184+
185+
ol_ns_config = getattr(ol_config, "namespace", "feast")
186+
namespaces = set()
187+
for p in allowed_projects:
188+
namespaces.add(resolve_namespace(ol_ns_config, p.name))
189+
return list(namespaces) if namespaces else None
190+
except Exception:
191+
return None
192+
193+
return _get_allowed
194+
195+
get_allowed_namespaces = _build_allowed_namespaces_fn(feast_store)
196+
145197
consumer_router = get_consumer_router(
146198
config=ol_config,
147199
store=ol_store,
148200
processor=processor,
201+
get_allowed_namespaces=get_allowed_namespaces,
149202
)
150203

151204
app.include_router(consumer_router)

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

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -86,6 +86,7 @@ def get_object_relationships_path(
8686
"featureView",
8787
"featureService",
8888
"feature",
89+
"savedDataset",
8990
]
9091
if object_type not in valid_types:
9192
raise ValueError(
@@ -175,6 +176,7 @@ def get_complete_registry_data(
175176
"featureServices": project_resources.get("featureServices", []),
176177
"features": project_resources.get("features", []),
177178
"labels": project_resources.get("labels", []),
179+
"savedDatasets": project_resources.get("savedDatasets", []),
178180
},
179181
"relationships": lineage_response.get("relationships", []),
180182
"indirectRelationships": lineage_response.get("indirectRelationships", []),
@@ -186,6 +188,7 @@ def get_complete_registry_data(
186188
"featureServices": pagination.get("featureServices", {}),
187189
"features": pagination.get("features", {}),
188190
"labels": pagination.get("labels", {}),
191+
"savedDatasets": pagination.get("savedDatasets", {}),
189192
"relationships": lineage_response.get("relationshipsPagination", {}),
190193
"indirectRelationships": lineage_response.get(
191194
"indirectRelationshipsPagination", {}
@@ -274,6 +277,8 @@ def get_complete_registry_data_all(
274277
feat["project"] = project_name
275278
for lbl in project_resources.get("labels", []):
276279
lbl["project"] = project_name
280+
for sd in project_resources.get("savedDatasets", []):
281+
sd["project"] = project_name
277282
all_data.append(
278283
{
279284
"project": project_name,
@@ -285,6 +290,7 @@ def get_complete_registry_data_all(
285290
"featureServices": project_resources.get("featureServices", []),
286291
"features": project_resources.get("features", []),
287292
"labels": project_resources.get("labels", []),
293+
"savedDatasets": project_resources.get("savedDatasets", []),
288294
},
289295
"relationships": lineage_response.get("relationships", []),
290296
"indirectRelationships": lineage_response.get(

‎sdk/python/feast/feature_store.py‎

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -376,6 +376,7 @@ def _init_openlineage_emitter(self) -> Optional[Any]:
376376
ol_config = self.config.openlineage.to_openlineage_config()
377377
emitter = FeastOpenLineageEmitter(ol_config)
378378
if emitter.is_enabled:
379+
self._wire_local_processor(emitter)
379380
return emitter
380381
except ImportError:
381382
# OpenLineage not installed, silently skip
@@ -384,6 +385,21 @@ def _init_openlineage_emitter(self) -> Optional[Any]:
384385
warnings.warn(f"Failed to initialize OpenLineage emitter: {e}")
385386
return None
386387

388+
def _wire_local_processor(self, emitter: Any) -> None:
389+
"""Wire the local OL consumer processor into the emitter so
390+
Feast-produced events are also stored in the consumer DB."""
391+
try:
392+
from feast.api.registry.rest import get_ol_processor
393+
394+
processor = get_ol_processor()
395+
if processor and hasattr(emitter, "_client") and emitter._client:
396+
emitter._client.set_local_processor(processor)
397+
_logger.info(
398+
"Feast OL emitter wired to local consumer processor (lazy)"
399+
)
400+
except Exception as e:
401+
_logger.debug(f"Could not wire emitter to local processor: {e}")
402+
387403
def __repr__(self) -> str:
388404
# Show lazy loading status without triggering initialization
389405
registry_status = "not loaded" if self._registry is None else "loaded"

‎sdk/python/feast/lineage/registry_lineage.py‎

Lines changed: 166 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -2,24 +2,110 @@
22
Registry lineage generation for Feast objects.
33
44
This module provides functionality to generate relationship graphs between
5-
Feast objects (entities, feature views, data sources, feature services)
6-
for lineage visualization.
5+
Feast objects (entities, feature views, data sources, feature services,
6+
saved datasets) for lineage visualization.
77
"""
88

99
from dataclasses import dataclass
1010
from enum import Enum
11-
from typing import Dict, List, Tuple
11+
from typing import Dict, List, Set, Tuple
1212

1313
from feast.protos.feast.core.Registry_pb2 import Registry
1414

1515

16+
def _extract_storage_identifiers(storage) -> Set[str]:
17+
"""Extract physical location identifiers from a SavedDatasetStorage proto.
18+
19+
Returns a set of non-empty strings (URIs, table names, paths) that can
20+
be matched against DataSource options.
21+
"""
22+
ids: Set[str] = set()
23+
if hasattr(storage, "file_storage") and storage.HasField("file_storage"):
24+
if storage.file_storage.uri:
25+
ids.add(storage.file_storage.uri)
26+
if hasattr(storage, "bigquery_storage") and storage.HasField("bigquery_storage"):
27+
if storage.bigquery_storage.table:
28+
ids.add(storage.bigquery_storage.table)
29+
if hasattr(storage, "redshift_storage") and storage.HasField("redshift_storage"):
30+
if storage.redshift_storage.table:
31+
ids.add(storage.redshift_storage.table)
32+
if hasattr(storage, "snowflake_storage") and storage.HasField("snowflake_storage"):
33+
if storage.snowflake_storage.table:
34+
ids.add(storage.snowflake_storage.table)
35+
if hasattr(storage, "spark_storage") and storage.HasField("spark_storage"):
36+
if storage.spark_storage.path:
37+
ids.add(storage.spark_storage.path)
38+
if storage.spark_storage.table:
39+
ids.add(storage.spark_storage.table)
40+
if hasattr(storage, "trino_storage") and storage.HasField("trino_storage"):
41+
if storage.trino_storage.table:
42+
ids.add(storage.trino_storage.table)
43+
if hasattr(storage, "athena_storage") and storage.HasField("athena_storage"):
44+
if storage.athena_storage.table:
45+
ids.add(storage.athena_storage.table)
46+
return ids
47+
48+
49+
def _extract_datasource_identifiers(data_source) -> Set[str]:
50+
"""Extract physical location identifiers from a DataSource proto.
51+
52+
Returns a set of non-empty strings (URIs, table names, paths) that can
53+
be compared against SavedDatasetStorage identifiers.
54+
"""
55+
ids: Set[str] = set()
56+
opts = (
57+
data_source.WhichOneof("options")
58+
if hasattr(data_source, "WhichOneof")
59+
else None
60+
)
61+
if opts == "file_options" and data_source.file_options.uri:
62+
ids.add(data_source.file_options.uri)
63+
elif opts == "bigquery_options" and data_source.bigquery_options.table:
64+
ids.add(data_source.bigquery_options.table)
65+
elif opts == "redshift_options" and data_source.redshift_options.table:
66+
ids.add(data_source.redshift_options.table)
67+
elif opts == "snowflake_options" and data_source.snowflake_options.table:
68+
ids.add(data_source.snowflake_options.table)
69+
elif opts == "spark_options":
70+
if data_source.spark_options.path:
71+
ids.add(data_source.spark_options.path)
72+
if data_source.spark_options.table:
73+
ids.add(data_source.spark_options.table)
74+
elif opts == "trino_options" and data_source.trino_options.table:
75+
ids.add(data_source.trino_options.table)
76+
elif opts == "athena_options" and data_source.athena_options.table:
77+
ids.add(data_source.athena_options.table)
78+
79+
# Also check batch_source if present (FeatureView's embedded source)
80+
if hasattr(data_source, "batch_source") and data_source.HasField("batch_source"):
81+
ids.update(_extract_datasource_identifiers(data_source.batch_source))
82+
83+
return ids
84+
85+
86+
def _build_datasource_location_index(registry: Registry) -> Dict[str, str]:
87+
"""Build a reverse index: physical location → DataSource name.
88+
89+
Scans all DataSources in the registry and maps each physical identifier
90+
(URI, table, path) to the DataSource's name.
91+
"""
92+
location_to_name: Dict[str, str] = {}
93+
for ds in registry.data_sources:
94+
if not (hasattr(ds, "name") and ds.name):
95+
continue
96+
for loc_id in _extract_datasource_identifiers(ds):
97+
location_to_name[loc_id] = ds.name
98+
return location_to_name
99+
100+
16101
class FeastObjectType(Enum):
17102
DATA_SOURCE = "dataSource"
18103
ENTITY = "entity"
19104
FEATURE_VIEW = "featureView"
20105
LABEL_VIEW = "labelView"
21106
FEATURE_SERVICE = "featureService"
22107
FEATURE = "feature"
108+
SAVED_DATASET = "savedDataset"
23109

24110

25111
@dataclass
@@ -390,6 +476,83 @@ def _parse_direct_relationships(self, registry: Registry) -> List[EntityRelation
390476
)
391477
)
392478

479+
# SavedDataset relationships
480+
ds_location_index = _build_datasource_location_index(registry)
481+
482+
for saved_dataset in registry.saved_datasets:
483+
if hasattr(saved_dataset, "spec") and saved_dataset.spec:
484+
# FeatureService -> SavedDataset (when created via a feature service)
485+
if (
486+
hasattr(saved_dataset.spec, "feature_service_name")
487+
and saved_dataset.spec.feature_service_name
488+
):
489+
relationships.append(
490+
EntityRelation(
491+
source=EntityReference(
492+
FeastObjectType.FEATURE_SERVICE,
493+
saved_dataset.spec.feature_service_name,
494+
),
495+
target=EntityReference(
496+
FeastObjectType.SAVED_DATASET,
497+
saved_dataset.spec.name,
498+
),
499+
)
500+
)
501+
502+
# FeatureView -> SavedDataset (derived from feature refs "view:feat")
503+
if (
504+
hasattr(saved_dataset.spec, "features")
505+
and saved_dataset.spec.features
506+
):
507+
from feast.utils import _parse_feature_ref
508+
509+
seen_views: set = set()
510+
for feat_ref in saved_dataset.spec.features:
511+
try:
512+
view_name, _, _ = _parse_feature_ref(feat_ref)
513+
except ValueError:
514+
continue
515+
if view_name and view_name not in seen_views:
516+
seen_views.add(view_name)
517+
relationships.append(
518+
EntityRelation(
519+
source=EntityReference(
520+
FeastObjectType.FEATURE_VIEW,
521+
view_name,
522+
),
523+
target=EntityReference(
524+
FeastObjectType.SAVED_DATASET,
525+
saved_dataset.spec.name,
526+
),
527+
)
528+
)
529+
530+
# DataSource -> SavedDataset (matched via storage location)
531+
if (
532+
hasattr(saved_dataset.spec, "storage")
533+
and saved_dataset.spec.storage
534+
):
535+
storage_ids = _extract_storage_identifiers(
536+
saved_dataset.spec.storage
537+
)
538+
matched_ds_names: set = set()
539+
for loc_id in storage_ids:
540+
ds_name = ds_location_index.get(loc_id)
541+
if ds_name and ds_name not in matched_ds_names:
542+
matched_ds_names.add(ds_name)
543+
relationships.append(
544+
EntityRelation(
545+
source=EntityReference(
546+
FeastObjectType.DATA_SOURCE,
547+
ds_name,
548+
),
549+
target=EntityReference(
550+
FeastObjectType.SAVED_DATASET,
551+
saved_dataset.spec.name,
552+
),
553+
)
554+
)
555+
393556
return relationships
394557

395558
def _parse_indirect_relationships(

‎sdk/python/feast/openlineage/client.py‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -275,6 +275,7 @@ def emit_job_event(
275275
inputs: Optional[List[Any]] = None,
276276
outputs: Optional[List[Any]] = None,
277277
job_facets: Optional[Dict[str, Any]] = None,
278+
namespace: Optional[str] = None,
278279
) -> bool:
279280
"""
280281
Emit a JobEvent for a Feast job definition.
@@ -284,6 +285,7 @@ def emit_job_event(
284285
inputs: List of input datasets
285286
outputs: List of output datasets
286287
job_facets: Job facets
288+
namespace: Optional namespace for the job (defaults to client namespace)
287289
288290
Returns:
289291
True if successful, False otherwise
@@ -297,7 +299,7 @@ def emit_job_event(
297299
event = JobEvent(
298300
eventTime=datetime.now(timezone.utc).isoformat(),
299301
job=Job(
300-
namespace=self.namespace,
302+
namespace=namespace or self.namespace,
301303
name=job_name,
302304
facets=job_facets or {},
303305
),

0 commit comments

Comments
 (0)