Fixes #27538: [OpenLineage] Add AWS Glue, Kusto, and Cosmos DB dataset naming support (#27533)

mohittilala · mohittilala · commit 1b8ea815ae3a · 2026-04-23T22:08:05.000+05:30
* feat(openlineage): add AWS Glue, Kusto, and Cosmos DB dataset naming support

* Add better logging for exceptions

* Add cosmos defensive check for `colls/` prefix in names

* Correct spell mistakes and wrap long line

* fix: guard None source attributes in _sort_array_entity_fields

Optional array fields (tasks, columns, fields) on existing OMD entities
can be None when the API omits empty arrays in responses. Iterating over
None crashed with TypeError. Mirrors the existing `or []` guard already
applied to destination_attributes on the line below.

* Add test for schema not found case retuning none instead of rasing exception
diff --git a/ingestion/src/metadata/ingestion/source/pipeline/openlineage/metadata.py b/ingestion/src/metadata/ingestion/source/pipeline/openlineage/metadata.py
@@ -13,6 +13,7 @@
 OpenLineage source to extract metadata from Kafka or Kinesis events
 """
 import json
+import re
 import time
 import traceback
 from collections import defaultdict
@@ -99,7 +100,8 @@ class OpenlineageSource(PipelineServiceSource):
     Works under the assumption that OpenLineage integrations produce events to Kafka topic or Kinesis stream,
     which is a source of events for this connector.
 
-    Only OpenLineage events that indicate successfull data movement (COMPLETE, RUNNING, START) are taken into account in this connector.
+    Only OpenLineage events that indicate successful data movement (COMPLETE, RUNNING, START) are taken into account
+    in this connector.
 
     Configuring OpenLineage integrations: https://openlineage.io/docs/integrations/about
     """
@@ -183,6 +185,27 @@ def _get_table_details(cls, data: Dict) -> TableDetails:
                     "input table name cannot be retrieved from name attribute."
                 )
 
+        namespace = data.get("namespace", "")
+
+        # AWS Glue: arn:aws:glue:{region}:{account} / table/{database}/{table}
+        # Source: https://openlineage.io/docs/spec/naming/
+        if namespace.startswith("arn:aws:glue:"):
+            result = OpenlineageSource._parse_glue_table_name(name)
+            if result:
+                return result
+
+        # Azure Data Explorer (Kusto): azurekusto://{host} / {database}/{table}
+        if namespace.startswith("azurekusto://"):
+            result = OpenlineageSource._parse_slash_table_name(name)
+            if result:
+                return result
+
+        # Azure Cosmos DB: azurecosmos://{host}/dbs/{db} / colls/{collection}
+        if namespace.startswith("azurecosmos://"):
+            result = OpenlineageSource._parse_cosmos_table_name(namespace, name)
+            if result:
+                return result
+
         name_parts = name.split(".")
 
         if len(name_parts) < 2:
@@ -228,6 +251,59 @@ def _get_topic_details(data: Dict) -> TopicDetails:
 
         return TopicDetails(name=name, broker_hostname=broker_hostname)
 
+    @staticmethod
+    def _parse_glue_table_name(name: str) -> Optional[TableDetails]:
+        """
+        Parse AWS Glue OL dataset name: ``table/{database}/{table}``.
+
+        Glue EMR jobs emit a slash-separated name with a ``table/`` prefix instead
+        of the dot-separated ``schema.table`` convention used by SQL engines.
+
+        Source: https://github.com/OpenLineage/OpenLineage/blob/main/client/java/
+                src/main/java/io/openlineage/client/dataset/Naming.java (GlueNaming)
+        """
+        if not name.startswith("table/"):
+            return None
+        parts = name[len("table/") :].split("/")
+        if len(parts) < 2:
+            return None
+        return TableDetails(name=parts[-1].lower(), schema=parts[-2].lower())
+
+    @staticmethod
+    def _parse_slash_table_name(name: str) -> Optional[TableDetails]:
+        """
+        Parse slash-separated ``{database}/{table}`` OL dataset names.
+
+        Used by Azure Data Explorer (Kusto):
+          namespace ``azurekusto://{host}`` / name ``{database}/{table}``
+
+        Source: https://github.com/OpenLineage/OpenLineage/blob/main/client/java/
+                src/main/java/io/openlineage/client/dataset/Naming.java (KustoNaming)
+        """
+        parts = name.split("/")
+        if len(parts) < 2:
+            return None
+        return TableDetails(name=parts[-1].lower(), schema=parts[-2].lower())
+
+    @staticmethod
+    def _parse_cosmos_table_name(namespace: str, name: str) -> Optional[TableDetails]:
+        """
+        Parse Azure Cosmos DB OL dataset names.
+
+        The database lives in the namespace path (``azurecosmos://{host}/dbs/{db}``)
+        while the name field is ``colls/{collection}``.
+
+        Source: https://github.com/OpenLineage/OpenLineage/blob/main/client/java/
+                src/main/java/io/openlineage/client/dataset/Naming.java (CosmosNaming)
+        """
+        db_match = re.search(r"/dbs/([^/]+)", namespace)
+        coll_match = re.fullmatch(r"colls/([^/]+)", name)
+        if not db_match or not coll_match:
+            return None
+        return TableDetails(
+            name=coll_match.group(1).lower(), schema=db_match.group(1).lower()
+        )
+
     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"):
@@ -323,6 +399,11 @@ def _get_table_fqn(
                     )
                     return f"{schema_fqn}.{table_details.name}"
                 except FQNNotFoundException:
+                    logger.debug(
+                        f"Table '{table_details.name}' in schema '{table_details.schema}' "
+                        f"not found in services {resolved_services or self.get_db_service_names()}. "
+                        "Skipping lineage edge."
+                    )
                     return None
         except Exception:
             logger.warning(
@@ -465,7 +546,7 @@ def _get_schema_fqn_from_om(
 
         if not result:
             raise FQNNotFoundException(
-                f"Schema FQN not found within services: {services}"
+                f"Schema '{schema}' not found in services: {services}"
             )
 
         return result
@@ -565,14 +646,13 @@ def get_create_table_request(self, table: Dict) -> Optional[Either]:
         if not om_table_fqn:
             try:
                 om_schema_fqn = self._get_schema_fqn_from_om(table_details.schema)
-            except FQNNotFoundException as e:
-                return Either(
-                    left=StackTraceError(
-                        name="",
-                        error=f"Failed to get fully qualified schema name: {e}",
-                        stackTrace=traceback.format_exc(),
-                    )
+            except FQNNotFoundException:
+                logger.warning(
+                    f"Schema '{table_details.schema}' not found in configured services "
+                    f"{self.get_db_service_names()}. Skipping table creation for "
+                    f"'{table_details.name}'."
                 )
+                return None
 
             # After finding schema fqn (based on partial schema name) we know where we can create table
             # and we move forward with creating request.
@@ -870,10 +950,13 @@ def _poll_kafka(self, broker: KafkaBrokerConfig) -> Iterable[OpenLineageEvent]:
                         if result:
                             yield result
                     except Exception as e:
-                        logger.debug(e)
+                        logger.warning(
+                            f"Failed to parse OpenLineage event from Kafka message: {e}"
+                        )
+                        logger.debug(traceback.format_exc())
 
         except Exception as e:
-            traceback.print_exc()
+            logger.debug(traceback.format_exc())
             raise InvalidSourceException(f"Failed to read from Kafka: {str(e)}")
 
         finally:
@@ -933,12 +1016,15 @@ def _poll_kinesis(self, broker: KinesisBrokerConfig) -> Iterable[OpenLineageEven
                             if result:
                                 yield result
                         except Exception as e:
-                            logger.debug(e)
+                            logger.warning(
+                                f"Failed to parse OpenLineage event from Kinesis record: {e}"
+                            )
+                            logger.debug(traceback.format_exc())
 
                     time.sleep(pool_timeout)
 
         except Exception as e:
-            traceback.print_exc()
+            logger.debug(traceback.format_exc())
             raise InvalidSourceException(f"Failed to read from Kinesis: {str(e)}")
 
     def get_pipeline_name(self, pipeline_details: OpenLineageEvent) -> str:
diff --git a/ingestion/tests/unit/topology/pipeline/test_openlineage.py b/ingestion/tests/unit/topology/pipeline/test_openlineage.py
@@ -799,6 +799,28 @@ def test_get_create_table_request(self, mock_get_schema_fqn, mock_get_table_fqn)
                 create_request.columns[i].dataTypeDisplay, expected_type_display
             )
 
+    @patch(
+        "metadata.ingestion.source.pipeline.openlineage.metadata.OpenlineageSource._get_table_fqn_from_om"
+    )
+    @patch(
+        "metadata.ingestion.source.pipeline.openlineage.metadata.OpenlineageSource._get_schema_fqn_from_om"
+    )
+    def test_get_create_table_request_schema_not_found_returns_none(
+        self, mock_get_schema_fqn, mock_get_table_fqn
+    ):
+        """Schema not found in any configured service — returns None without raising."""
+        mock_get_table_fqn.side_effect = FQNNotFoundException("Table not found")
+        mock_get_schema_fqn.side_effect = FQNNotFoundException("Schema not found")
+        table_data = {
+            "name": "unknown_schema.employees",
+            "namespace": "bigquery",
+            "facets": {},
+        }
+
+        result = self.open_lineage_source.get_create_table_request(table_data)
+
+        assert result is None
+
     @patch("confluent_kafka.Consumer")
     def test_get_pipelines_list_filters_complete_events(self, mock_consumer_class):
         """Test that get_pipelines_list returns COMPLETE events"""
@@ -1689,16 +1711,18 @@ def test_yield_pipeline_lineage_topic_not_found_skips_gracefully(self):
         mock_pipeline = Mock()
         mock_pipeline.id.root = pipeline_id
 
-        with patch.object(
-            self.open_lineage_source, "metadata"
-        ) as mock_metadata, patch.object(
-            self.open_lineage_source,
-            "_get_table_fqn",
-            return_value="db-service.public.some_table",
-        ), patch.object(
-            self.open_lineage_source,
-            "get_create_table_request",
-            return_value=None,
+        with (
+            patch.object(self.open_lineage_source, "metadata") as mock_metadata,
+            patch.object(
+                self.open_lineage_source,
+                "_get_table_fqn",
+                return_value="db-service.public.some_table",
+            ),
+            patch.object(
+                self.open_lineage_source,
+                "get_create_table_request",
+                return_value=None,
+            ),
         ):
             # Empty messaging services list — no broker match for unknown-broker
             mock_metadata.list_all_entities.return_value = iter([])
@@ -1735,6 +1759,134 @@ def mock_get_by_name(entity, fqn, **kwargs):
             "No lineage edges should be produced when input topic cannot be resolved",
         )
 
+    def test_parse_glue_table_name_trino_glue_catalog_schema(self):
+        """Trino backed by AWS Glue Data Catalog uses the public schema and underscore-separated table names.
+        Verifies the parser handles the common Glue catalog table naming pattern correctly.
+        """
+        result = OpenlineageSource._parse_glue_table_name(
+            "table/public/order_line_items"
+        )
+        self.assertEqual(result.name, "order_line_items")
+        self.assertEqual(result.schema, "public")
+
+    def test_parse_glue_table_name_happy_path(self):
+        """Glue OL naming: table/{database}/{table} — source: Naming.java GlueNaming."""
+        result = OpenlineageSource._parse_glue_table_name("table/sales/users")
+        self.assertEqual(result.name, "users")
+        self.assertEqual(result.schema, "sales")
+
+    def test_parse_glue_table_name_normalizes_to_lowercase(self):
+        """Glue table and database names are normalized to lowercase for FQN matching."""
+        result = OpenlineageSource._parse_glue_table_name("table/Sales/Users")
+        self.assertEqual(result.name, "users")
+        self.assertEqual(result.schema, "sales")
+
+    def test_parse_glue_table_name_not_glue_format_returns_none(self):
+        """Names without the table/ prefix are not Glue format and return None."""
+        self.assertIsNone(OpenlineageSource._parse_glue_table_name("sales.users"))
+
+    def test_parse_glue_table_name_missing_table_part_returns_none(self):
+        """table/ prefix with only one path segment is malformed and returns None."""
+        self.assertIsNone(OpenlineageSource._parse_glue_table_name("table/only_db"))
+
+    def test_parse_slash_table_name_happy_path(self):
+        """Kusto OL naming: {database}/{table} — source: Naming.java KustoNaming."""
+        result = OpenlineageSource._parse_slash_table_name("mydb/mytable")
+        self.assertEqual(result.name, "mytable")
+        self.assertEqual(result.schema, "mydb")
+
+    def test_parse_slash_table_name_normalizes_to_lowercase(self):
+        """Kusto table and database names are normalized to lowercase for FQN matching."""
+        result = OpenlineageSource._parse_slash_table_name("MyDB/MyTable")
+        self.assertEqual(result.name, "mytable")
+        self.assertEqual(result.schema, "mydb")
+
+    def test_parse_slash_table_name_single_part_returns_none(self):
+        """A single path segment without a slash cannot be split into db/table and returns None."""
+        self.assertIsNone(OpenlineageSource._parse_slash_table_name("only_table"))
+
+    def test_parse_cosmos_table_name_happy_path(self):
+        """Cosmos OL naming: db from namespace /dbs/{db}, name colls/{coll} — source: Naming.java CosmosNaming."""
+        result = OpenlineageSource._parse_cosmos_table_name(
+            "azurecosmos://myaccount.documents.azure.com/dbs/mydb",
+            "colls/mycollection",
+        )
+        self.assertEqual(result.name, "mycollection")
+        self.assertEqual(result.schema, "mydb")
+
+    def test_parse_cosmos_table_name_normalizes_to_lowercase(self):
+        """Cosmos database and collection names are normalized to lowercase for FQN matching."""
+        result = OpenlineageSource._parse_cosmos_table_name(
+            "azurecosmos://host/dbs/MyDB", "colls/MyCollection"
+        )
+        self.assertEqual(result.name, "mycollection")
+        self.assertEqual(result.schema, "mydb")
+
+    def test_parse_cosmos_table_name_no_dbs_segment_returns_none(self):
+        """A Cosmos namespace without /dbs/{db} cannot provide the database name and returns None."""
+        self.assertIsNone(
+            OpenlineageSource._parse_cosmos_table_name(
+                "azurecosmos://host", "colls/mycoll"
+            )
+        )
+
+    def test_parse_cosmos_table_name_non_colls_name_returns_none(self):
+        """A Cosmos name not matching colls/{collection} is non-conformant and returns None."""
+        self.assertIsNone(
+            OpenlineageSource._parse_cosmos_table_name(
+                "azurecosmos://host/dbs/mydb", "mycollection"
+            )
+        )
+
+    def test_get_table_details_glue_namespace_parses_slash_name(self):
+        """AWS Glue EMR events use arn:aws:glue namespace + table/{db}/{table} name."""
+        data = {
+            "namespace": "arn:aws:glue:us-east-1:123456789012",
+            "name": "table/sales/users",
+        }
+        result = OpenlineageSource._get_table_details(data)
+        self.assertEqual(result.name, "users")
+        self.assertEqual(result.schema, "sales")
+
+    def test_get_table_details_kusto_namespace_parses_slash_name(self):
+        """Azure Kusto events use azurekusto namespace + {db}/{table} name."""
+        data = {
+            "namespace": "azurekusto://mycluster.kusto.windows.net",
+            "name": "mydb/mytable",
+        }
+        result = OpenlineageSource._get_table_details(data)
+        self.assertEqual(result.name, "mytable")
+        self.assertEqual(result.schema, "mydb")
+
+    def test_get_table_details_cosmos_namespace_parses_colls_name(self):
+        """Azure Cosmos DB events carry the database in the namespace path."""
+        data = {
+            "namespace": "azurecosmos://host.documents.azure.com/dbs/mydb",
+            "name": "colls/orders",
+        }
+        result = OpenlineageSource._get_table_details(data)
+        self.assertEqual(result.name, "orders")
+        self.assertEqual(result.schema, "mydb")
+
+    def test_get_entity_details_glue_namespace_resolves_to_table(self):
+        """Glue ARN namespace + table/{db}/{table} name resolves to a table entity."""
+        data = {
+            "namespace": "arn:aws:glue:us-east-1:123456789012",
+            "name": "table/sales/users",
+            "facets": {},
+        }
+        result = OpenlineageSource._get_entity_details(data)
+        self.assertIsNotNone(result)
+        self.assertEqual(result.entity_type, "table")
+        self.assertEqual(result.table_details.name, "users")
+        self.assertEqual(result.table_details.schema, "sales")
+
+    def test_get_entity_details_unparseable_name_raises_value_error(self):
+        """Unrecognised name formats raise ValueError so callers can surface the error."""
+        data = {"namespace": "trino://host:8080", "name": "invalidname"}
+        with self.assertRaises(ValueError):
+            OpenlineageSource._get_entity_details(data)
+
 
 if __name__ == "__main__":
     unittest.main()
