Add ChainRetryPolicy to common.compat for Airflow 3.3

providers/common/compat/provider.yaml

@@ -66,3 +66,8 @@ integrations:
   - integration-name: Common Compat
     external-doc-url: https://airflow.apache.org/docs/apache-airflow-providers-common-compat/
     tags: [software]
+
+retry-policies:
+  - integration-name: Common Compat
+    python-modules:
+      - airflow.providers.common.compat._retry_policy

providers/common/compat/src/airflow/providers/common/compat/_retry_policy.py

@@ -0,0 +1,127 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""
+Backport of ``airflow.sdk.definitions.retry_policy.ChainRetryPolicy`` for Airflow 3.3.
+
+Retry policies arrived in Airflow 3.3 and ``ChainRetryPolicy`` in 3.4. This copy lets a Dag that
+imports it from :mod:`airflow.providers.common.compat.sdk` run on 3.3 with the same behaviour. That
+module tries the SDK first, so wherever the SDK has the class (3.4 and later) it hands out the SDK
+class and this one is never imported.
+
+The worker only calls ``evaluate`` on whatever ``retry_policy`` holds, so a policy built from this
+class behaves the same as one built from the SDK class. Keep the logic here in step with the SDK;
+the compat tests compare the source of both classes and fail on drift.
+"""
+
+from __future__ import annotations
+
+import logging
+from typing import TYPE_CHECKING
+
+from airflow.sdk.definitions.retry_policy import RetryAction, RetryDecision, RetryPolicy
+
+if TYPE_CHECKING:
+    from collections.abc import Sequence
+
+    from airflow.sdk.definitions.context import Context
+
+log = logging.getLogger(__name__)
+
+__all__ = ["ChainRetryPolicy"]
+
+
+class ChainRetryPolicy(RetryPolicy):
+    """
+    Consult policies in order; the first RETRY or FAIL wins.
+
+    Every policy sees the original task exception. A policy that returns DEFAULT has nothing
+    to add, whatever reason it attached, and the next one is asked. When every policy returns
+    DEFAULT the chain does too, and the task's own ``retries`` and ``retry_delay`` apply.
+
+    This makes the order of a fallback explicit and lets any policies compose, including ones
+    from different packages::
+
+        retry_policy = ChainRetryPolicy(
+            [
+                # Known failures first, so no model call is spent on them.
+                ExceptionRetryPolicy(rules=[RetryRule(exception=PermissionError, action=RetryAction.FAIL)]),
+                # A policy that consults a model (airflow.providers.common.ai.policies.retry).
+                LLMRetryPolicy(llm_conn_id="pydanticai_default"),
+                # The floor when the model is unreachable.
+                ExceptionRetryPolicy(
+                    rules=[RetryRule(exception=ConnectionError, retry_delay=timedelta(seconds=30))]
+                ),
+            ]
+        )
+
+    A policy that raises an ordinary exception, or returns something other than a
+    :class:`~airflow.sdk.definitions.retry_policy.RetryDecision`, is logged and treated as DEFAULT.
+    ``BaseException`` propagates.
+
+    The winning decision's reason names the policy that decided, then what every earlier policy
+    said: ``HTTPStatusRetryPolicy: HTTP 404 (after ExceptionRetryPolicy: no decision)``. The
+    worker stores it as ``retry_reason`` on a RETRY and logs it otherwise.
+
+    :param policies: The policies to consult, in order. At least one.
+    """
+
+    def __init__(self, policies: Sequence[RetryPolicy]) -> None:
+        if isinstance(policies, RetryPolicy):
+            raise TypeError(
+                "ChainRetryPolicy takes a sequence of policies; wrap the single policy in a list."
+            )
+        policies = list(policies)
+        if not policies:
+            raise ValueError("ChainRetryPolicy needs at least one policy.")
+        for position, policy in enumerate(policies):
+            if not isinstance(policy, RetryPolicy):
+                raise TypeError(
+                    f"ChainRetryPolicy policies[{position}] must be a RetryPolicy, got {type(policy).__name__}."
+                )
+        self.policies = policies
+
+    def evaluate(
+        self,
+        exception: BaseException,
+        try_number: int,
+        max_tries: int,
+        context: Context | None = None,
+    ) -> RetryDecision:
+        trail: list[str] = []
+        for policy in self.policies:
+            name = type(policy).__name__
+            try:
+                decision = policy.evaluate(
+                    exception=exception, try_number=try_number, max_tries=max_tries, context=context
+                )
+            except Exception as exc:
+                log.exception("%s raised while evaluating the retry policy; treated as no decision", name)
+                trail.append(f"{name}: raised {type(exc).__name__}")
+                continue
+            if not isinstance(decision, RetryDecision):
+                log.error("%s returned %r instead of a RetryDecision; treated as no decision", name, decision)
+                trail.append(f"{name}: invalid decision")
+                continue
+            if decision.action is RetryAction.DEFAULT:
+                trail.append(f"{name}: {decision.reason or 'no decision'}")
+                continue
+            reason = f"{name}: {decision.reason or decision.action.value}"
+            if trail:
+                reason = f"{reason} (after {'; '.join(trail)})"
+            return RetryDecision(action=decision.action, retry_delay=decision.retry_delay, reason=reason)
+        # The worker logs this reason as the policy decision; it is not stored, since nothing is retried by it.
+        return RetryDecision(action=RetryAction.DEFAULT, reason=f"no policy decided ({'; '.join(trail)})")

providers/common/compat/src/airflow/providers/common/compat/sdk.py

@@ -25,7 +25,7 @@
 
 from typing import TYPE_CHECKING
 
-from airflow.providers.common.compat.version_compat import AIRFLOW_V_3_0_PLUS
+from airflow.providers.common.compat.version_compat import AIRFLOW_V_3_0_PLUS, AIRFLOW_V_3_3_PLUS
 
 if TYPE_CHECKING:
     import airflow.sdk.io as io  # noqa: F401
@@ -116,6 +116,16 @@
             DagRunTriggerException as DagRunTriggerException,
             DownstreamTasksSkipped as DownstreamTasksSkipped,
         )
+    # Retry policies exist from Airflow 3.3 (conditionally imported)
+    if AIRFLOW_V_3_3_PLUS:
+        from airflow.sdk.definitions.retry_policy import (
+            ChainRetryPolicy as ChainRetryPolicy,
+            ExceptionRetryPolicy as ExceptionRetryPolicy,
+            RetryAction as RetryAction,
+            RetryDecision as RetryDecision,
+            RetryPolicy as RetryPolicy,
+            RetryRule as RetryRule,
+        )
     from airflow.sdk.execution_time.context import (
         AIRFLOW_VAR_NAME_FORMAT_MAPPING as AIRFLOW_VAR_NAME_FORMAT_MAPPING,
         context_to_airflow_vars as context_to_airflow_vars,
@@ -332,6 +342,24 @@
     # 3.2+: airflow.sdk.lineage.AssetLineageInfo
     _IMPORT_MAP["AssetLineageInfo"] = ("airflow.sdk.lineage", "airflow.lineage.hook")
 
+# Retry policies arrived in Airflow 3.3 (AIP-105) and ChainRetryPolicy in 3.4. The SDK module is tried
+# first, so wherever it has the class (3.4 and later) that one is used and the copy in this provider is
+# never imported; on 3.3 the copy stands in for it.
+_AIRFLOW_3_3_ONLY_RETRY_POLICIES: dict[str, str | tuple[str, ...]] = {
+    "ChainRetryPolicy": (
+        "airflow.sdk.definitions.retry_policy",
+        "airflow.providers.common.compat._retry_policy",
+    ),
+    "ExceptionRetryPolicy": "airflow.sdk.definitions.retry_policy",
+    "RetryAction": "airflow.sdk.definitions.retry_policy",
+    "RetryDecision": "airflow.sdk.definitions.retry_policy",
+    "RetryPolicy": "airflow.sdk.definitions.retry_policy",
+    "RetryRule": "airflow.sdk.definitions.retry_policy",
+}
+
+if AIRFLOW_V_3_3_PLUS:
+    _IMPORT_MAP.update(_AIRFLOW_3_3_ONLY_RETRY_POLICIES)
+
 # Module map: module_name -> module_path(s)
 # For entire modules that have been moved (e.g., timezone)
 # Usage: from airflow.providers.common.compat.sdk import timezone