fix(ingest/tableau): honor the key projectNameWithin in pagination (datahub-project#12107)

Author: sid-acryl

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

@@ -1290,7 +1290,6 @@ def get_connection_objects(
         page_size = page_size_override or self.config.page_size
 
         filter_pages = get_filter_pages(query_filter, page_size)
-
         for filter_page in filter_pages:
             has_next_page = 1
             current_cursor: Optional[str] = None

metadata-ingestion/src/datahub/ingestion/source/tableau/tableau_common.py

@@ -975,15 +975,22 @@ def get_filter_pages(query_filter: dict, page_size: int) -> List[dict]:
     # a few ten thousand, then tableau server responds with empty response
     # causing below error:
     # tableauserverclient.server.endpoint.exceptions.NonXMLResponseError: b''
+
+    # in practice, we only do pagination if len(query_filter.keys()) == 1
+    if len(query_filter.keys()) != 1:
+        return filter_pages
+
+    current_key = (list(query_filter.keys()))[0]
+
     if (
-        len(query_filter.keys()) == 1
-        and query_filter.get(c.ID_WITH_IN)
-        and isinstance(query_filter[c.ID_WITH_IN], list)
+        current_key in [c.ID_WITH_IN, c.PROJECT_NAME_WITH_IN]
+        and query_filter.get(current_key)
+        and isinstance(query_filter[current_key], list)
     ):
-        ids = query_filter[c.ID_WITH_IN]
+        ids = query_filter[current_key]
         filter_pages = [
             {
-                c.ID_WITH_IN: ids[
+                current_key: ids[
                     start : (
                         start + page_size if start + page_size < len(ids) else len(ids)
                     )

metadata-ingestion/tests/integration/tableau/test_tableau_ingest.py

@@ -61,6 +61,7 @@
     "projects": ["default", "Project 2", "Samples"],
     "extract_project_hierarchy": False,
     "page_size": 1000,
+    "workbook_page_size": 1000,
     "ingest_tags": True,
     "ingest_owner": True,
     "ingest_tables_external": True,
@@ -674,6 +675,7 @@ def test_tableau_ingest_with_platform_instance(
         "platform_instance": "acryl_site1",
         "projects": ["default", "Project 2"],
         "page_size": 1000,
+        "workbook_page_size": 1000,
         "ingest_tags": True,
         "ingest_owner": True,
         "ingest_tables_external": True,

metadata-ingestion/tests/unit/test_tableau_source.py

@@ -182,8 +182,14 @@ def test_get_filter_pages_simple():
     assert get_filter_pages(filter_dict, 10) == [filter_dict]
 
 
-def test_get_filter_pages_non_id_large_filter_passthrough():
-    projects = [f"project{i}" for i in range(20000)]
+def test_get_filter_pages_non_id_large_filter():
+    projects = [f"project{i}" for i in range(10)]
+    filter_dict = {c.PROJECT_NAME_WITH_IN: projects}
+    assert get_filter_pages(filter_dict, 10) == [filter_dict]
+
+
+def test_get_filter_pages_for_single_key():
+    projects = ["project1"]
     filter_dict = {c.PROJECT_NAME_WITH_IN: projects}
     assert get_filter_pages(filter_dict, 10) == [filter_dict]