fix(ingestion/tableau): restructure the tableau graphql datasource query

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

@@ -0,0 +1,6 @@
+from datahub.ingestion.source.tableau.tableau import (
+    TableauConfig,
+    TableauSiteSource,
+    TableauSource,
+    TableauSourceReport,
+)

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

@@ -69,7 +69,6 @@
     TestConnectionReport,
 )
 from datahub.ingestion.api.workunit import MetadataWorkUnit
-from datahub.ingestion.source import tableau_constant as c
 from datahub.ingestion.source.common.subtypes import (
     BIContainerSubTypes,
     DatasetSubTypes,
@@ -83,7 +82,8 @@
     StatefulIngestionConfigBase,
     StatefulIngestionSourceBase,
 )
-from datahub.ingestion.source.tableau_common import (
+from datahub.ingestion.source.tableau import tableau_constant as c
+from datahub.ingestion.source.tableau.tableau_common import (
     FIELD_TYPE_MAPPING,
     MetadataQueryException,
     TableauLineageOverrides,
@@ -93,6 +93,7 @@
     dashboard_graphql_query,
     database_servers_graphql_query,
     database_tables_graphql_query,
+    datasource_upstream_fields_graphql_query,
     embedded_datasource_graphql_query,
     get_filter_pages,
     get_overridden_info,
@@ -101,7 +102,7 @@
     make_fine_grained_lineage_class,
     make_upstream_class,
     published_datasource_graphql_query,
-    query_metadata,
+    query_metadata_cursor_based_pagination,
     sheet_graphql_query,
     tableau_field_to_schema_field,
     workbook_graphql_query,
@@ -345,6 +346,12 @@ class TableauConfig(
         default=10,
         description="[advanced] Number of metadata objects (e.g. CustomSQLTable, PublishedDatasource, etc) to query at a time using the Tableau API.",
     )
+
+    fetch_size: int = Field(
+        default=250,
+        description="Specifies the number of records to retrieve in each batch during a query execution.",
+    )
+
     # We've found that even with a small workbook page size (e.g. 10), the Tableau API often
     # returns warnings like this:
     # {
@@ -972,22 +979,28 @@ def get_connection_object_page(
         query: str,
         connection_type: str,
         query_filter: str,
-        count: int = 0,
-        offset: int = 0,
+        current_cursor: Optional[str],  # initial value is None
+        fetch_size: int = 0,
         retry_on_auth_error: bool = True,
         retries_remaining: Optional[int] = None,
-    ) -> Tuple[dict, int, int]:
+    ) -> Tuple[dict, Optional[str], int]:
         retries_remaining = retries_remaining or self.config.max_retries
 
         logger.debug(
-            f"Query {connection_type} to get {count} objects with offset {offset}"
+            f"Query {connection_type} to get {fetch_size} objects with offset {current_cursor}"
             f" and filter {query_filter}"
         )
         try:
             assert self.server is not None
-            query_data = query_metadata(
-                self.server, query, connection_type, count, offset, query_filter
+            query_data = query_metadata_cursor_based_pagination(
+                server=self.server,
+                main_query=query,
+                connection_name=connection_type,
+                first=fetch_size,
+                after=current_cursor,
+                qry_filter=query_filter,
             )
+
         except REAUTHENTICATE_ERRORS:
             if not retry_on_auth_error:
                 raise
@@ -997,11 +1010,11 @@ def get_connection_object_page(
             # will be thrown and we need to re-authenticate and retry.
             self._re_authenticate()
             return self.get_connection_object_page(
-                query,
-                connection_type,
-                query_filter,
-                count,
-                offset,
+                query=query,
+                connection_type=connection_type,
+                query_filter=query_filter,
+                fetch_size=fetch_size,
+                current_cursor=current_cursor,
                 retry_on_auth_error=False,
                 retries_remaining=retries_remaining,
             )
@@ -1015,13 +1028,13 @@ def get_connection_object_page(
             if retries_remaining <= 0:
                 raise
             return self.get_connection_object_page(
-                query,
-                connection_type,
-                query_filter,
-                count,
-                offset,
+                query=query,
+                connection_type=connection_type,
+                query_filter=query_filter,
+                fetch_size=fetch_size,
+                current_cursor=current_cursor,
                 retry_on_auth_error=False,
-                retries_remaining=retries_remaining - 1,
+                retries_remaining=retries_remaining,
             )
 
         if c.ERRORS in query_data:
@@ -1035,13 +1048,19 @@ def get_connection_object_page(
                 # filter out PERMISSIONS_MODE_SWITCHED to report error in human-readable format
                 other_errors = []
                 permission_mode_errors = []
+                node_limit_errors = []
                 for error in errors:
                     if (
                         error.get("extensions")
                         and error["extensions"].get("code")
                         == "PERMISSIONS_MODE_SWITCHED"
                     ):
                         permission_mode_errors.append(error)
+                    elif (
+                        error.get("extensions")
+                        and error["extensions"].get("code") == "NODE_LIMIT_EXCEEDED"
+                    ):
+                        node_limit_errors.append(error)
                     else:
                         other_errors.append(error)
 
@@ -1060,6 +1079,18 @@ def get_connection_object_page(
                         context=f"{permission_mode_errors}",
                     )
 
+                if node_limit_errors:
+                    logger.debug(f"Node Limit Error. query_data {query_data}")
+                    self.report.warning(
+                        title="Node Limit Errors",
+                        message="Increase your tableau node limit",
+                        context=f"""{{
+                                "errors": {node_limit_errors},
+                                "connection_type": {connection_type},
+                                "query_filter": {query_filter},
+                                "query": {query},
+                        }}""",
+                    )
             else:
                 # As of Tableau Server 2024.2, the metadata API sporadically returns a 30 second
                 # timeout error. It doesn't reliably happen, so retrying a couple times makes sense.
@@ -1082,23 +1113,28 @@ def get_connection_object_page(
                     )
                     time.sleep(backoff_time)
                     return self.get_connection_object_page(
-                        query,
-                        connection_type,
-                        query_filter,
-                        count,
-                        offset,
+                        query=query,
+                        connection_type=connection_type,
+                        query_filter=query_filter,
+                        fetch_size=fetch_size,
+                        current_cursor=current_cursor,
                         retry_on_auth_error=False,
-                        retries_remaining=retries_remaining - 1,
+                        retries_remaining=retries_remaining,
                     )
                 raise RuntimeError(f"Query {connection_type} error: {errors}")
 
         connection_object = query_data.get(c.DATA, {}).get(connection_type, {})
 
-        total_count = connection_object.get(c.TOTAL_COUNT, 0)
         has_next_page = connection_object.get(c.PAGE_INFO, {}).get(
             c.HAS_NEXT_PAGE, False
         )
-        return connection_object, total_count, has_next_page
+
+        next_cursor = connection_object.get(c.PAGE_INFO, {}).get(
+            "endCursor",
+            None,
+        )
+
+        return connection_object, next_cursor, has_next_page
 
     def get_connection_objects(
         self,
@@ -1114,29 +1150,23 @@ def get_connection_objects(
         filter_pages = get_filter_pages(query_filter, page_size)
 
         for filter_page in filter_pages:
-            total_count = page_size
             has_next_page = 1
-            offset = 0
+            current_cursor: Optional[str] = None
             while has_next_page:
-                count = (
-                    page_size
-                    if offset + page_size < total_count
-                    else total_count - offset
-                )
+                filter_: str = make_filter(filter_page)
+
                 (
                     connection_objects,
-                    total_count,
+                    current_cursor,
                     has_next_page,
                 ) = self.get_connection_object_page(
-                    query,
-                    connection_type,
-                    make_filter(filter_page),
-                    count,
-                    offset,
+                    query=query,
+                    connection_type=connection_type,
+                    query_filter=filter_,
+                    fetch_size=self.config.fetch_size,
+                    current_cursor=current_cursor,
                 )
 
-                offset += count
-
                 yield from connection_objects.get(c.NODES) or []
 
     def emit_workbooks(self) -> Iterable[MetadataWorkUnit]:
@@ -2291,6 +2321,41 @@ def _get_datasource_container_key(
 
         return container_key
 
+    def update_datasource_for_field_upstream(
+        self,
+        datasource: dict,
+        field_upstream_query: str,
+    ) -> dict:
+        # Collect field ids to fetch field upstreams
+        field_ids: List[str] = []
+        for field in datasource.get(c.FIELDS, []):
+            if field.get(c.ID):
+                field_ids.append(field.get(c.ID))
+
+        # Fetch field upstreams and arrange them in map
+        field_vs_upstream: Dict[str, dict] = {}
+        for field_upstream in self.get_connection_objects(
+            field_upstream_query,
+            c.FIELDS_CONNECTION,
+            {c.ID_WITH_IN: field_ids},
+        ):
+            if field_upstream.get(c.ID):
+                field_id = field_upstream[c.ID]
+                # delete the field id, we don't need it for further processing
+                del field_upstream[c.ID]
+                field_vs_upstream[field_id] = field_upstream
+
+        # update datasource's field for its upstream
+        for field_dict in datasource.get(c.FIELDS, []):
+            field_upstream_dict: Optional[dict] = field_vs_upstream.get(
+                field_dict.get(c.ID)
+            )
+            if field_upstream_dict:
+                # Add upstream fields to field
+                field_dict.update(field_upstream_dict)
+
+        return datasource
+
     def emit_published_datasources(self) -> Iterable[MetadataWorkUnit]:
         datasource_filter = {c.ID_WITH_IN: self.datasource_ids_being_used}
 
@@ -2299,6 +2364,11 @@ def emit_published_datasources(self) -> Iterable[MetadataWorkUnit]:
             c.PUBLISHED_DATA_SOURCES_CONNECTION,
             datasource_filter,
         ):
+            datasource = self.update_datasource_for_field_upstream(
+                datasource=datasource,
+                field_upstream_query=datasource_upstream_fields_graphql_query,
+            )
+
             yield from self.emit_datasource(datasource)
 
     def emit_upstream_tables(self) -> Iterable[MetadataWorkUnit]:
@@ -2940,6 +3010,10 @@ def emit_embedded_datasources(self) -> Iterable[MetadataWorkUnit]:
             c.EMBEDDED_DATA_SOURCES_CONNECTION,
             datasource_filter,
         ):
+            datasource = self.update_datasource_for_field_upstream(
+                datasource=datasource,
+                field_upstream_query=datasource_upstream_fields_graphql_query,
+            )
             yield from self.emit_datasource(
                 datasource,
                 datasource.get(c.WORKBOOK),