feat(ingest/tableau): add retry on timeout

metadata-ingestion/setup.py

@@ -449,7 +449,7 @@
  },
  # FIXME: I don't think tableau uses sqllineage anymore so we should be able
  # to remove that dependency.
- "tableau": {"tableauserverclient>=0.17.0"} | sqllineage_lib | sqlglot_lib,
+ "tableau": {"tableauserverclient>=0.24.0"} | sqllineage_lib | sqlglot_lib,
  "teradata": sql_common
  | usage_common
  | sqlglot_lib

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

@@ -137,7 +137,10 @@ def report_log(
  )
 
  # Add the simple exception details to the context.
- context = f"{context}: {exc}"
+ if context:
+ context = f"{context} {type(exc)}: {exc}"
+ else:
+ context = f"{type(exc)}: {exc}"
  elif log:
  logger.log(level=level.value, msg=log_content, stacklevel=stacklevel)
 

metadata-ingestion/src/datahub/ingestion/source/state/stale_entity_removal_handler.py

@@ -298,7 +298,7 @@ def gen_removed_entity_workunits(self) -> Iterable[MetadataWorkUnit]:
 
  if copy_previous_state_and_fail:
  logger.info(
- f"Copying urns from last state (size {last_checkpoint_state.urns}) to current state (size {cur_checkpoint_state.urns}) "
+ f"Copying urns from last state (size {len(last_checkpoint_state.urns)}) to current state (size {len(cur_checkpoint_state.urns)}) "
  "to ensure stale entities from previous runs are deleted on the next successful run."
  )
  for urn in last_checkpoint_state.urns:

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

@@ -1,5 +1,6 @@
 import logging
 import re
+import time
 from collections import OrderedDict
 from dataclasses import dataclass
 from datetime import datetime
@@ -13,6 +14,7 @@
  Optional,
  Set,
  Tuple,
+ Type,
  Union,
  cast,
 )
@@ -158,6 +160,19 @@
 from datahub.utilities import config_clean
 from datahub.utilities.urns.dataset_urn import DatasetUrn
 
+try:
+ # On earlier versions of the tableauserverclient, the NonXMLResponseError
+ # was thrown when reauthentication was needed. We'll keep both exceptions
+ # around for now, but can remove this in the future.
+ from tableauserverclient.server.endpoint.exceptions import NotSignedInError
+
+ REAUTHENTICATE_ERRORS: Tuple[Type[Exception], ...] = (
+ NotSignedInError,
+ NonXMLResponseError,
+ )
+except ImportError:
+ REAUTHENTICATE_ERRORS = (NonXMLResponseError,)
+
 logger: logging.Logger = logging.getLogger(__name__)
 
 # Replace / with |
@@ -965,7 +980,7 @@ def get_connection_object_page(
  query_data = query_metadata(
  self.server, query, connection_type, count, offset, query_filter
  )
- except NonXMLResponseError:
+ except REAUTHENTICATE_ERRORS:
  if not retry_on_auth_error:
  raise
 
@@ -1038,6 +1053,35 @@ def get_connection_object_page(
  )
 
  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.
+ if all(
+ error.get("message")
+ == "Execution canceled because timeout of 30000 millis was reached"
+ for error in errors
+ ):
+ # If it was only a timeout error, we can retry.
+ if retries_remaining <= 0:
+ raise
+
+ # This is a pretty dumb backoff mechanism, but it's good enough for now.
+ backoff_time = min(
+ (self.config.max_retries - retries_remaining + 1) ** 2, 60
+ )
+ logger.info(
+ f"Query {connection_type} received a 30 second timeout error - will retry in {backoff_time} seconds. "
+ f"Retries remaining: {retries_remaining}"
+ )
+ time.sleep(backoff_time)
+ return self.get_connection_object_page(
+ query,
+ connection_type,
+ query_filter,
+ count,
+ offset,
+ retry_on_auth_error=False,
+ retries_remaining=retries_remaining - 1,
+ )
  raise RuntimeError(f"Query {connection_type} error: {errors}")
 
  connection_object = query_data.get(c.DATA, {}).get(connection_type, {})