feat: Enhancements to Feast OpenLineage UI

Signed-off-by: ntkathole <nikhilkathole2683@gmail.com>
ntkathole · Aug 18, 2026 · Aug 5, 2026 · Aug 9, 2026 · Aug 10, 2026 · Aug 18, 2026
commit d80d3cc0d941f79088875aaa87db9ba6409b6ab0
@@ -538,6 +538,7 @@ def _submit_and_process_materialization_jobs(
         regular_fvs: list,
         previous_states: dict,
         date_range: "_MaterializationDateRange",
+        openlineage_run_id: Optional[str] = None,
     ) -> None:
         """
         Submit all tasks to the engine in one call and process the results.
@@ -550,8 +551,22 @@ def _submit_and_process_materialization_jobs(
         )
 
         batch_start = time.monotonic()
+        materialize_kwargs: Dict[str, Any] = {}
+        if openlineage_run_id and self.openlineage_emitter is not None:
+            from feast.openlineage.identity import (
+                LineageParentContext,
+                materialize_job_name,
+            )
+
+            materialize_kwargs["lineage_parent"] = LineageParentContext(
+                job_namespace=self.openlineage_emitter.namespace_for(self.project),
+                job_name=materialize_job_name(self.project),
+                run_id=openlineage_run_id,
+            )
         try:
-            jobs = provider.batch_engine.materialize(self.registry, tasks)
+            jobs = provider.batch_engine.materialize(
+                self.registry, tasks, **materialize_kwargs
+            )
         except Exception:
             self._rollback_fv_states(regular_fvs, previous_states)
             raise
@@ -608,6 +623,7 @@ def _materialize_fvs_batch(
         end_date: datetime,
         tqdm_builder,
         disable_event_timestamp: bool = False,
+        openlineage_run_id: Optional[str] = None,
     ) -> None:
         """Batch path: collect all FVs, submit to engine in one call.
 
@@ -644,6 +660,7 @@ def _materialize_fvs_batch(
                 regular_fvs,
                 previous_states,
                 date_range,
+                openlineage_run_id=openlineage_run_id,
             )
 
     @property
@@ -1921,21 +1938,9 @@ def teardown(self):
     def _teardown_openlineage(self):
         """Clean up OpenLineage data for this project's namespace during teardown."""
         try:
-            if (
-                hasattr(self.config, "openlineage")
-                and self.config.openlineage is not None
-                and self.config.openlineage.enabled
-            ):
-                ol_config = self.config.openlineage.to_openlineage_config()
-                consumer_cfg = getattr(ol_config, "consumer", None)
-                if consumer_cfg and getattr(consumer_cfg, "enabled", False):
-                    conn_str = getattr(consumer_cfg, "connection_string", None)
-                    if conn_str:
-                        from feast.openlineage.store import OpenLineageStore
-
-                        ol_store = OpenLineageStore(connection_string=conn_str)
-                        namespace = f"{self.project}/{self.project}"
-                        ol_store.purge_namespace(namespace)
+            emitter = self.openlineage_emitter
+            if emitter is not None:
+                emitter.teardown_project(self.project)
         except Exception as e:
             warnings.warn(f"Failed to clean up OpenLineage data during teardown: {e}")
 
@@ -2696,6 +2701,7 @@ def tqdm_builder(length):
                     regular_fvs_with_dates,
                     end_date_tz,
                     tqdm_builder,
+                    openlineage_run_id=ol_run_id,
                 )
             else:
                 for feature_view, start_date in regular_fvs_with_dates:
@@ -2909,6 +2915,7 @@ def tqdm_builder(length):
                     end_date,
                     tqdm_builder,
                     disable_event_timestamp=disable_event_timestamp,
+                    openlineage_run_id=ol_run_id,
                 )
             else:
                 for feature_view, fv_start in regular_fvs_with_dates:
@@ -3033,7 +3040,11 @@ def _emit_openlineage_materialize_start(
             return None
         try:
             run_id, success = self.openlineage_emitter.emit_materialize_start(
-                feature_views, start_date, end_date, self.project
+                feature_views,
+                start_date,
+                end_date,
+                self.project,
+                online_store=getattr(self.config, "online_store", None),
             )
             # Return run_id only if START was successfully emitted
             # This prevents orphaned COMPLETE/FAIL events
@@ -3052,7 +3063,10 @@ def _emit_openlineage_materialize_complete(
             return
         try:
             self.openlineage_emitter.emit_materialize_complete(
-                run_id, feature_views, self.project
+                run_id,
+                feature_views,
+                self.project,
+                online_store=getattr(self.config, "online_store", None),
             )
         except Exception as e:
             warnings.warn(f"Failed to emit OpenLineage materialize complete event: {e}")

@@ -1,5 +1,5 @@
 from abc import ABC, abstractmethod
-from typing import List, Sequence, Union
+from typing import TYPE_CHECKING, List, Optional, Sequence, Union
 
 import pyarrow as pa
 
@@ -19,6 +19,9 @@
 from feast.on_demand_feature_view import OnDemandFeatureView
 from feast.stream_feature_view import StreamFeatureView
 
+if TYPE_CHECKING:
+    from feast.openlineage.identity import LineageParentContext
+
 
 class ComputeEngine(ABC):
     """
@@ -104,10 +107,24 @@ def materialize(
         self,
         registry: BaseRegistry,
         tasks: Union[MaterializationTask, List[MaterializationTask]],
+        *,
+        lineage_parent: Optional["LineageParentContext"] = None,
         **kwargs,
     ) -> List[MaterializationJob]:
+        """Materialize features for the given tasks.
+
+        Args:
+            registry: Feature registry
+            tasks: One or more materialization tasks
+            lineage_parent: Optional OpenLineage parent run. Engines that emit
+                their own lineage (e.g. SparkApplication) should attach this as
+                a parentRun; others ignore it.
+            **kwargs: Engine-specific options
+        """
         if isinstance(tasks, MaterializationTask):
             tasks = [tasks]
+        if lineage_parent is not None:
+            kwargs["lineage_parent"] = lineage_parent
         return [self._materialize_one(registry, task, **kwargs) for task in tasks]
 
     def _materialize_one(

@@ -1,7 +1,7 @@
 import logging
 import time
 import uuid
-from typing import List, Optional, Sequence, Union
+from typing import Any, Dict, List, Optional, Sequence, Union
 
 import pyarrow as pa
 import yaml
@@ -173,6 +173,11 @@ def materialize(
             tasks = [tasks]
 
         job_id = uuid.uuid4().hex[:8]
+        from feast.openlineage.identity import coerce_lineage_parent
+
+        lineage_parent = coerce_lineage_parent(
+            kwargs.get("lineage_parent") or kwargs.get("openlineage_parent")
+        )
 
         try:
             self._create_with_retry(
@@ -193,7 +198,7 @@ def materialize(
             return [job for _ in tasks]
 
         try:
-            cr = self._build_spark_application_cr(job_id)
+            cr = self._build_spark_application_cr(job_id, lineage_parent=lineage_parent)
             self._create_with_retry(
                 lambda: self.custom_api.create_namespaced_custom_object(
                     group="sparkoperator.k8s.io",
@@ -223,7 +228,14 @@ def materialize(
         )
         try:
             self._wait_for_completion(job)
-            return self._build_per_fv_jobs(registry, tasks, job_id)
+            # Freeze success before finally-cleanup deletes the CR; otherwise
+            # FeatureStore's later job.status() poll gets 404 → false FAIL.
+            if (
+                job.error() is None
+                and job.status() == MaterializationJobStatus.SUCCEEDED
+            ):
+                job.mark_succeeded()
+            return self._build_per_fv_jobs(registry, tasks, job_id, job)
         finally:
             self._cleanup(job_id)
 
@@ -252,6 +264,7 @@ def _build_per_fv_jobs(
         registry: BaseRegistry,
         tasks: List[MaterializationTask],
         job_id: str,
+        job: SparkApplicationMaterializationJob,
     ) -> List[MaterializationJob]:
         """Build one independent job object per FV from registry state.
 
@@ -270,6 +283,14 @@ def _build_per_fv_jobs(
         that actually succeeded (#6673). ``CompletedMaterializationJob`` reports
         SUCCEEDED without any Kubernetes call, so it is safe after cleanup.
         """
+        if len(tasks) <= 1:
+            if (
+                job.error() is None
+                and job.status() == MaterializationJobStatus.SUCCEEDED
+            ):
+                return [CompletedMaterializationJob(job_id)]
+            return [job]
+
         jobs: List[MaterializationJob] = []
         for task in tasks:
             fv = registry.get_feature_view(task.feature_view.name, task.project)
@@ -324,7 +345,9 @@ def _create_configmap(self, job_id: str, tasks: List[MaterializationTask]):
             namespace=self.config.namespace, body=manifest
         )
 
-    def _build_spark_application_cr(self, job_id: str) -> dict:
+    def _build_spark_application_cr(
+        self, job_id: str, lineage_parent: Optional[Any] = None
+    ) -> dict:
         driver_env_conf = {
             "spark.kubernetes.driverEnv.FEAST_CONFIGMAP_NAME": f"feast-sa-{job_id}",
             "spark.kubernetes.driverEnv.FEAST_CONFIGMAP_NAMESPACE": self.config.namespace,
@@ -336,6 +359,22 @@ def _build_spark_application_cr(self, job_id: str) -> dict:
                     entry["value"]
                 )
 
+        from feast.openlineage.identity import (
+            coerce_lineage_parent,
+            spark_compute_job_name,
+        )
+
+        # K8s names stay unique per run; OL job identity is stable.
+        k8s_app_name = f"feast-sa-{job_id}"
+        project = getattr(self.repo_config, "project", None) or "feast"
+        parent_conf: Dict[str, str] = {
+            "spark.app.name": k8s_app_name,
+            "spark.openlineage.appName": spark_compute_job_name(project),
+        }
+        parent = coerce_lineage_parent(lineage_parent)
+        if parent:
+            parent_conf.update(parent.to_spark_openlineage_conf())
+
         spec = {
             "type": "Python",
             "mode": "cluster",
@@ -347,6 +386,7 @@ def _build_spark_application_cr(self, job_id: str) -> dict:
             "sparkConf": {
                 "spark.scheduler.mode": "FAIR",
                 **(self.config.spark_conf or {}),
+                **parent_conf,
                 **driver_env_conf,
             },
             "restartPolicy": {

@@ -90,8 +90,16 @@ def __init__(
         self.namespace = namespace
         self.custom_api = custom_api
         self._error: Optional[BaseException] = error
+        # Once we observe SUCCEEDED/ERROR on the CR, keep it. Feast (and the
+        # Spark Operator TTL) delete the CR after success; a later status()
+        # poll must not turn that into a false "not found" failure.
+        self._terminal_status: Optional[MaterializationJobStatus] = None
+        if error is not None:
+            self._terminal_status = MaterializationJobStatus.ERROR
 
     def status(self) -> MaterializationJobStatus:
+        if self._terminal_status is not None:
+            return self._terminal_status
         if self._error is not None:
             return MaterializationJobStatus.ERROR
 
@@ -112,8 +120,16 @@ def status(self) -> MaterializationJobStatus:
                 .get("errorMessage", f"SparkApplication failed: {state}")
             )
             self._error = Exception(msg)
+            self._terminal_status = MaterializationJobStatus.ERROR
+        elif result == MaterializationJobStatus.SUCCEEDED:
+            self._terminal_status = MaterializationJobStatus.SUCCEEDED
         return result
 
+    def mark_succeeded(self) -> None:
+        """Record successful completion before the CR is deleted/TTL'd."""
+        self._error = None
+        self._terminal_status = MaterializationJobStatus.SUCCEEDED
+
     def _get_cr_with_retry(self) -> Optional[dict]:
         """Fetch SparkApplication CR with exponential backoff on transient errors."""
         last_exc = None
@@ -128,6 +144,9 @@ def _get_cr_with_retry(self) -> Optional[dict]:
                 )
             except ApiException as e:
                 if e.status == 404:
+                    # Already completed jobs are cleaned up; don't invent a failure.
+                    if self._terminal_status == MaterializationJobStatus.SUCCEEDED:
+                        return None
                     self._error = Exception(
                         f"SparkApplication feast-sa-{self._job_id} not found"
                     )

@@ -66,21 +66,38 @@
     FeastEntityFacet,
     FeastFeatureServiceFacet,
     FeastFeatureViewFacet,
+    FeastJobKindFacet,
     FeastMaterializationFacet,
+    FeastOnlineStoreFacet,
     FeastProjectFacet,
 )
+from feast.openlineage.identity import (
+    FeastJobKind,
+    LineageParentContext,
+    materialize_job_name,
+    resolve_namespace,
+    spark_compute_job_name,
+)
 
 __all__ = [
     # Main classes (used internally by native integration)
     "FeastOpenLineageClient",
     "FeastOpenLineageEmitter",
     "OpenLineageConfig",
     "OpenLineageConsumerConfig",
+    # Identity / context
+    "FeastJobKind",
+    "LineageParentContext",
+    "materialize_job_name",
+    "resolve_namespace",
+    "spark_compute_job_name",
     # Facets (custom Feast metadata in lineage events)
     "FeastFeatureViewFacet",
     "FeastFeatureServiceFacet",
     "FeastDataSourceFacet",
     "FeastEntityFacet",
+    "FeastOnlineStoreFacet",
     "FeastMaterializationFacet",
     "FeastProjectFacet",
+    "FeastJobKindFacet",
 ]