fix(ingestion/glue): manage table names from resource_links from near…

…est catalog correctly (#11578)
datahub-project · Oct 23, 2024 · 326afc6 · 326afc6
1 parent 581a39a
commit 326afc6
Show file tree

Hide file tree

Showing 3 changed files with 34 additions and 35 deletions.
diff --git a/metadata-ingestion/src/datahub/ingestion/source/aws/glue.py b/metadata-ingestion/src/datahub/ingestion/source/aws/glue.py
@@ -678,12 +678,19 @@ def get_all_databases(self) -> Iterable[Mapping[str, Any]]:
         else:
             paginator_response = paginator.paginate()
 
-        for page in paginator_response:
-            yield from page["DatabaseList"]
+        pattern = "DatabaseList"
+        if self.source_config.ignore_resource_links:
+            # exclude resource links by using a JMESPath conditional query against the TargetDatabase struct key
+            pattern += "[?!TargetDatabase]"
+
+        for database in paginator_response.search(pattern):
+            if self.source_config.database_pattern.allowed(database["Name"]):
+                yield database
 
-    def get_tables_from_database(self, database_name: str) -> Iterable[Dict]:
+    def get_tables_from_database(self, database: Mapping[str, Any]) -> Iterable[Dict]:
         # see https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/glue/paginator/GetTables.html
         paginator = self.glue_client.get_paginator("get_tables")
+        database_name = database["Name"]
 
         if self.source_config.catalog_id:
             paginator_response = paginator.paginate(
@@ -692,34 +699,28 @@ def get_tables_from_database(self, database_name: str) -> Iterable[Dict]:
         else:
             paginator_response = paginator.paginate(DatabaseName=database_name)
 
-        for page in paginator_response:
-            yield from page["TableList"]
+        for table in paginator_response.search("TableList"):
+            # if resource links are detected, re-use database names from the current catalog
+            # otherwise, external names are picked up instead of aliased ones when creating full table names later
+            # This will cause an incoherent situation when creating full table names later
+            # Note: use an explicit source_config check but it is useless actually (filtering has already been done)
+            if (
+                not self.source_config.ignore_resource_links
+                and "TargetDatabase" in database
+            ):
+                table["DatabaseName"] = database["Name"]
+            yield table
 
     def get_all_databases_and_tables(
         self,
-    ) -> Tuple[Dict, List[Dict]]:
-        all_databases = self.get_all_databases()
-
-        if self.source_config.ignore_resource_links:
-            all_databases = [
-                database
-                for database in all_databases
-                if "TargetDatabase" not in database
-            ]
-
-        allowed_databases = {
-            database["Name"]: database
-            for database in all_databases
-            if self.source_config.database_pattern.allowed(database["Name"])
-        }
-
+    ) -> Tuple[List[Mapping[str, Any]], List[Dict]]:
+        all_databases = [*self.get_all_databases()]
         all_tables = [
-            table
-            for database_name in allowed_databases
-            for table in self.get_tables_from_database(database_name)
+            tables
+            for database in all_databases
+            for tables in self.get_tables_from_database(database)
         ]
-
-        return allowed_databases, all_tables
+        return all_databases, all_tables
 
     def get_lineage_if_enabled(
         self, mce: MetadataChangeEventClass
@@ -1039,7 +1040,7 @@ def get_workunit_processors(self) -> List[Optional[MetadataWorkUnitProcessor]]:
     def get_workunits_internal(self) -> Iterable[MetadataWorkUnit]:
         databases, tables = self.get_all_databases_and_tables()
 
-        for database in databases.values():
+        for database in databases:
             yield from self.gen_database_containers(database)
 
         for table in tables:

diff --git a/metadata-ingestion/tests/unit/glue/test_glue_source.py b/metadata-ingestion/tests/unit/glue/test_glue_source.py
@@ -267,8 +267,8 @@ def test_platform_config():
 @pytest.mark.parametrize(
     "ignore_resource_links, all_databases_and_tables_result",
     [
-        (True, ({}, [])),
-        (False, ({"test-database": resource_link_database}, target_database_tables)),
+        (True, ([], [])),
+        (False, ([resource_link_database], target_database_tables)),
     ],
 )
 def test_ignore_resource_links(ignore_resource_links, all_databases_and_tables_result):
@@ -289,7 +289,7 @@ def test_ignore_resource_links(ignore_resource_links, all_databases_and_tables_r
         glue_stubber.add_response(
             "get_tables",
             get_tables_response_for_target_database,
-            {"DatabaseName": "test-database"},
+            {"DatabaseName": "resource-link-test-database"},
         )
 
         assert source.get_all_databases_and_tables() == all_databases_and_tables_result

diff --git a/metadata-ingestion/tests/unit/glue/test_glue_source_stubs.py b/metadata-ingestion/tests/unit/glue/test_glue_source_stubs.py
@@ -5,7 +5,7 @@
 from botocore.response import StreamingBody
 
 resource_link_database = {
-    "Name": "test-database",
+    "Name": "resource-link-test-database",
     "CreateTime": datetime.datetime(2021, 6, 9, 14, 14, 19),
     "CreateTableDefaultPermissions": [],
     "TargetDatabase": {"CatalogId": "432143214321", "DatabaseName": "test-database"},
@@ -92,10 +92,8 @@
         },
     ]
 }
-databases_1 = {
-    "flights-database": {"Name": "flights-database", "CatalogId": "123412341234"}
-}
-databases_2 = {"test-database": {"Name": "test-database", "CatalogId": "123412341234"}}
+databases_1 = [{"Name": "flights-database", "CatalogId": "123412341234"}]
+databases_2 = [{"Name": "test-database", "CatalogId": "123412341234"}]
 tables_1 = [
     {
         "Name": "avro",