Fixes 26434: [openlinege] Resolve openlineage pipeline/job to it's integration type (#26821)

jsingh-yelp · ayush-shah · mohittilala · commit 00ded0e064a8 · 2026-04-23T22:08:05.000+05:30
* resolve openlineage pipeline/job to it's integration type

* include kafka broker port for _build_broker_to_service_map function

---------

Co-authored-by: Ayush Shah &lt;ayush02shah12@gmail.com&gt;
diff --git a/ingestion/src/metadata/ingestion/source/pipeline/openlineage/metadata.py b/ingestion/src/metadata/ingestion/source/pipeline/openlineage/metadata.py
@@ -62,6 +62,13 @@
     TopicDetails,
     TopicFQN,
 )
+from metadata.ingestion.source.pipeline.openlineage.service_resolver import (
+    build_service_name,
+    extract_integration_type,
+    find_pipeline_by_namespace,
+    get_or_create_pipeline_service,
+    resolve_pipeline_service_type,
+)
 from metadata.ingestion.source.pipeline.openlineage.utils import (
     FQNNotFoundException,
     message_to_open_lineage_event,
@@ -87,6 +94,8 @@ class OpenlineageSource(PipelineServiceSource):
     """
 
     _db_service_names_warned: bool = False
+    _service_cache: Dict[str, str]
+    _current_pipeline_service: Optional[str] = None
 
     @classmethod
     def create(
@@ -102,7 +111,8 @@ def create(
         return cls(config, metadata)
 
     def prepare(self):
-        """Nothing to prepare"""
+        self._service_cache = {}
+        self._current_pipeline_service = None
 
     def close(self) -> None:
         self.metadata.compute_percentile(Pipeline, self.today)
@@ -192,12 +202,16 @@ def _get_topic_details(data: Dict) -> TopicDetails:
         except KeyError:
             raise ValueError("Topic name is not present")
 
-        broker_hostname = urlparse(namespace).hostname
+        parsed = urlparse(namespace)
+        broker_hostname = parsed.hostname
         if not broker_hostname:
             raise ValueError(
                 f"Could not extract broker hostname from namespace: {namespace}"
             )
 
+        if parsed.port:
+            broker_hostname = f"{broker_hostname}:{parsed.port}"
+
         return TopicDetails(name=name, broker_hostname=broker_hostname)
 
     def _get_table_fqn(self, table_details: TableDetails) -> Optional[str]:
@@ -240,9 +254,9 @@ def _build_broker_to_service_map(self) -> Dict[str, str]:
                         bootstrap_servers = svc.connection.config.bootstrapServers or ""
                         svc_fqn = svc.fullyQualifiedName.root
                         for broker in bootstrap_servers.split(","):
-                            hostname = broker.strip().split(":")[0]
-                            if hostname:
-                                self._broker_to_service[hostname] = svc_fqn
+                            broker = broker.strip()
+                            if broker:
+                                self._broker_to_service[broker] = svc_fqn
                     except Exception:
                         logger.debug(
                             f"Could not extract bootstrapServers from service {svc.name}"
@@ -514,16 +528,50 @@ def _get_column_lineage(
 
         return OpenlineageSource._create_output_lineage_dict(_result)
 
+    def _resolve_pipeline_service(self, pipeline_details: OpenLineageEvent) -> str:
+        """
+        Resolve the pipeline service for the current event.
+
+        Resolution order:
+        1. **Namespace fallback** — try ``namespace.jobName`` as a pipeline
+           FQN.  If a pipeline already exists (e.g. ingested by a native
+           Airflow connector), reuse its service.
+        2. **Integration type** — extract from
+           ``job.facets.jobType.integration`` and create a typed service
+           (e.g. ``spark_openlineage``).
+        3. **Default** — fall back to the configured OpenLineage service.
+        """
+        fallback = self.context.get().pipeline_service
+
+        ns_result = find_pipeline_by_namespace(self.metadata, pipeline_details)
+        if ns_result:
+            service_name, _ = ns_result
+            return service_name
+
+        integration = extract_integration_type(pipeline_details)
+        service_name = build_service_name(integration, fallback)
+
+        if service_name != fallback:
+            service_type = resolve_pipeline_service_type(integration)
+            get_or_create_pipeline_service(
+                self.metadata, service_name, service_type, self._service_cache
+            )
+
+        return service_name
+
     def yield_pipeline(
         self, pipeline_details: OpenLineageEvent
     ) -> Iterable[Either[CreatePipelineRequest]]:
         pipeline_name = self.get_pipeline_name(pipeline_details)
+        self._current_pipeline_service = self._resolve_pipeline_service(
+            pipeline_details
+        )
         try:
             description = f"""```json
             {json.dumps(pipeline_details.run_facet, indent=4).strip()}```"""
             request = CreatePipelineRequest(
                 name=pipeline_name,
-                service=self.context.get().pipeline_service,
+                service=self._current_pipeline_service,
                 description=description,
             )
 
@@ -597,10 +645,13 @@ def yield_pipeline_lineage_details(
             for n in product(input_edges, output_edges)
         ]
 
+        service_name = (
+            self._current_pipeline_service or self.context.get().pipeline_service
+        )
         pipeline_fqn = fqn.build(
             metadata=self.metadata,
             entity_type=Pipeline,
-            service_name=self.context.get().pipeline_service,
+            service_name=service_name,
             pipeline_name=self.context.get().pipeline,
         )
 
diff --git a/ingestion/src/metadata/ingestion/source/pipeline/openlineage/service_resolver.py b/ingestion/src/metadata/ingestion/source/pipeline/openlineage/service_resolver.py
@@ -0,0 +1,151 @@
+#  Copyright 2025 Collate
+#  Licensed under the Collate Community License, Version 1.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#  https://github.com/open-metadata/OpenMetadata/blob/main/ingestion/LICENSE
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+
+"""
+Resolves the pipeline service type and name from OpenLineage event metadata.
+
+OpenLineage events carry integration identity (Spark, Flink, Airflow, etc.)
+in job facets. This module extracts that information and maps it to the
+appropriate OMD PipelineServiceType, creating services as needed.
+"""
+
+from typing import Dict, Optional, Tuple
+
+from metadata.generated.schema.api.services.createPipelineService import (
+    CreatePipelineServiceRequest,
+)
+from metadata.generated.schema.entity.data.pipeline import Pipeline
+from metadata.generated.schema.entity.services.pipelineService import (
+    PipelineService,
+    PipelineServiceType,
+)
+from metadata.generated.schema.type.basic import EntityName
+from metadata.ingestion.ometa.ometa_api import OpenMetadata
+from metadata.ingestion.source.pipeline.openlineage.models import OpenLineageEvent
+from metadata.utils.logger import ingestion_logger
+
+logger = ingestion_logger()
+
+INTEGRATION_TO_SERVICE_TYPE: Dict[str, PipelineServiceType] = {
+    "spark": PipelineServiceType.Spark,
+    "flink": PipelineServiceType.Flink,
+    "airflow": PipelineServiceType.Airflow,
+    "dbt": PipelineServiceType.DBTCloud,
+    "dagster": PipelineServiceType.Dagster,
+}
+
+SERVICE_NAME_SUFFIX = "_openlineage"
+
+
+def extract_integration_type(event: OpenLineageEvent) -> Optional[str]:
+    """
+    Extract the integration type from an OpenLineage event via the
+    standard ``job.facets.jobType.integration`` field.
+
+    Returns a lowercase string like "spark", "flink", "airflow", or None.
+    """
+    try:
+        integration = event.job["facets"]["jobType"]["integration"]
+        if integration:
+            return integration.strip().lower()
+    except (KeyError, TypeError, AttributeError):
+        pass
+
+    return None
+
+
+def find_pipeline_by_namespace(
+    metadata: OpenMetadata,
+    event: OpenLineageEvent,
+) -> Optional[Tuple[str, Pipeline]]:
+    """
+    Try to find an existing pipeline using ``namespace.jobName`` as FQN.
+
+    When pipelines are ingested by native connectors (e.g. Airflow), the OL
+    job namespace often matches the pipeline service name already present in
+    OMD.  Checking this first avoids creating duplicate pipelines under a new
+    service.
+
+    Returns ``(service_name, pipeline_entity)`` on hit, or ``None``.
+    """
+    try:
+        namespace = event.job.get("namespace")
+        name = event.job.get("name")
+    except (AttributeError, TypeError):
+        return None
+
+    if not namespace or not name:
+        return None
+
+    fallback_fqn = f"{namespace}.{name}"
+    existing = metadata.get_by_name(Pipeline, fallback_fqn)
+    if existing:
+        logger.info(f"Resolved pipeline via namespace fallback: {fallback_fqn}")
+        return namespace, existing
+
+    return None
+
+
+def resolve_pipeline_service_type(
+    integration: Optional[str],
+) -> PipelineServiceType:
+    """Map an integration string to a PipelineServiceType enum."""
+    if integration and integration in INTEGRATION_TO_SERVICE_TYPE:
+        return INTEGRATION_TO_SERVICE_TYPE[integration]
+    return PipelineServiceType.OpenLineage
+
+
+def build_service_name(integration: Optional[str], fallback_service: str) -> str:
+    """
+    Build the pipeline service name.
+
+    If integration is recognized, returns "{integration}_openlineage".
+    Otherwise returns the fallback (the configured OpenLineage service name).
+    """
+    if integration and integration in INTEGRATION_TO_SERVICE_TYPE:
+        return f"{integration}{SERVICE_NAME_SUFFIX}"
+    return fallback_service
+
+
+def get_or_create_pipeline_service(
+    metadata: OpenMetadata,
+    service_name: str,
+    service_type: PipelineServiceType,
+    _cache: Optional[Dict[str, str]] = None,
+) -> str:
+    """
+    Ensure a PipelineService with the given name and type exists in OMD.
+
+    Uses an external cache dict to avoid repeated API calls within a session.
+    Returns the service name.
+    """
+    if _cache is not None and service_name in _cache:
+        return _cache[service_name]
+
+    existing = metadata.get_by_name(PipelineService, service_name)
+    if existing:
+        if _cache is not None:
+            _cache[service_name] = service_name
+        return service_name
+
+    logger.info(
+        f"Creating pipeline service '{service_name}' with type '{service_type.value}'"
+    )
+    request = CreatePipelineServiceRequest(
+        name=EntityName(service_name),
+        serviceType=service_type,
+    )
+    metadata.create_or_update(request)
+
+    if _cache is not None:
+        _cache[service_name] = service_name
+
+    return service_name
diff --git a/ingestion/tests/unit/topology/pipeline/test_openlineage.py b/ingestion/tests/unit/topology/pipeline/test_openlineage.py
@@ -1169,7 +1169,7 @@ def test_entity_detection_kafka_namespace_returns_topic(self):
         self.assertIsNotNone(result.topic_details)
         self.assertIsNone(result.table_details)
         self.assertEqual(result.topic_details.name, "my-events-topic")
-        self.assertEqual(result.topic_details.broker_hostname, "broker-host")
+        self.assertEqual(result.topic_details.broker_hostname, "broker-host:9092")
 
     def test_entity_detection_non_kafka_namespace_returns_table(self):
         """Test that non-kafka namespaces (e.g. bigquery, hive) are detected as tables."""
@@ -1193,7 +1193,7 @@ def test_topic_details_extraction_various_broker_formats(self):
             {"name": "topic1", "namespace": "kafka://my-broker:9092"}
         )
         self.assertEqual(result.name, "topic1")
-        self.assertEqual(result.broker_hostname, "my-broker")
+        self.assertEqual(result.broker_hostname, "my-broker:9092")
 
         # Broker without port
         result = OpenlineageSource._get_topic_details(
diff --git a/ingestion/tests/unit/topology/pipeline/test_service_resolver.py b/ingestion/tests/unit/topology/pipeline/test_service_resolver.py
