feat: add automatic retry for transient dbt command errors

elementary/clients/dbt/command_line_dbt_runner.py

@@ -5,16 +5,38 @@
 from typing import Any, Dict, List, Optional
 
 import yaml
+from tenacity import (
+    RetryCallState,
+    retry,
+    retry_if_exception,
+    stop_after_attempt,
+    wait_exponential,
+)
 
 from elementary.clients.dbt.base_dbt_runner import BaseDbtRunner
 from elementary.clients.dbt.dbt_log import parse_dbt_output
+from elementary.clients.dbt.transient_errors import is_transient_error
 from elementary.exceptions.exceptions import DbtCommandError, DbtLsCommandError
 from elementary.monitor.dbt_project_utils import is_dbt_package_up_to_date
 from elementary.utils.env_vars import is_debug
 from elementary.utils.log import get_logger
 
 logger = get_logger(__name__)
 
+# Retry configuration for transient errors.
+_TRANSIENT_MAX_RETRIES = 3
+_TRANSIENT_WAIT_MULTIPLIER = 10  # seconds
+_TRANSIENT_WAIT_MAX = 60  # seconds
+
+
+class DbtTransientError(Exception):
+    """Raised internally to signal a transient dbt failure that should be retried."""
+
+    def __init__(self, result: "DbtCommandResult", message: str) -> None:
+        super().__init__(message)
+        self.result = result
+
+
 MACRO_RESULT_PATTERN = re.compile(
     "Elementary: --ELEMENTARY-MACRO-OUTPUT-START--(.*)--ELEMENTARY-MACRO-OUTPUT-END--"
 )
@@ -112,23 +134,140 @@ def _run_command(
         else:
             logger.debug(log_msg)
 
-        result = self._inner_run_command(
-            dbt_command_args,
+        result = self._execute_inner_command(
+            dbt_command_args=dbt_command_args,
+            log_command_args=log_command_args,
             capture_output=capture_output,
             quiet=quiet,
             log_output=log_output,
             log_format=log_format,
         )
 
-        if capture_output and result.output:
-            logger.debug(
-                f"Result bytes size for command '{log_command_args}' is {len(result.output)}"
+        return result
+
+    def _execute_inner_command(
+        self,
+        dbt_command_args: List[str],
+        log_command_args: List[str],
+        capture_output: bool,
+        quiet: bool,
+        log_output: bool,
+        log_format: str,
+    ) -> DbtCommandResult:
+        """Execute ``_inner_run_command`` with automatic retries for transient errors.
+
+        This method wraps the actual command execution, checks the result
+        for known transient error patterns (per adapter), and retries using
+        ``tenacity`` when appropriate.  Non-transient failures are returned
+        immediately without retrying.
+        """
+
+        def _before_retry(retry_state: RetryCallState) -> None:
+            attempt = retry_state.attempt_number
+            logger.warning(
+                "Transient error detected for dbt command '%s' "
+                "(attempt %d/%d). Retrying...",
+                " ".join(log_command_args),
+                attempt,
+                _TRANSIENT_MAX_RETRIES,
             )
-            if log_output or is_debug():
-                for log in parse_dbt_output(result.output, log_format):
-                    logger.info(log.msg)
 
-        return result
+        @retry(
+            retry=retry_if_exception(lambda exc: isinstance(exc, DbtTransientError)),
+            stop=stop_after_attempt(_TRANSIENT_MAX_RETRIES),
+            wait=wait_exponential(
+                multiplier=_TRANSIENT_WAIT_MULTIPLIER, max=_TRANSIENT_WAIT_MAX
+            ),
+            before_sleep=_before_retry,
+            reraise=True,
+        )
+        def _attempt() -> DbtCommandResult:
+            # Always capture output so transient-error detection can inspect
+            # stdout/stderr.  The original ``capture_output`` flag is still
+            # honoured for logging behaviour (see below).
+            try:
+                result = self._inner_run_command(
+                    dbt_command_args,
+                    capture_output=True,
+                    quiet=quiet,
+                    log_output=log_output,
+                    log_format=log_format,
+                )
+            except DbtCommandError as exc:
+                # DbtCommandError is raised when raise_on_failure=True and
+                # the command exits with a non-zero return code.  Extract
+                # actual output/stderr from the underlying process error
+                # for more accurate transient-error detection.
+                output_text = str(exc)
+                stderr_text: Optional[str] = None
+                if exc.proc_err is not None:
+                    if exc.proc_err.output:
+                        output_text = (
+                            exc.proc_err.output.decode()
+                            if isinstance(exc.proc_err.output, bytes)
+                            else str(exc.proc_err.output)
+                        )
+                    if exc.proc_err.stderr:
+                        stderr_text = (
+                            exc.proc_err.stderr.decode()
+                            if isinstance(exc.proc_err.stderr, bytes)
+                            else str(exc.proc_err.stderr)
+                        )
+                if is_transient_error(
+                    self.target, output=output_text, stderr=stderr_text
+                ):
+                    raise DbtTransientError(
+                        result=DbtCommandResult(
+                            success=False,
+                            output=output_text,
+                            stderr=stderr_text,
+                        ),
+                        message=(f"Transient error during dbt command: {exc}"),
+                    ) from exc
+                raise
+
+            if capture_output and result.output:
+                logger.debug(
+                    "Result bytes size for command '%s' is %d",
+                    log_command_args,
+                    len(result.output),
+                )
+                if log_output or is_debug():
+                    for log in parse_dbt_output(result.output, log_format):
+                        logger.info(log.msg)
+
+            # Command completed but may have failed (raise_on_failure=False).
+            if not result.success and is_transient_error(
+                self.target, output=result.output, stderr=result.stderr
+            ):
+                raise DbtTransientError(
+                    result=result,
+                    message=(
+                        f"Transient error during dbt command: "
+                        f"{' '.join(log_command_args)}"
+                    ),
+                )
+
+            return result
+
+        try:
+            return _attempt()
+        except DbtTransientError as exc:
+            # All retry attempts exhausted.
+            logger.exception(
+                "dbt command '%s' failed after %d attempts due to " "transient errors.",
+                " ".join(log_command_args),
+                _TRANSIENT_MAX_RETRIES,
+            )
+            # Preserve the raise_on_failure contract: if the original
+            # failure was a DbtCommandError (i.e. raise_on_failure=True),
+            # re-raise it so callers relying on exception handling still
+            # see the expected exception type.
+            if isinstance(exc.__cause__, DbtCommandError):
+                raise exc.__cause__ from exc
+            # Otherwise (raise_on_failure=False path), return the last
+            # failed result so callers that check result.success work.
+            return exc.result
 
     def deps(self, quiet: bool = False, capture_output: bool = True) -> bool:
         result = self._run_command(

elementary/clients/dbt/transient_errors.py

@@ -0,0 +1,159 @@
+"""Per-adapter transient error patterns for automatic retry.
+
+Each adapter may produce transient errors that are safe to retry.  This
+module centralises those patterns so that the runner can decide whether a
+failed dbt command should be retried transparently.
+
+To add patterns for a new adapter, append a new entry to
+``_ADAPTER_PATTERNS`` with the **adapter type** as key (e.g.
+``"bigquery"``, ``"snowflake"``) and a tuple of **plain, lowercase**
+substrings that appear in the error output.  Matching is
+case-insensitive substring search so regex is not needed.
+
+Note: The ``target`` argument accepted by :func:`is_transient_error` may
+be either the dbt adapter type *or* the profile target name (e.g.
+``"dev"``, ``"prod"``).  When it does not match any known adapter key,
+**all** adapter patterns are checked defensively.  This is safe because
+adapter-specific error messages only appear in output from that adapter.
+"""
+
+from typing import Dict, Optional, Sequence, Tuple
+
+# ---------------------------------------------------------------------------
+# Per-adapter transient error substrings (all lowercase).
+#
+# A command failure is considered *transient* when the dbt output
+# (stdout + stderr, lowercased) contains **any** of the substrings
+# listed for the active adapter **or** in the ``_COMMON`` list.
+# ---------------------------------------------------------------------------
+
+_COMMON: Tuple[str, ...] = (
+    # Generic connection / HTTP errors that any adapter can surface.
+    "connection reset by peer",
+    "connection was closed",
+    "remotedisconnected",
+    "connectionerror",
+    "brokenpipeerror",
+    "connection aborted",
+    "read timed out",
+)
+
+_DATABRICKS_PATTERNS: Tuple[str, ...] = (
+    "temporarily_unavailable",
+    "504 gateway timeout",
+    "502 bad gateway",
+    "service unavailable",
+)
+
+_ADAPTER_PATTERNS: Dict[str, Tuple[str, ...]] = {
+    "bigquery": (
+        # Streaming-buffer delay after a streaming insert.
+        "streaming data from",
+        "is temporarily unavailable",
+        # Generic transient backend error (500).
+        "retrying may solve the problem",
+        "backenderror",
+        # Rate-limit / quota errors.
+        "exceeded rate limits",
+        "rateLimitExceeded".lower(),
+        "quota exceeded",
+        # Internal errors surfaced as 503 / "internal error".
+        "internal error encountered",
+        "503 service unavailable",
+        "http 503",
+    ),
+    "snowflake": (
+        "could not connect to snowflake backend",
+        "authentication token has expired",
+        "incident id:",
+        "service is unavailable",
+    ),
+    "redshift": (
+        "connection timed out",
+        "could not connect to the server",
+        "ssl syscall error",
+    ),
+    "databricks": _DATABRICKS_PATTERNS,
+    "databricks_catalog": _DATABRICKS_PATTERNS,
+    "athena": (
+        "throttlingexception",
+        "toomanyrequestsexception",
+        "service unavailable",
+    ),
+    "dremio": (
+        # Common patterns (remotedisconnected, connection was closed) already
+        # cover the most frequent Dremio transient errors.  Add Dremio-specific
+        # patterns here as they are identified.
+    ),
+    "postgres": (
+        "could not connect to server",
+        "connection timed out",
+        "server closed the connection unexpectedly",
+        "ssl syscall error",
+    ),
+    "trino": (
+        "service unavailable",
+        "server returned http response code: 503",
+    ),
+    "clickhouse": (
+        "connection timed out",
+        "broken pipe",
+    ),
+}
+
+# Pre-computed union of all adapter-specific patterns for the unknown-target
+# fallback path.  Built once at import time to avoid repeated iteration.
+_ALL_ADAPTER_PATTERNS: Tuple[str, ...] = tuple(
+    pattern for patterns in _ADAPTER_PATTERNS.values() for pattern in patterns
+)
+
+
+def is_transient_error(
+    target: Optional[str],
+    output: Optional[str] = None,
+    stderr: Optional[str] = None,
+) -> bool:
+    """Return ``True`` if *output*/*stderr* contain a known transient error.
+
+    Parameters
+    ----------
+    target:
+        The dbt adapter type (e.g. ``"bigquery"``, ``"snowflake"``) **or**
+        the dbt profile target name (e.g. ``"dev"``, ``"prod"``).
+        When the value matches a key in ``_ADAPTER_PATTERNS``, only that
+        adapter's patterns (plus ``_COMMON``) are used.  When it does
+        **not** match any known adapter, **all** adapter patterns are
+        checked defensively to avoid missing transient errors.
+        When ``None`` only the common patterns are checked.
+    output:
+        The captured stdout of the dbt command (may be ``None``).
+    stderr:
+        The captured stderr of the dbt command (may be ``None``).
+    """
+    haystack = _build_haystack(output, stderr)
+    if not haystack:
+        return False
+
+    if isinstance(target, str):
+        adapter_patterns = _ADAPTER_PATTERNS.get(target.lower())
+        if adapter_patterns is not None:
+            # Known adapter — use common + adapter-specific patterns.
+            patterns: Sequence[str] = (*_COMMON, *adapter_patterns)
+        else:
+            # Unknown target key (e.g. profile target name). Check all adapters.
+            patterns = (*_COMMON, *_ALL_ADAPTER_PATTERNS)
+    else:
+        # No target provided; still check all adapters defensively.
+        patterns = (*_COMMON, *_ALL_ADAPTER_PATTERNS)
+
+    return any(pattern in haystack for pattern in patterns)
+
+
+def _build_haystack(output: Optional[str] = None, stderr: Optional[str] = None) -> str:
+    """Concatenate and lowercase *output* + *stderr* for matching."""
+    parts = []
+    if output and isinstance(output, str):
+        parts.append(output)
+    if stderr and isinstance(stderr, str):
+        parts.append(stderr)
+    return "\n".join(parts).lower()

pyproject.toml

@@ -42,6 +42,7 @@ packaging = ">=20.9"
 azure-storage-blob = ">=12.11.0"
 pymsteams = ">=0.2.2,<1.0.0"
 tabulate = ">= 0.9.0"
+tenacity = ">=8.0,<10.0"
 pytz = ">= 2025.1"
 
 dbt-snowflake = {version = ">=0.20,<2.0.0", optional = true}