This is an automated email from the ASF dual-hosted git repository. dabla pushed a commit to branch dabla/durable-operators/01-common-compat-resumable-job-mixin in repository https://gitbox.apache.org/repos/asf/airflow.git
commit e87496114052879b26e60f1c5983f4031c222a47 Author: David Blain <[email protected]> AuthorDate: Wed Oct 14 01:03:04 2026 +0200 Use one ResumableJobMixin fallback from common.compat in providers Seven providers each carried the same fallback for Airflow versions without ResumableJobMixin, and every operator made durable next would add another. One shared stand-in keeps the fallback and its tests in one place, and lets a provider make an operator durable without writing its own compatibility code. --- providers/amazon/pyproject.toml | 2 +- .../airflow/providers/amazon/aws/operators/glue.py | 37 +-------- .../amazon/aws/operators/redshift_data.py | 36 +------- .../tests/unit/amazon/aws/operators/test_glue.py | 18 ---- .../amazon/aws/operators/test_redshift_data.py | 18 ---- providers/apache/livy/pyproject.toml | 2 +- .../providers/apache/livy/operators/livy.py | 40 +-------- .../tests/unit/apache/livy/operators/test_livy.py | 18 ---- providers/apache/spark/pyproject.toml | 2 +- .../apache/spark/operators/spark_submit.py | 35 +------- .../apache/spark/operators/test_spark_submit.py | 17 ---- .../common/compat/_resumable_job_mixin.py | 58 +++++++++++++ .../src/airflow/providers/common/compat/sdk.py | 8 ++ .../common/compat/test__resumable_job_mixin.py | 95 ++++++++++++++++++++++ providers/databricks/pyproject.toml | 2 +- .../providers/databricks/operators/databricks.py | 42 ++-------- .../unit/databricks/operators/test_databricks.py | 18 ---- providers/google/pyproject.toml | 2 +- .../providers/google/cloud/operators/bigquery.py | 40 ++------- .../unit/google/cloud/operators/test_bigquery.py | 18 ---- providers/snowflake/pyproject.toml | 2 +- .../providers/snowflake/operators/snowflake.py | 36 +------- .../unit/snowflake/operators/test_snowflake.py | 18 ---- 23 files changed, 186 insertions(+), 378 deletions(-) diff --git a/providers/amazon/pyproject.toml b/providers/amazon/pyproject.toml index f81c0da3a31..423f0c2fdbe 100644 --- a/providers/amazon/pyproject.toml +++ b/providers/amazon/pyproject.toml @@ -59,7 +59,7 @@ requires-python = ">=3.11" # After you modify the dependencies, and rebuild your Breeze CI image with ``breeze ci-image build`` dependencies = [ "apache-airflow>=2.11.0", - "apache-airflow-providers-common-compat>=1.18.0", + "apache-airflow-providers-common-compat>=1.18.0", # use next version "apache-airflow-providers-common-sql>=1.32.0", "apache-airflow-providers-http", # The boto3/botocore floors are raised automatically by `breeze ci upgrade` to limit `pip` backtracking diff --git a/providers/amazon/src/airflow/providers/amazon/aws/operators/glue.py b/providers/amazon/src/airflow/providers/amazon/aws/operators/glue.py index ca2ff750d30..c8b1d617e0f 100644 --- a/providers/amazon/src/airflow/providers/amazon/aws/operators/glue.py +++ b/providers/amazon/src/airflow/providers/amazon/aws/operators/glue.py @@ -43,42 +43,7 @@ from airflow.providers.common.compat.openlineage.utils.spark import ( inject_parent_job_information_into_glue_arguments, inject_transport_information_into_glue_arguments, ) -from airflow.providers.common.compat.sdk import AirflowException, conf - -_DURABLE_UNSET = object() - - -def _warn_and_disable_durable_pre_3_3(durable: Any) -> bool: - """Shared by the <3.3 compat stub: durable has no effect below 3.3, warn if it was set.""" - if durable is not _DURABLE_UNSET: - warnings.warn( - "`durable` has no effect on Airflow versions below 3.3.", - UserWarning, - stacklevel=3, - ) - return False - - -# ResumableJobMixin only exists on Airflow 3.3+; this provider still targets >=2.11. Drop this -# fallback once the provider's minimum Airflow version is >=3.3. -try: - from airflow.sdk import ResumableJobMixin -except ImportError: - - class ResumableJobMixin: # type: ignore[no-redef] - """Airflow <3.3 stub, task_state_store unavailable, always submits fresh.""" - - external_id_key: str = "glue_job_run_id" - - def __init__(self, *, durable: Any = _DURABLE_UNSET, **kwargs: Any) -> None: - super().__init__(**kwargs) - self.durable = _warn_and_disable_durable_pre_3_3(durable) - - def execute_resumable(self, context): - external_id = self.submit_job(context) - self.poll_until_complete(external_id, context) - return self.get_job_result(external_id, context) - +from airflow.providers.common.compat.sdk import AirflowException, ResumableJobMixin, conf if TYPE_CHECKING: from pydantic import JsonValue diff --git a/providers/amazon/src/airflow/providers/amazon/aws/operators/redshift_data.py b/providers/amazon/src/airflow/providers/amazon/aws/operators/redshift_data.py index f385fcc1761..8853d36958c 100644 --- a/providers/amazon/src/airflow/providers/amazon/aws/operators/redshift_data.py +++ b/providers/amazon/src/airflow/providers/amazon/aws/operators/redshift_data.py @@ -17,7 +17,6 @@ # under the License. from __future__ import annotations -import warnings from typing import TYPE_CHECKING, Any, cast import botocore.exceptions @@ -27,40 +26,7 @@ from airflow.providers.amazon.aws.operators.base_aws import AwsBaseOperator from airflow.providers.amazon.aws.triggers.redshift_data import RedshiftDataTrigger from airflow.providers.amazon.aws.utils import validate_execute_complete_event from airflow.providers.amazon.aws.utils.mixins import aws_template_fields -from airflow.providers.common.compat.sdk import AirflowException, conf - -_DURABLE_UNSET = object() - - -def _warn_and_disable_durable_pre_3_3(durable: Any) -> bool: - """Shared by the <3.3 compat stub: durable has no effect below 3.3, warn if it was set.""" - if durable is not _DURABLE_UNSET: - warnings.warn( - "`durable` has no effect on Airflow versions below 3.3.", - UserWarning, - stacklevel=3, - ) - return False - - -try: - from airflow.sdk import ResumableJobMixin -except ImportError: - - class ResumableJobMixin: # type: ignore[no-redef] - """Airflow <3.3 stub, task_state_store unavailable, always submits fresh.""" - - external_id_key: str = "redshift_statement_id" - - def __init__(self, *, durable: Any = _DURABLE_UNSET, **kwargs: Any) -> None: - super().__init__(**kwargs) - self.durable = _warn_and_disable_durable_pre_3_3(durable) - - def execute_resumable(self, context): - external_id = self.submit_job(context) - self.poll_until_complete(external_id, context) - return self.get_job_result(external_id, context) - +from airflow.providers.common.compat.sdk import AirflowException, ResumableJobMixin, conf if TYPE_CHECKING: from mypy_boto3_redshift_data.type_defs import ( diff --git a/providers/amazon/tests/unit/amazon/aws/operators/test_glue.py b/providers/amazon/tests/unit/amazon/aws/operators/test_glue.py index f1ecd61bbe8..4a1852b919c 100644 --- a/providers/amazon/tests/unit/amazon/aws/operators/test_glue.py +++ b/providers/amazon/tests/unit/amazon/aws/operators/test_glue.py @@ -17,7 +17,6 @@ from __future__ import annotations import re -import warnings from collections.abc import Generator from datetime import datetime from typing import TYPE_CHECKING @@ -36,12 +35,10 @@ from airflow.providers.amazon.aws.hooks.glue import GlueDataQualityHook, GlueJob from airflow.providers.amazon.aws.hooks.s3 import S3Hook from airflow.providers.amazon.aws.links.glue import GlueJobRunDetailsLink from airflow.providers.amazon.aws.operators.glue import ( - _DURABLE_UNSET, GlueDataQualityOperator, GlueDataQualityRuleRecommendationRunOperator, GlueDataQualityRuleSetEvaluationRunOperator, GlueJobOperator, - _warn_and_disable_durable_pre_3_3, ) from airflow.providers.common.compat.sdk import AirflowException, TaskDeferred @@ -928,21 +925,6 @@ class TestGlueJobOperatorOpenLineageInjection: assert GlueJobOperator.TASK_UUID_ARG in call_args -class TestWarnAndDisableDurableAirflowPre3_3: - def test_no_warning_when_unset(self): - with warnings.catch_warnings(record=True) as caught: - warnings.simplefilter("always") - result = _warn_and_disable_durable_pre_3_3(_DURABLE_UNSET) - assert result is False - assert caught == [] - - @pytest.mark.parametrize("value", [True, False]) - def test_warns_and_disables_when_explicitly_set(self, value): - with pytest.warns(UserWarning, match="durable.*no effect"): - result = _warn_and_disable_durable_pre_3_3(value) - assert result is False - - class TestGlueJobOperatorDeprecation: @pytest.mark.parametrize("resume_value", [True, False]) def test_warns_and_maps_to_durable_old_flag(self, resume_value): diff --git a/providers/amazon/tests/unit/amazon/aws/operators/test_redshift_data.py b/providers/amazon/tests/unit/amazon/aws/operators/test_redshift_data.py index 7831cf51021..7e9145b43d1 100644 --- a/providers/amazon/tests/unit/amazon/aws/operators/test_redshift_data.py +++ b/providers/amazon/tests/unit/amazon/aws/operators/test_redshift_data.py @@ -17,7 +17,6 @@ # under the License. from __future__ import annotations -import warnings from datetime import datetime from unittest import mock @@ -27,9 +26,7 @@ import pytest from airflow.models.dag import DAG from airflow.providers.amazon.aws.hooks.redshift_data import QueryExecutionOutput from airflow.providers.amazon.aws.operators.redshift_data import ( - _DURABLE_UNSET, RedshiftDataOperator, - _warn_and_disable_durable_pre_3_3, ) from airflow.providers.amazon.aws.triggers.redshift_data import RedshiftDataTrigger from airflow.providers.common.compat.sdk import AirflowException, TaskDeferred @@ -738,18 +735,3 @@ class TestRedshiftDataOperatorDurable: ): operator = self._make_operator() assert operator.durable is False - - -class TestWarnAndDisableDurableAirflowPre3_3: - def test_no_warning_when_unset(self): - with warnings.catch_warnings(record=True) as caught: - warnings.simplefilter("always") - result = _warn_and_disable_durable_pre_3_3(_DURABLE_UNSET) - assert result is False - assert caught == [] - - @pytest.mark.parametrize("value", [True, False]) - def test_warns_and_disables_when_explicitly_set(self, value): - with pytest.warns(UserWarning, match="durable.*no effect"): - result = _warn_and_disable_durable_pre_3_3(value) - assert result is False diff --git a/providers/apache/livy/pyproject.toml b/providers/apache/livy/pyproject.toml index 00560a8e195..f2fb5be3489 100644 --- a/providers/apache/livy/pyproject.toml +++ b/providers/apache/livy/pyproject.toml @@ -60,7 +60,7 @@ requires-python = ">=3.11" dependencies = [ "apache-airflow>=2.11.0", "apache-airflow-providers-http>=6.0.1", - "apache-airflow-providers-common-compat>=1.12.0", + "apache-airflow-providers-common-compat>=1.12.0", # use next version "aiohttp>=3.14.0", ] diff --git a/providers/apache/livy/src/airflow/providers/apache/livy/operators/livy.py b/providers/apache/livy/src/airflow/providers/apache/livy/operators/livy.py index b95b63cbd00..ed65acd7cac 100644 --- a/providers/apache/livy/src/airflow/providers/apache/livy/operators/livy.py +++ b/providers/apache/livy/src/airflow/providers/apache/livy/operators/livy.py @@ -17,7 +17,6 @@ from __future__ import annotations import time -import warnings from collections.abc import Sequence from functools import cached_property from typing import TYPE_CHECKING, Any, cast @@ -28,44 +27,7 @@ from airflow.providers.common.compat.openlineage.utils.spark import ( inject_parent_job_information_into_spark_properties, inject_transport_information_into_spark_properties, ) -from airflow.providers.common.compat.sdk import AirflowException, BaseOperator, conf - -_DURABLE_UNSET = object() - - -def _warn_and_disable_durable_pre_3_3(durable: Any) -> bool: - """Shared by the <3.3 compat stub: durable has no effect below 3.3, warn if it was set.""" - if durable is not _DURABLE_UNSET: - warnings.warn( - "`durable` has no effect on Airflow versions below 3.3.", - UserWarning, - stacklevel=3, - ) - return False - - -# ResumableJobMixin ships in airflow.sdk, which only exists on Airflow 3, while this provider -# still targets apache-airflow>=2.11. Guard the import and fall back to a stub on Airflow 2; -# drop the fallback once the provider's minimum Airflow version is >=3.0. -try: - from airflow.sdk import ResumableJobMixin -except ImportError: - - class ResumableJobMixin: # type: ignore[no-redef] - """Airflow <3.3 stub, task_state_store unavailable, always submits fresh.""" - - external_id_key: str = "livy_batch_id" - - def __init__(self, *, durable: Any = _DURABLE_UNSET, **kwargs: Any) -> None: - # Swallow ``durable`` so it doesn't reach BaseOperator; crash recovery is a no-op here. - super().__init__(**kwargs) - self.durable = _warn_and_disable_durable_pre_3_3(durable) - - def execute_resumable(self, context): - external_id = self.submit_job(context) - self.poll_until_complete(external_id, context) - return self.get_job_result(external_id, context) - +from airflow.providers.common.compat.sdk import AirflowException, BaseOperator, ResumableJobMixin, conf if TYPE_CHECKING: from pydantic import JsonValue diff --git a/providers/apache/livy/tests/unit/apache/livy/operators/test_livy.py b/providers/apache/livy/tests/unit/apache/livy/operators/test_livy.py index 7fb6d22b459..674255c0a41 100644 --- a/providers/apache/livy/tests/unit/apache/livy/operators/test_livy.py +++ b/providers/apache/livy/tests/unit/apache/livy/operators/test_livy.py @@ -17,7 +17,6 @@ from __future__ import annotations import logging -import warnings from typing import Any from unittest.mock import MagicMock, patch @@ -27,9 +26,7 @@ from airflow.models import Connection from airflow.models.dag import DAG from airflow.providers.apache.livy.hooks.livy import BatchState, LivyHook from airflow.providers.apache.livy.operators.livy import ( - _DURABLE_UNSET, LivyOperator, - _warn_and_disable_durable_pre_3_3, ) from airflow.providers.common.compat.sdk import AirflowException, timezone @@ -742,18 +739,3 @@ class TestLivyOperatorResumable: operator.poll_until_complete(55, {}) assert operator._batch_id == 55 - - -class TestWarnAndDisableDurableAirflowPre3_3: - def test_no_warning_when_unset(self): - with warnings.catch_warnings(record=True) as caught: - warnings.simplefilter("always") - result = _warn_and_disable_durable_pre_3_3(_DURABLE_UNSET) - assert result is False - assert caught == [] - - @pytest.mark.parametrize("value", [True, False]) - def test_warns_and_disables_when_explicitly_set(self, value): - with pytest.warns(UserWarning, match="durable.*no effect"): - result = _warn_and_disable_durable_pre_3_3(value) - assert result is False diff --git a/providers/apache/spark/pyproject.toml b/providers/apache/spark/pyproject.toml index 5a635f3de24..d0532022fec 100644 --- a/providers/apache/spark/pyproject.toml +++ b/providers/apache/spark/pyproject.toml @@ -59,7 +59,7 @@ requires-python = ">=3.11" # After you modify the dependencies, and rebuild your Breeze CI image with ``breeze ci-image build`` dependencies = [ "apache-airflow>=2.11.0", - "apache-airflow-providers-common-compat>=1.12.0", + "apache-airflow-providers-common-compat>=1.12.0", # use next version "pyspark-client>=4.0.0", "grpcio-status>=1.67.0", "requests>=2.32.0", diff --git a/providers/apache/spark/src/airflow/providers/apache/spark/operators/spark_submit.py b/providers/apache/spark/src/airflow/providers/apache/spark/operators/spark_submit.py index 1a9d2748a65..7e32f799901 100644 --- a/providers/apache/spark/src/airflow/providers/apache/spark/operators/spark_submit.py +++ b/providers/apache/spark/src/airflow/providers/apache/spark/operators/spark_submit.py @@ -30,46 +30,13 @@ from airflow.providers.common.compat.openlineage.utils.spark import ( inject_parent_job_information_into_spark_properties, inject_transport_information_into_spark_properties, ) -from airflow.providers.common.compat.sdk import BaseOperator, conf +from airflow.providers.common.compat.sdk import BaseOperator, ResumableJobMixin, conf try: from airflow.providers.cncf.kubernetes import kube_client except ImportError: kube_client = None # type: ignore[assignment] -_DURABLE_UNSET = object() - - -def _warn_and_disable_durable_pre_3_3(durable: Any) -> bool: - """Shared by the <3.3 compat stub: durable has no effect below 3.3, warn if it was set.""" - if durable is not _DURABLE_UNSET: - warnings.warn( - "`durable` has no effect on Airflow versions below 3.3.", - UserWarning, - stacklevel=3, - ) - return False - - -try: - from airflow.sdk import ResumableJobMixin -except ImportError: - # ResumableJobMixin only exists on Airflow 3.3+; this provider still targets older - # versions. Drop this fallback once the provider's minimum Airflow version is >=3.3. - class ResumableJobMixin: # type: ignore[no-redef] - """Airflow <3.3 stub, task_state_store unavailable, always submits fresh.""" - - external_id_key: str = "remote_job_id" - - def __init__(self, *, durable: Any = _DURABLE_UNSET, **kwargs: Any) -> None: - super().__init__(**kwargs) - self.durable = _warn_and_disable_durable_pre_3_3(durable) - - def execute_resumable(self, context): - external_id = self.submit_job(context) - self.poll_until_complete(external_id, context) - return self.get_job_result(external_id, context) - if TYPE_CHECKING: from pydantic import JsonValue diff --git a/providers/apache/spark/tests/unit/apache/spark/operators/test_spark_submit.py b/providers/apache/spark/tests/unit/apache/spark/operators/test_spark_submit.py index c55cb10f3eb..f3cb0945923 100644 --- a/providers/apache/spark/tests/unit/apache/spark/operators/test_spark_submit.py +++ b/providers/apache/spark/tests/unit/apache/spark/operators/test_spark_submit.py @@ -29,9 +29,7 @@ from airflow.exceptions import AirflowProviderDeprecationWarning from airflow.models import DagRun, TaskInstance from airflow.models.dag import DAG from airflow.providers.apache.spark.operators.spark_submit import ( - _DURABLE_UNSET, SparkSubmitOperator, - _warn_and_disable_durable_pre_3_3, ) from airflow.providers.common.compat.sdk import timezone from airflow.utils.types import DagRunType @@ -1136,18 +1134,3 @@ class TestSparkSubmitOperatorK8sTracking: operator.execute(context={"task_state_store": task_store}) assert task_store.get("spark_job_id") is None - - -class TestWarnAndDisableDurableAirflowPre3_3: - def test_no_warning_when_unset(self): - with warnings.catch_warnings(record=True) as caught: - warnings.simplefilter("always") - result = _warn_and_disable_durable_pre_3_3(_DURABLE_UNSET) - assert result is False - assert caught == [] - - @pytest.mark.parametrize("value", [True, False]) - def test_warns_and_disables_when_explicitly_set(self, value): - with pytest.warns(UserWarning, match="durable.*no effect"): - result = _warn_and_disable_durable_pre_3_3(value) - assert result is False diff --git a/providers/common/compat/src/airflow/providers/common/compat/_resumable_job_mixin.py b/providers/common/compat/src/airflow/providers/common/compat/_resumable_job_mixin.py new file mode 100644 index 00000000000..aa2b9c2f4c3 --- /dev/null +++ b/providers/common/compat/src/airflow/providers/common/compat/_resumable_job_mixin.py @@ -0,0 +1,58 @@ +# 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. +""" +Stand-in for ``airflow.sdk.ResumableJobMixin`` on Airflow versions below 3.3. + +``ResumableJobMixin`` and the ``task_state_store`` it relies on arrived in Airflow 3.3. An operator +that imports the mixin from :mod:`airflow.providers.common.compat.sdk` keeps working on older +versions: this stand-in submits a new job on every try, exactly what the operator did before it +became durable. That module tries the SDK first, so on 3.3 and later it hands out the SDK class and +this one is never imported. +""" + +from __future__ import annotations + +import warnings +from typing import TYPE_CHECKING, Any + +if TYPE_CHECKING: + from airflow.sdk.definitions.context import Context + +# Tells "durable left at its default" apart from an explicit durable=True/False, so only an explicit +# value warns that it has no effect here. +_DURABLE_UNSET: Any = object() + + +class ResumableJobMixin: + """Airflow < 3.3 stand-in: no ``task_state_store`` to resume from, so every try submits a new job.""" + + external_id_key: str = "remote_job_id" + + def __init__(self, *, durable: bool = _DURABLE_UNSET, **kwargs: Any) -> None: + if durable is not _DURABLE_UNSET: + warnings.warn( + "`durable` has no effect on Airflow versions below 3.3.", + UserWarning, + stacklevel=2, + ) + super().__init__(**kwargs) + self.durable = False + + def execute_resumable(self, context: Context) -> Any: + external_id = self.submit_job(context) # type: ignore[attr-defined] + self.poll_until_complete(external_id, context) # type: ignore[attr-defined] + return self.get_job_result(external_id, context) # type: ignore[attr-defined] diff --git a/providers/common/compat/src/airflow/providers/common/compat/sdk.py b/providers/common/compat/src/airflow/providers/common/compat/sdk.py index f37a194b337..d123fcc5269 100644 --- a/providers/common/compat/src/airflow/providers/common/compat/sdk.py +++ b/providers/common/compat/src/airflow/providers/common/compat/sdk.py @@ -127,6 +127,7 @@ if TYPE_CHECKING: RetryPolicy as RetryPolicy, RetryRule as RetryRule, ) + from airflow.sdk import ResumableJobMixin as ResumableJobMixin 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, @@ -372,6 +373,13 @@ _AIRFLOW_3_3_ONLY_RETRY_POLICIES: dict[str, str | tuple[str, ...]] = { if AIRFLOW_V_3_3_PLUS: _IMPORT_MAP.update(_AIRFLOW_3_3_ONLY_RETRY_POLICIES) +# ResumableJobMixin (AIP-103) arrived in Airflow 3.3. The SDK is tried first; below 3.3 the stand-in +# submits a new job on every try, so a durable operator still runs there, just without resuming. +_IMPORT_MAP["ResumableJobMixin"] = ( + "airflow.sdk", + "airflow.providers.common.compat._resumable_job_mixin", +) + # 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 diff --git a/providers/common/compat/tests/unit/common/compat/test__resumable_job_mixin.py b/providers/common/compat/tests/unit/common/compat/test__resumable_job_mixin.py new file mode 100644 index 00000000000..03592018b1a --- /dev/null +++ b/providers/common/compat/tests/unit/common/compat/test__resumable_job_mixin.py @@ -0,0 +1,95 @@ +# 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. +from __future__ import annotations + +import warnings +from unittest.mock import MagicMock, call + +import pytest + +from airflow.providers.common.compat import _resumable_job_mixin, sdk + +from tests_common.test_utils.version_compat import AIRFLOW_V_3_3_PLUS + + +class _Base: + def __init__(self, **kwargs): + self.base_kwargs = kwargs + + +class _JobOperator(_resumable_job_mixin.ResumableJobMixin, _Base): + external_id_key = "job_id" + + def __init__(self, **kwargs): + super().__init__(**kwargs) + self.calls = MagicMock() + + def submit_job(self, context): + self.calls.submit_job(context) + return "job-1" + + def poll_until_complete(self, external_id, context): + self.calls.poll_until_complete(external_id, context) + + def get_job_result(self, external_id, context): + self.calls.get_job_result(external_id, context) + return "result" + + +def test_resumable_job_mixin_is_exported(): + assert "ResumableJobMixin" in sdk.__all__ + + [email protected](not AIRFLOW_V_3_3_PLUS, reason="ResumableJobMixin arrived in Airflow 3.3") +def test_sdk_class_is_used_from_airflow_3_3(): + from airflow.sdk import ResumableJobMixin + + assert sdk.ResumableJobMixin is ResumableJobMixin + + [email protected](AIRFLOW_V_3_3_PLUS, reason="The SDK has ResumableJobMixin from Airflow 3.3") +def test_stand_in_is_used_below_airflow_3_3(): + assert sdk.ResumableJobMixin is _resumable_job_mixin.ResumableJobMixin + + +class TestStandIn: + def test_durable_left_at_default_does_not_warn(self): + with warnings.catch_warnings(): + warnings.simplefilter("error") + operator = _JobOperator(task_id="t") + + assert operator.durable is False + assert operator.base_kwargs == {"task_id": "t"} + + @pytest.mark.parametrize("durable", [True, False]) + def test_explicit_durable_warns_and_is_disabled(self, durable): + with pytest.warns(UserWarning, match="`durable` has no effect on Airflow versions below 3.3"): + operator = _JobOperator(task_id="t", durable=durable) + + assert operator.durable is False + assert operator.base_kwargs == {"task_id": "t"} + + def test_execute_resumable_submits_polls_and_returns_result(self): + operator = _JobOperator(task_id="t") + context = {} + + assert operator.execute_resumable(context) == "result" + assert operator.calls.mock_calls == [ + call.submit_job(context), + call.poll_until_complete("job-1", context), + call.get_job_result("job-1", context), + ] diff --git a/providers/databricks/pyproject.toml b/providers/databricks/pyproject.toml index 28910f667f5..0729c8c18a9 100644 --- a/providers/databricks/pyproject.toml +++ b/providers/databricks/pyproject.toml @@ -59,7 +59,7 @@ requires-python = ">=3.11" # After you modify the dependencies, and rebuild your Breeze CI image with ``breeze ci-image build`` dependencies = [ "apache-airflow>=2.11.0", - "apache-airflow-providers-common-compat>=1.21.0", + "apache-airflow-providers-common-compat>=1.21.0", # use next version "apache-airflow-providers-common-sql>=1.32.0", "requests>=2.32.0,<3", "databricks-sql-connector>=4.4.0", diff --git a/providers/databricks/src/airflow/providers/databricks/operators/databricks.py b/providers/databricks/src/airflow/providers/databricks/operators/databricks.py index 8b3e3fcd8ae..0f01fa2ae55 100644 --- a/providers/databricks/src/airflow/providers/databricks/operators/databricks.py +++ b/providers/databricks/src/airflow/providers/databricks/operators/databricks.py @@ -24,13 +24,19 @@ import copy import hashlib import json as json_utils import time -import warnings from abc import ABC, abstractmethod from collections.abc import Mapping, Sequence from functools import cached_property from typing import TYPE_CHECKING, Any, cast -from airflow.providers.common.compat.sdk import AirflowException, BaseOperator, BaseOperatorLink, XCom, conf +from airflow.providers.common.compat.sdk import ( + AirflowException, + BaseOperator, + BaseOperatorLink, + ResumableJobMixin, + XCom, + conf, +) from airflow.providers.databricks.exceptions import DatabricksApiError, DatabricksOperatorPayloadError from airflow.providers.databricks.hooks.databricks import ( DatabricksHook, @@ -69,38 +75,6 @@ if TYPE_CHECKING: from airflow.sdk import TaskGroup from airflow.sdk.types import Context, Logger -_DURABLE_UNSET = object() - - -def _warn_and_disable_durable_pre_3_3(durable: Any) -> bool: - """Shared by the <3.3 compat stub: durable has no effect below 3.3, warn if it was set.""" - if durable is not _DURABLE_UNSET: - warnings.warn( - "`durable` has no effect on Airflow versions below 3.3.", - UserWarning, - stacklevel=3, - ) - return False - - -try: - from airflow.sdk import ResumableJobMixin -except ImportError: - - class ResumableJobMixin: # type: ignore[no-redef] - """Airflow <3.3 stub, task_state_store unavailable, always submits fresh.""" - - external_id_key: str = "databricks_run_id" - - def __init__(self, *, durable: Any = _DURABLE_UNSET, **kwargs: Any) -> None: - super().__init__(**kwargs) - self.durable = _warn_and_disable_durable_pre_3_3(durable) - - def execute_resumable(self, context): - external_id = self.submit_job(context) - self.poll_until_complete(external_id, context) - return self.get_job_result(external_id, context) - DEFER_METHOD_NAME = "execute_complete" XCOM_RUN_ID_KEY = "run_id" diff --git a/providers/databricks/tests/unit/databricks/operators/test_databricks.py b/providers/databricks/tests/unit/databricks/operators/test_databricks.py index 354bdf8e880..40672c42f2e 100644 --- a/providers/databricks/tests/unit/databricks/operators/test_databricks.py +++ b/providers/databricks/tests/unit/databricks/operators/test_databricks.py @@ -19,7 +19,6 @@ from __future__ import annotations import copy import hashlib -import warnings from datetime import datetime, timedelta from typing import Any from unittest import mock @@ -41,7 +40,6 @@ from airflow.providers.common.compat.sdk import AirflowException, BaseOperator, from airflow.providers.databricks.exceptions import DatabricksApiError from airflow.providers.databricks.hooks.databricks import RunState, SQLStatementState from airflow.providers.databricks.operators.databricks import ( - _DURABLE_UNSET, DatabricksCreateJobsOperator, DatabricksNotebookOperator, DatabricksRunNowOperator, @@ -49,7 +47,6 @@ from airflow.providers.databricks.operators.databricks import ( DatabricksSubmitRunOperator, DatabricksTaskBaseOperator, DatabricksTaskOperator, - _warn_and_disable_durable_pre_3_3, ) from airflow.providers.databricks.triggers.databricks import ( DatabricksExecutionTrigger, @@ -5270,18 +5267,3 @@ class TestDatabricksTaskOperator: operator.monitor_databricks_job() assert mock_databricks_hook.return_value.get_run.call_count == 3 - - -class TestWarnAndDisableDurableAirflowPre3_3: - def test_no_warning_when_unset(self): - with warnings.catch_warnings(record=True) as caught: - warnings.simplefilter("always") - result = _warn_and_disable_durable_pre_3_3(_DURABLE_UNSET) - assert result is False - assert caught == [] - - @pytest.mark.parametrize("value", [True, False]) - def test_warns_and_disables_when_explicitly_set(self, value): - with pytest.warns(UserWarning, match="durable.*no effect"): - result = _warn_and_disable_durable_pre_3_3(value) - assert result is False diff --git a/providers/google/pyproject.toml b/providers/google/pyproject.toml index 644fbef1610..7706552f0b3 100644 --- a/providers/google/pyproject.toml +++ b/providers/google/pyproject.toml @@ -59,7 +59,7 @@ requires-python = ">=3.11" # After you modify the dependencies, and rebuild your Breeze CI image with ``breeze ci-image build`` dependencies = [ "apache-airflow>=2.11.0", - "apache-airflow-providers-common-compat>=1.13.0", + "apache-airflow-providers-common-compat>=1.13.0", # use next version "apache-airflow-providers-common-sql>=1.32.0", "asgiref>=3.5.2; python_version < '3.14'", "asgiref>=3.11.1; python_version >= '3.14'", diff --git a/providers/google/src/airflow/providers/google/cloud/operators/bigquery.py b/providers/google/src/airflow/providers/google/cloud/operators/bigquery.py index 8b60e461ead..7d1adc4faa9 100644 --- a/providers/google/src/airflow/providers/google/cloud/operators/bigquery.py +++ b/providers/google/src/airflow/providers/google/cloud/operators/bigquery.py @@ -34,7 +34,12 @@ from google.cloud.bigquery.routine import Routine from google.cloud.bigquery.table import RowIterator, Table, TableListItem, TableReference from airflow.exceptions import AirflowProviderDeprecationWarning -from airflow.providers.common.compat.sdk import AirflowException, AirflowSkipException, conf +from airflow.providers.common.compat.sdk import ( + AirflowException, + AirflowSkipException, + ResumableJobMixin, + conf, +) from airflow.providers.common.sql.operators.sql import ( # for _parse_boolean SQLCheckOperator, SQLColumnCheckOperator, @@ -73,39 +78,6 @@ except ImportError: return value is not NOTSET -_DURABLE_UNSET = object() - - -def _warn_and_disable_durable_pre_3_3(durable: Any) -> bool: - """Shared by the <3.3 compat stub: durable has no effect below 3.3, warn if it was set.""" - if durable is not _DURABLE_UNSET: - warnings.warn( - "`durable` has no effect on Airflow versions below 3.3.", - UserWarning, - stacklevel=3, - ) - return False - - -try: - from airflow.sdk import ResumableJobMixin -except ImportError: - - class ResumableJobMixin: # type: ignore[no-redef] - """Airflow <3.3 stub, task_state_store unavailable, always submits fresh.""" - - external_id_key: str = "bigquery_job_id" - - def __init__(self, *, durable: Any = _DURABLE_UNSET, **kwargs: Any) -> None: - super().__init__(**kwargs) - self.durable = _warn_and_disable_durable_pre_3_3(durable) - - def execute_resumable(self, context): - external_id = self.submit_job(context) - self.poll_until_complete(external_id, context) - return self.get_job_result(external_id, context) - - if TYPE_CHECKING: from google.api_core.retry import Retry from google.cloud.bigquery import UnknownJob diff --git a/providers/google/tests/unit/google/cloud/operators/test_bigquery.py b/providers/google/tests/unit/google/cloud/operators/test_bigquery.py index c44947bf1f6..b959d14b8a4 100644 --- a/providers/google/tests/unit/google/cloud/operators/test_bigquery.py +++ b/providers/google/tests/unit/google/cloud/operators/test_bigquery.py @@ -20,7 +20,6 @@ from __future__ import annotations import json import logging import os -import warnings from contextlib import suppress from unittest import mock from unittest.mock import ANY, MagicMock @@ -53,7 +52,6 @@ from airflow.providers.common.compat.sdk import ( ) from airflow.providers.google.cloud.openlineage.utils import BIGQUERY_NAMESPACE from airflow.providers.google.cloud.operators.bigquery import ( - _DURABLE_UNSET, BigQueryCheckOperator, BigQueryColumnCheckOperator, BigQueryCreateEmptyDatasetOperator, @@ -76,7 +74,6 @@ from airflow.providers.google.cloud.operators.bigquery import ( BigQueryUpdateTableSchemaOperator, BigQueryUpsertTableOperator, BigQueryValueCheckOperator, - _warn_and_disable_durable_pre_3_3, ) from airflow.providers.google.cloud.triggers.bigquery import ( BigQueryCheckTrigger, @@ -3446,18 +3443,3 @@ class TestBigQueryListRoutinesOperator: call_kwargs = mock_hook.return_value.list_routines.call_args.kwargs assert call_kwargs["dataset_id"] == TEST_DATASET assert call_kwargs["max_results"] == 10 - - -class TestWarnAndDisableDurableAirflowPre3_3: - def test_no_warning_when_unset(self): - with warnings.catch_warnings(record=True) as caught: - warnings.simplefilter("always") - result = _warn_and_disable_durable_pre_3_3(_DURABLE_UNSET) - assert result is False - assert caught == [] - - @pytest.mark.parametrize("value", [True, False]) - def test_warns_and_disables_when_explicitly_set(self, value): - with pytest.warns(UserWarning, match="durable.*no effect"): - result = _warn_and_disable_durable_pre_3_3(value) - assert result is False diff --git a/providers/snowflake/pyproject.toml b/providers/snowflake/pyproject.toml index 84de0a45667..9f7a906e299 100644 --- a/providers/snowflake/pyproject.toml +++ b/providers/snowflake/pyproject.toml @@ -59,7 +59,7 @@ requires-python = ">=3.11" # After you modify the dependencies, and rebuild your Breeze CI image with ``breeze ci-image build`` dependencies = [ "apache-airflow>=2.11.0", - "apache-airflow-providers-common-compat>=1.12.0", + "apache-airflow-providers-common-compat>=1.12.0", # use next version "apache-airflow-providers-common-sql>=1.32.0", "asgiref>=2.3.0; python_version < '3.14'", "asgiref>=3.11.1; python_version >= '3.14'", diff --git a/providers/snowflake/src/airflow/providers/snowflake/operators/snowflake.py b/providers/snowflake/src/airflow/providers/snowflake/operators/snowflake.py index dc5aa5c9711..2855815843e 100644 --- a/providers/snowflake/src/airflow/providers/snowflake/operators/snowflake.py +++ b/providers/snowflake/src/airflow/providers/snowflake/operators/snowflake.py @@ -18,7 +18,6 @@ from __future__ import annotations import time -import warnings from collections.abc import Iterable, Mapping, Sequence from datetime import timedelta from functools import cached_property @@ -26,7 +25,7 @@ from typing import TYPE_CHECKING, Any, ClassVar, SupportsAbs, cast import requests -from airflow.providers.common.compat.sdk import conf +from airflow.providers.common.compat.sdk import ResumableJobMixin, conf from airflow.providers.common.sql.operators.sql import ( SQLCheckOperator, SQLExecuteQueryOperator, @@ -36,39 +35,6 @@ from airflow.providers.common.sql.operators.sql import ( from airflow.providers.snowflake.hooks.snowflake_sql_api import SnowflakeSqlApiHook from airflow.providers.snowflake.triggers.snowflake_trigger import SnowflakeSqlApiTrigger -_DURABLE_UNSET = object() - - -def _warn_and_disable_durable_pre_3_3(durable: Any) -> bool: - """Shared by the <3.3 compat stub: durable has no effect below 3.3, warn if it was set.""" - if durable is not _DURABLE_UNSET: - warnings.warn( - "`durable` has no effect on Airflow versions below 3.3.", - UserWarning, - stacklevel=3, - ) - return False - - -try: - from airflow.sdk import ResumableJobMixin -except ImportError: - - class ResumableJobMixin: # type: ignore[no-redef] - """Airflow <3.3 stub, task_state_store unavailable, always submits fresh.""" - - external_id_key: str = "snowflake_query_ids" - - def __init__(self, *, durable: Any = _DURABLE_UNSET, **kwargs: Any) -> None: - super().__init__(**kwargs) - self.durable = _warn_and_disable_durable_pre_3_3(durable) - - def execute_resumable(self, context): - external_id = self.submit_job(context) - self.poll_until_complete(external_id, context) - return self.get_job_result(external_id, context) - - if TYPE_CHECKING: from pydantic import JsonValue diff --git a/providers/snowflake/tests/unit/snowflake/operators/test_snowflake.py b/providers/snowflake/tests/unit/snowflake/operators/test_snowflake.py index c8a75157fe8..d37a74272e1 100644 --- a/providers/snowflake/tests/unit/snowflake/operators/test_snowflake.py +++ b/providers/snowflake/tests/unit/snowflake/operators/test_snowflake.py @@ -17,7 +17,6 @@ # under the License. from __future__ import annotations -import warnings from unittest import mock from unittest.mock import MagicMock, call @@ -32,13 +31,11 @@ from airflow.models.taskinstance import TaskInstance from airflow.providers.common.compat.sdk import TaskDeferred from airflow.providers.common.sql.operators.sql import SQLExecuteQueryOperator from airflow.providers.snowflake.operators.snowflake import ( - _DURABLE_UNSET, SnowflakeCheckOperator, SnowflakeIntervalCheckOperator, SnowflakeNotebookOperator, SnowflakeSqlApiOperator, SnowflakeValueCheckOperator, - _warn_and_disable_durable_pre_3_3, ) from airflow.providers.snowflake.triggers.snowflake_trigger import SnowflakeSqlApiTrigger from airflow.utils.types import DagRunType @@ -997,21 +994,6 @@ class TestSnowflakeSqlApiOperatorDurable: assert operator.durable is False -class TestWarnAndDisableDurableAirflowPre3_3: - def test_no_warning_when_unset(self): - with warnings.catch_warnings(record=True) as caught: - warnings.simplefilter("always") - result = _warn_and_disable_durable_pre_3_3(_DURABLE_UNSET) - assert result is False - assert caught == [] - - @pytest.mark.parametrize("value", [True, False]) - def test_warns_and_disables_when_explicitly_set(self, value): - with pytest.warns(UserWarning, match="durable.*no effect"): - result = _warn_and_disable_durable_pre_3_3(value) - assert result is False - - class TestSnowflakeNotebookOperatorSQL: """Tests for SQL query building."""
