fix(ingestion): improve robustness of glue ingestion source

fixes: #2625

metadata-ingestion/src/datahub/ingestion/source/glue.py

@@ -140,38 +140,32 @@ def create(cls, config_dict, ctx):
 
     def get_workunits(self) -> Iterable[MetadataWorkUnit]:
         def get_all_tables() -> List[dict]:
-            def get_tables_from_database(
-                database_name: str, tables: List
-            ) -> List[dict]:
-                kwargs = {"DatabaseName": database_name}
-                while True:
-                    data = self.glue_client.get_tables(**kwargs)
-                    tables += data["TableList"]
-                    if "NextToken" in data:
-                        kwargs["NextToken"] = data["NextToken"]
-                    else:
-                        break
-                return tables
-
-            def get_tables_from_all_databases() -> List[dict]:
-                tables = []
-                kwargs: Dict = {}
-                while True:
-                    data = self.glue_client.search_tables(**kwargs)
-                    tables += data["TableList"]
-                    if "NextToken" in data:
-                        kwargs["NextToken"] = data["NextToken"]
-                    else:
-                        break
-                return tables
+            def get_tables_from_database(database_name: str) -> List[dict]:
+                new_tables = []
+                paginator = self.glue_client.get_paginator("get_tables")
+                for page in paginator.paginate(DatabaseName=database_name):
+                    new_tables += page["TableList"]
+
+                return new_tables
+
+            def get_database_names() -> List[str]:
+                database_names = []
+                paginator = self.glue_client.get_paginator("get_databases")
+                for page in paginator.paginate():
+                    for db in page["DatabaseList"]:
+                        if self.source_config.database_pattern.allowed(db["Name"]):
+                            database_names.append(db["Name"])
+
+                return database_names
 
             if self.source_config.database_pattern.is_fully_specified_allow_list():
-                all_tables: List[dict] = []
                 database_names = self.source_config.database_pattern.get_allowed_list()
-                for database in database_names:
-                    all_tables += get_tables_from_database(database, all_tables)
             else:
-                all_tables = get_tables_from_all_databases()
+                database_names = get_database_names()
+
+            all_tables: List[dict] = []
+            for database in database_names:
+                all_tables += get_tables_from_database(database)
             return all_tables
 
         tables = get_all_tables()

metadata-ingestion/tests/unit/test_glue_source.py

@@ -89,7 +89,20 @@ def test_turn_boto_glue_data_to_metadata_event(self):
         stringy_timestamp = datetime.strptime(FROZEN_TIME, "%Y-%m-%d %H:%M:%S")
         timestamp = int(datetime.timestamp(stringy_timestamp) * 1000)
 
-        response = {
+        get_databases_response = {
+            "DatabaseList": [
+                {
+                    "Name": "datalake_grilled",
+                    "Description": "irrelevant",
+                    "LocationUri": "irrelevant",
+                    "Parameters": {},
+                    "CreateTime": datetime(2015, 1, 1),
+                    "CreateTableDefaultPermissions": [],
+                    "CatalogId": "irrelevant",
+                },
+            ],
+        }
+        get_tables_response = {
             "TableList": [
                 {
                     "Name": "Barbeque",
@@ -110,7 +123,10 @@ def test_turn_boto_glue_data_to_metadata_event(self):
         }
 
         with Stubber(self.glue_source.glue_client) as stubber:
-            stubber.add_response("search_tables", response, {})
+            stubber.add_response("get_databases", get_databases_response, {})
+            stubber.add_response(
+                "get_tables", get_tables_response, {"DatabaseName": "datalake_grilled"}
+            )
             actual_work_unit = list(self.glue_source.get_workunits())[0]
 
         expected_metadata_work_unit = create_metadata_work_unit(timestamp)