add namespace-based DB service resolution for openLineage db_table lookups

jsingh-yelp · jsingh-yelp · commit ba89dc759694 · 2026-04-02T18:26:18.000-04:00
diff --git a/ingestion/src/metadata/ingestion/source/pipeline/openlineage/metadata.py b/ingestion/src/metadata/ingestion/source/pipeline/openlineage/metadata.py
@@ -32,6 +32,7 @@
     KinesisBrokerConfig,
     OpenLineageConnection,
 )
+from metadata.generated.schema.entity.services.databaseService import DatabaseService
 from metadata.generated.schema.entity.services.ingestionPipelines.status import (
     StackTraceError,
 )
@@ -69,6 +70,11 @@
     get_or_create_pipeline_service,
     resolve_pipeline_service_type,
 )
+from metadata.ingestion.source.pipeline.openlineage.table_resolver import (
+    extract_db_scheme_from_namespace,
+    find_service_by_namespace_mapping,
+    find_services_by_scheme,
+)
 from metadata.ingestion.source.pipeline.openlineage.utils import (
     FQNNotFoundException,
     message_to_open_lineage_event,
@@ -113,6 +119,9 @@ def create(
     def prepare(self):
         self._service_cache = {}
         self._current_pipeline_service = None
+        self._entity_cache: Dict[str, Any] = {}
+        self._namespace_to_service_cache: Dict[str, Optional[List[str]]] = {}
+        self._db_service_type_map: Dict[str, str] = self._build_db_service_type_map()
 
     def close(self) -> None:
         self.metadata.compute_percentile(Pipeline, self.today)
@@ -214,7 +223,94 @@ def _get_topic_details(data: Dict) -> TopicDetails:
 
         return TopicDetails(name=name, broker_hostname=broker_hostname)
 
-    def _get_table_fqn(self, table_details: TableDetails) -> Optional[str]:
+    def _get_by_name_cached(self, entity_class, fqn_str: str, **kwargs):
+        """Wrapper around metadata.get_by_name with in-memory caching."""
+        if not hasattr(self, "_entity_cache"):
+            return self.metadata.get_by_name(entity_class, fqn_str, **kwargs)
+        key = f"{entity_class.__name__}:{fqn_str}"
+        if key not in self._entity_cache:
+            self._entity_cache[key] = self.metadata.get_by_name(
+                entity_class, fqn_str, **kwargs
+            )
+        return self._entity_cache[key]
+
+    def _build_db_service_type_map(self):
+        """Build a map of {service_name: DatabaseServiceType} filtered to configured dbServiceNames."""
+        type_map = {}
+        for service_name in self.get_db_service_names():
+            try:
+                svc = self.metadata.get_by_name(DatabaseService, service_name)
+                if svc and svc.serviceType:
+                    type_map[service_name] = svc.serviceType
+            except Exception:
+                logger.debug(f"Could not fetch DB service: {service_name}")
+        return type_map
+
+    def _resolve_db_services_for_namespace(self, namespace: str) -> Optional[List[str]]:
+        """
+        Resolve which DB services to search for a given OL dataset namespace.
+
+        Resolution order:
+        1. Check namespaceToServiceMapping config (exact then prefix match).
+        2. Extract scheme from namespace, filter services by matching DB type.
+           If exactly one match -> use it. If multiple -> log warning and return all.
+        3. Return None -> caller falls back to all dbServiceNames.
+        """
+        if not hasattr(self, "_namespace_to_service_cache"):
+            return None
+
+        if namespace in self._namespace_to_service_cache:
+            return self._namespace_to_service_cache[namespace]
+
+        result = None
+        configured = set(self.get_db_service_names() or [])
+
+        mapping = (
+            getattr(self.service_connection, "namespaceToServiceMapping", None) or {}
+        )
+        mapped_service = find_service_by_namespace_mapping(namespace, mapping)
+        if mapped_service and mapped_service in configured:
+            result = [mapped_service]
+        else:
+            # Auto-discover by extracting the DB scheme from the namespace URL
+            db_scheme = extract_db_scheme_from_namespace(namespace)
+            if db_scheme:
+                matched = find_services_by_scheme(db_scheme, self._db_service_type_map)
+                if len(matched) == 1:
+                    result = matched
+                elif len(matched) > 1:
+                    logger.warning(
+                        f"Namespace '{namespace}' (scheme={db_scheme}) matches "
+                        f"multiple DB services: {matched}. Configure "
+                        f"'namespaceToServiceMapping' to disambiguate."
+                    )
+                    result = matched
+
+        self._namespace_to_service_cache[namespace] = result
+        return result
+
+    def _find_table_fqn(
+        self, table_details: TableDetails, services: Optional[List[str]] = None
+    ) -> str:
+        search_services = services or self.get_db_service_names()
+        for db_service in search_services:
+            result = fqn.build(
+                metadata=self.metadata,
+                entity_type=Table,
+                service_name=db_service,
+                database_name=table_details.database,
+                schema_name=table_details.schema,
+                table_name=table_details.name,
+            )
+            if result:
+                return result
+        raise FQNNotFoundException(
+            f"Table FQN not found for {table_details} in services {search_services}"
+        )
+
+    def _get_table_fqn(
+        self, table_details: TableDetails, namespace: Optional[str] = None
+    ) -> Optional[str]:
         if not self.get_db_service_names():
             if not self._db_service_names_warned:
                 logger.warning(
@@ -224,16 +320,39 @@ def _get_table_fqn(self, table_details: TableDetails) -> Optional[str]:
                 )
                 self._db_service_names_warned = True
             return None
+
+        resolved_services = self._resolve_db_services_for_namespace(namespace)
+
         try:
-            return self._get_table_fqn_from_om(table_details)
+            return self._find_table_fqn(table_details, services=resolved_services)
         except FQNNotFoundException:
             try:
-                schema_fqn = self._get_schema_fqn_from_om(table_details.schema)
-
+                schema_fqn = self._get_schema_fqn_from_om(
+                    table_details.schema, services=resolved_services
+                )
                 return f"{schema_fqn}.{table_details.name}"
             except FQNNotFoundException:
                 return None
 
+    def _get_table_fqn_from_om(
+        self, table_details: TableDetails, services: Optional[List[str]] = None
+    ) -> str:
+        """
+        Looks for matching Table entity in OM across all configured DB services.
+        """
+        for db_service in services or self.get_db_service_names():
+            result = fqn.build(
+                metadata=self.metadata,
+                entity_type=Table,
+                service_name=db_service,
+                database_name=table_details.database,
+                schema_name=table_details.schema,
+                table_name=table_details.name,
+            )
+            if result:
+                return result
+        raise FQNNotFoundException(f"Table FQN not found for {table_details}")
+
     def _build_broker_to_service_map(self) -> Dict[str, str]:
         """
         Build a cache mapping broker hostnames to messaging service FQNs.
@@ -309,15 +428,18 @@ def _get_topic_entity(self, topic_details: TopicDetails) -> Optional[Topic]:
             logger.warning(f"Error finding topic for {topic_details.name}: {exc}")
             return None
 
-    def _get_schema_fqn_from_om(self, schema: str) -> Optional[str]:
+    def _get_schema_fqn_from_om(
+        self, schema: str, services: Optional[List[str]] = None
+    ) -> Optional[str]:
         """
         Based on partial schema name look for any matching DatabaseSchema object in open metadata.
 
         :param schema: schema name
+        :param services: optional list of service names to search
         :return: fully qualified name of a DatabaseSchema in Open Metadata
         """
         result = None
-        services = self.get_db_service_names()
+        services = services or self.get_db_service_names()
 
         for db_service in services:
             result = fqn.build(
@@ -469,7 +591,10 @@ def _build_ol_name_to_fqn_map(self, tables: List):
             entity_details = self._get_entity_details(table)
             if entity_details.entity_type != "table":
                 continue
-            table_fqn = self._get_table_fqn(entity_details.table_details)
+            table_fqn = self._get_table_fqn(
+                entity_details.table_details,
+                namespace=table.get("namespace"),
+            )
 
             if table_fqn:
                 result[OpenlineageSource._get_ol_table_name(table)] = table_fqn
@@ -505,7 +630,10 @@ def _get_column_lineage(
             if entity_details.entity_type != "table":
                 continue
 
-            output_table_fqn = self._get_table_fqn(entity_details.table_details)
+            output_table_fqn = self._get_table_fqn(
+                entity_details.table_details,
+                namespace=table.get("namespace"),
+            )
             for field_name, field_spec in (
                 table.get("facets", {})
                 .get("columnLineage", {})
@@ -604,15 +732,16 @@ def yield_pipeline_lineage_details(
                     if create_table_request:
                         yield create_table_request
 
-                    table_fqn = self._get_table_fqn(entity_details.table_details)
+                    table_fqn = self._get_table_fqn(
+                        entity_details.table_details,
+                        namespace=entity_data.get("namespace"),
+                    )
 
                     if table_fqn:
                         entity_list.append(
                             LineageNode(
                                 fqn=TableFQN(value=table_fqn),
-                                uuid=self.metadata.get_by_name(
-                                    Table, table_fqn
-                                ).id.root,
+                                uuid=self._get_by_name_cached(Table, table_fqn).id.root,
                                 node_type="table",
                             )
                         )
diff --git a/ingestion/src/metadata/ingestion/source/pipeline/openlineage/table_resolver.py b/ingestion/src/metadata/ingestion/source/pipeline/openlineage/table_resolver.py
@@ -0,0 +1,147 @@
+#  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 a database service name from an OpenLineage dataset namespace.
+
+OpenLineage dataset namespaces carry the data store type in the URL scheme
+(e.g. ``mysql://host:3306/db``, ``redshift://cluster:5439/db``).  This module
+uses that scheme to narrow lineage table lookups to services of the matching
+type, eliminating false matches when the same table name exists in multiple
+services.
+
+Resolution order (per namespace):
+1. Explicit ``namespaceToServiceMapping`` config — exact then prefix match.
+2. Scheme-based auto-discovery — if exactly one configured DB service has the
+   matching type, use it.  If multiple match, log a warning and fall back.
+3. Caller falls back to existing suffix search across all ``dbServiceNames``.
+"""
+
+from typing import Dict, List, Optional
+from urllib.parse import urlparse
+
+from metadata.generated.schema.entity.services.databaseService import (
+    DatabaseServiceType,
+)
+from metadata.utils.logger import ingestion_logger
+
+logger = ingestion_logger()
+
+# Maps OpenLineage dataset namespace URI schemes to OMD DatabaseServiceType.
+# See: https://openlineage.io/docs/spec/naming/
+NAMESPACE_SCHEME_TO_SERVICE_TYPE: Dict[str, DatabaseServiceType] = {
+    "awsathena": DatabaseServiceType.Athena,
+    "bigquery": DatabaseServiceType.BigQuery,
+    "cassandra": DatabaseServiceType.Cassandra,
+    "db2": DatabaseServiceType.Db2,
+    "hive": DatabaseServiceType.Hive,
+    "mssql": DatabaseServiceType.Mssql,
+    "mysql": DatabaseServiceType.Mysql,
+    "oracle": DatabaseServiceType.Oracle,
+    "postgres": DatabaseServiceType.Postgres,
+    "redshift": DatabaseServiceType.Redshift,
+    "snowflake": DatabaseServiceType.Snowflake,
+    "sqlserver": DatabaseServiceType.Synapse,
+    "teradata": DatabaseServiceType.Teradata,
+    "trino": DatabaseServiceType.Trino,
+}
+
+
+def extract_db_scheme_from_namespace(namespace: str) -> Optional[str]:
+    """
+    Extract the URL scheme from an OpenLineage dataset namespace.
+
+    >>> extract_namespace_scheme("mysql://host:3306/db")
+    'mysql'
+    >>> extract_namespace_scheme("redshift://cluster:5439/db")
+    'redshift'
+    >>> extract_namespace_scheme("airflow")
+    None
+    """
+    try:
+        scheme = urlparse(namespace).scheme
+        return scheme.lower() if scheme else None
+    except Exception:
+        return None
+
+
+def find_service_by_namespace_mapping(
+    namespace: str,
+    mapping: Dict[str, str],
+) -> Optional[str]:
+    """
+    Look up a database service name from a user-configured
+    ``namespaceToServiceMapping`` dict (namespace-prefix → OMD service name).
+
+    Resolution order:
+    1. Exact match — ``mapping[namespace]``
+    2. Prefix match — bidirectional ``startsWith`` check.
+
+    Example mapping::
+
+        {
+            "mysql://cluster-a:3306": "mysql-cluster-a",
+            "redshift://prod.us-east-1:5439": "redshift-prod",
+        }
+
+    With namespace ``"mysql://cluster-a:3306/mydb"`` → returns ``"mysql-cluster-a"``
+    (namespace starts with the key).
+
+    Returns the mapped service name, or ``None`` if no entry matches.
+    """
+    if not namespace or not mapping:
+        return None
+
+    if namespace in mapping:
+        return mapping[namespace]
+
+    # check if there is a partial match
+    for key, service_name in mapping.items():
+        if namespace.startswith(key) or key.startswith(namespace):
+            return service_name
+
+    return None
+
+
+def find_services_by_scheme(
+    scheme: str,
+    db_service_type_map: Dict[str, DatabaseServiceType],
+) -> List[str]:
+    """
+    Filter a pre-built ``{service_name: DatabaseServiceType}`` map to only
+    those whose type matches the given URL scheme.
+
+    If the scheme is recognized (present in ``NAMESPACE_SCHEME_TO_SERVICE_TYPE``),
+    returns only services of that exact type.
+
+    If the scheme is **not** recognized, returns services whose
+    types are *not* in the known map — i.e. custom non-standard
+    services that are more likely to be the correct match.
+
+    Returns a list of matching service names (may be empty or contain
+    multiple entries when several services of the same type are configured).
+    """
+    target_type = NAMESPACE_SCHEME_TO_SERVICE_TYPE.get(scheme)
+
+    if target_type:
+        return [
+            name
+            for name, svc_type in db_service_type_map.items()
+            if svc_type == target_type
+        ]
+
+    known_types = set(NAMESPACE_SCHEME_TO_SERVICE_TYPE.values())
+
+    return [
+        name
+        for name, svc_type in db_service_type_map.items()
+        if svc_type not in known_types
+    ]
diff --git a/ingestion/tests/unit/topology/pipeline/test_openlineage.py b/ingestion/tests/unit/topology/pipeline/test_openlineage.py
