Fixes #21886: allow all success openlineage event types (#24889)

Author: jsingh-yelp

ingestion/src/metadata/ingestion/source/pipeline/openlineage/metadata.py

@@ -72,7 +72,7 @@ class OpenlineageSource(PipelineServiceSource):
     Works under the assumption that OpenLineage integrations produce events to Kafka topic, which is a source of events
     for this connector.
 
-    Only 'SUCCESS' OpenLineage events are taken into account in this connector.
+    Only OpenLineage events that indicate successfull data movement (COMPLETE, RUNNING, START) are taken into account in this connector.
 
     Configuring OpenLineage integrations: https://openlineage.io/docs/integrations/about
     """
@@ -195,18 +195,18 @@ def _render_pipeline_name(cls, pipeline_details: OpenLineageEvent) -> str:
         return f"{namespace}-{name}"
 
     @classmethod
-    def _filter_event_by_type(
-        cls, event: OpenLineageEvent, event_type: EventType
+    def _filter_event_by_types(
+        cls, event: OpenLineageEvent, event_types: List[EventType]
     ) -> Optional[Dict]:
         """
-        returns event if it's of particular event_type.
+        returns event if it's of one of the particular event_types.
         for example - for lineage events we will be only looking for EventType.COMPLETE event type.
 
         :param event: Open Lineage raw event.
-        :param event_type: type of event we are looking for.
-        :return: Open Lineage event if matches event_type, otherwise None
+        :param event_types: list of event types we are looking for.
+        :return: Open Lineage event if matches one of the event_types, otherwise None
         """
-        return event if event.event_type == event_type else {}
+        return event if event.event_type in event_types else {}
 
     @classmethod
     def _get_om_table_columns(cls, table_input: Dict) -> Optional[List]:
@@ -471,7 +471,10 @@ def get_pipelines_list(self) -> Optional[List[Any]]:
                         _result = message_to_open_lineage_event(
                             json.loads(message.value())
                         )
-                        result = self._filter_event_by_type(_result, EventType.COMPLETE)
+                        result = self._filter_event_by_types(
+                            _result,
+                            [EventType.COMPLETE, EventType.RUNNING, EventType.START],
+                        )
                         if result:
                             yield result
                     except Exception as e:

ingestion/src/metadata/ingestion/source/pipeline/openlineage/models.py

@@ -86,4 +86,9 @@ class EventType(str, Enum):
     List of used OpenLineage event types.
     """
 
+    START = "START"
+    RUNNING = "RUNNING"
     COMPLETE = "COMPLETE"
+    ABORT = "ABORT"
+    FAIL = "FAIL"
+    OTHER = "OTHER"

ingestion/tests/unit/topology/pipeline/test_openlineage.py

@@ -6,6 +6,7 @@
 from uuid import UUID
 
 from metadata.generated.schema.api.data.createTable import CreateTableRequest
+from metadata.generated.schema.api.lineage.addLineage import AddLineageRequest
 from metadata.generated.schema.entity.data.pipeline import Pipeline, Task
 from metadata.generated.schema.entity.services.connections.metadata.openMetadataConnection import (
     OpenMetadataConnection,
@@ -597,6 +598,187 @@ def test_get_create_table_request(self, mock_get_schema_fqn, mock_get_table_fqn)
                 create_request.columns[i].dataTypeDisplay, expected_type_display
             )
 
+    @patch("confluent_kafka.Consumer")
+    def test_get_pipelines_list_filters_complete_events(self, mock_consumer_class):
+        """Test that get_pipelines_list returns COMPLETE events"""
+        event = copy.deepcopy(VALID_EVENT)
+        event["eventType"] = "COMPLETE"
+        self.setup_mock_consumer_with_kafka_event(event)
+
+        result_generator = self.open_lineage_source.get_pipelines_list()
+        results = list(result_generator)
+
+        self.assertEqual(len(results), 1)
+        self.assertIsInstance(results[0], OpenLineageEvent)
+        self.assertEqual(results[0].event_type, "COMPLETE")
+
+    @patch("confluent_kafka.Consumer")
+    def test_get_pipelines_list_filters_running_events(self, mock_consumer_class):
+        """Test that get_pipelines_list returns RUNNING events"""
+        event = copy.deepcopy(VALID_EVENT)
+        event["eventType"] = "RUNNING"
+        self.setup_mock_consumer_with_kafka_event(event)
+
+        result_generator = self.open_lineage_source.get_pipelines_list()
+        results = list(result_generator)
+
+        self.assertEqual(len(results), 1)
+        self.assertIsInstance(results[0], OpenLineageEvent)
+        self.assertEqual(results[0].event_type, "RUNNING")
+
+    @patch("confluent_kafka.Consumer")
+    def test_get_pipelines_list_filters_start_events(self, mock_consumer_class):
+        """Test that get_pipelines_list returns START events"""
+        event = copy.deepcopy(VALID_EVENT)
+        event["eventType"] = "START"
+        self.setup_mock_consumer_with_kafka_event(event)
+
+        result_generator = self.open_lineage_source.get_pipelines_list()
+        results = list(result_generator)
+
+        self.assertEqual(len(results), 1)
+        self.assertIsInstance(results[0], OpenLineageEvent)
+        self.assertEqual(results[0].event_type, "START")
+
+    @patch("confluent_kafka.Consumer")
+    def test_get_pipelines_list_filters_out_fail_events(self, mock_consumer_class):
+        """Test that get_pipelines_list filters out FAIL events"""
+        event = copy.deepcopy(VALID_EVENT)
+        event["eventType"] = "FAIL"
+        self.setup_mock_consumer_with_kafka_event(event)
+
+        result_generator = self.open_lineage_source.get_pipelines_list()
+        results = list(result_generator)
+
+        self.assertEqual(len(results), 0)
+
+    @patch("confluent_kafka.Consumer")
+    def test_get_pipelines_list_filters_out_abort_events(self, mock_consumer_class):
+        """Test that get_pipelines_list filters out ABORT events"""
+        event = copy.deepcopy(VALID_EVENT)
+        event["eventType"] = "ABORT"
+        self.setup_mock_consumer_with_kafka_event(event)
+
+        result_generator = self.open_lineage_source.get_pipelines_list()
+        results = list(result_generator)
+
+        self.assertEqual(len(results), 0)
+
+    @patch(
+        "metadata.ingestion.source.pipeline.openlineage.metadata.OpenlineageSource._get_table_fqn_from_om"
+    )
+    def test_lineage_merge_start_with_data_running_without(self, mock_get_table_fqn):
+        """
+        Test that START event with lineage data followed by RUNNING event without
+        lineage data does not overwrite existing lineage in the database.
+
+        This simulates Flink streaming jobs where:
+        - START event contains initial lineage
+        - RUNNING events are heartbeats with no/empty lineage
+
+        The test verifies the complete flow:
+        1. START event creates lineage with column details
+        2. RUNNING event with empty data is processed
+        3. Query back the lineage - it should still have the original data
+        """
+        # Create START event with lineage data
+        start_event = copy.deepcopy(FULL_OL_KAFKA_EVENT)
+        start_event["eventType"] = "START"
+
+        # Create RUNNING event with same job but no lineage (empty inputs/outputs)
+        running_event = copy.deepcopy(FULL_OL_KAFKA_EVENT)
+        running_event["eventType"] = "RUNNING"
+        running_event["inputs"] = []
+        running_event["outputs"] = []
+
+        # Mock table FQN lookup
+        def mock_fqn_side_effect(table_details):
+            return f"testService.shopify.{table_details.name}"
+
+        mock_get_table_fqn.side_effect = mock_fqn_side_effect
+
+        # Mock metadata.get_by_name for table lookups
+        from_table_id = "69fc8906-4a4a-45ab-9a54-9cc2d399e10e"
+        to_table_id = "59fc8906-4a4a-45ab-9a54-9cc2d399e10e"
+
+        def mock_get_uuid_by_name(entity, fqn):
+            if fqn == "testService.shopify.raw_product_catalog":
+                return Mock(id=from_table_id)
+            elif fqn == "testService.shopify.fact_order_new5":
+                return Mock(id=to_table_id)
+            elif "openlineage_source" in fqn:  # Pipeline entity
+                return Mock(id=Mock(root="79fc8906-4a4a-45ab-9a54-9cc2d399e10e"))
+            return None
+
+        # Process START event with lineage
+        start_ol_event = message_to_open_lineage_event(start_event)
+        with patch.object(
+            OpenMetadataConnection,
+            "get_by_name",
+            create=True,
+            side_effect=mock_get_uuid_by_name,
+        ):
+            start_lineage_results = list(
+                self.open_lineage_source.yield_pipeline_lineage_details(start_ol_event)
+            )
+
+        # Process RUNNING event without lineage
+        running_ol_event = message_to_open_lineage_event(running_event)
+        with patch.object(
+            OpenMetadataConnection,
+            "get_by_name",
+            create=True,
+            side_effect=mock_get_uuid_by_name,
+        ):
+            running_lineage_results = list(
+                self.open_lineage_source.yield_pipeline_lineage_details(
+                    running_ol_event
+                )
+            )
+
+        # Extract lineage requests from START event
+        start_lineage_requests = [
+            r.right
+            for r in start_lineage_results
+            if r.right and isinstance(r.right, AddLineageRequest)
+        ]
+
+        # Extract lineage requests from RUNNING event
+        running_lineage_requests = [
+            r.right
+            for r in running_lineage_results
+            if r.right and isinstance(r.right, AddLineageRequest)
+        ]
+
+        # Verify START event produced lineage with column details
+        start_requests_with_columns = [
+            req
+            for req in start_lineage_requests
+            if req.edge.lineageDetails and req.edge.lineageDetails.columnsLineage
+        ]
+        self.assertGreater(
+            len(start_requests_with_columns),
+            0,
+            "START event should produce lineage requests with column details",
+        )
+
+        # Count column lineage entries from START
+        start_column_count = sum(
+            len(req.edge.lineageDetails.columnsLineage)
+            for req in start_requests_with_columns
+        )
+        self.assertGreater(
+            start_column_count, 0, "START event should have column lineage"
+        )
+
+        # Key assertion: RUNNING event with empty inputs/outputs produces no lineage requests
+        # This prevents empty data from being sent to the database
+        self.assertEqual(
+            len(running_lineage_requests),
+            0,
+            "RUNNING event with empty inputs/outputs should not produce any lineage requests",
+        )
+
 
 if __name__ == "__main__":
     unittest.main()