This is an automated email from the ASF dual-hosted git repository.
kaxilnaik pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new 16ff7dcb3b8 Move ``SetupTeardownContext`` to Task SDK (#53160)
16ff7dcb3b8 is described below
commit 16ff7dcb3b80cbb71bb3e8494e5664b8c0ff2445
Author: Kaxil Naik <[email protected]>
AuthorDate: Fri Jul 11 12:03:37 2025 +0530
Move ``SetupTeardownContext`` to Task SDK (#53160)
---
airflow-core/src/airflow/utils/__init__.py | 11 +++++++++++
task-sdk/src/airflow/sdk/bases/operator.py | 2 +-
.../src/airflow/sdk/definitions/_internal/abstractoperator.py | 2 +-
.../src/airflow/sdk/definitions/_internal}/setup_teardown.py | 4 ++--
.../src/airflow/sdk/definitions/decorators/setup_teardown.py | 7 +++----
task-sdk/src/airflow/sdk/definitions/xcom_arg.py | 2 +-
6 files changed, 19 insertions(+), 9 deletions(-)
diff --git a/airflow-core/src/airflow/utils/__init__.py
b/airflow-core/src/airflow/utils/__init__.py
index 8a1573c8457..574fbb642f3 100644
--- a/airflow-core/src/airflow/utils/__init__.py
+++ b/airflow-core/src/airflow/utils/__init__.py
@@ -18,6 +18,8 @@
from __future__ import annotations
+from airflow.utils.deprecation_tools import add_deprecated_classes
+
def _deprecate_this_module(message: str, **shims: tuple[str, str]):
import warnings
@@ -34,3 +36,12 @@ def _deprecate_this_module(message: str, **shims: tuple[str,
str]):
return getattr(__import__(impa), attr)
return __getattr__
+
+
+__deprecated_classes = {
+ "setup_teardown": {
+ "BaseSetupTeardownContext":
"airflow.sdk.definitions._internal.setup_teardown.BaseSetupTeardownContext",
+ "SetupTeardownContext":
"airflow.sdk.definitions._internal.setup_teardown.SetupTeardownContext",
+ },
+}
+add_deprecated_classes(__deprecated_classes, __name__)
diff --git a/task-sdk/src/airflow/sdk/bases/operator.py
b/task-sdk/src/airflow/sdk/bases/operator.py
index 3f74694aaa7..8cab2861e92 100644
--- a/task-sdk/src/airflow/sdk/bases/operator.py
+++ b/task-sdk/src/airflow/sdk/bases/operator.py
@@ -55,6 +55,7 @@ from airflow.sdk.definitions._internal.abstractoperator
import (
)
from airflow.sdk.definitions._internal.decorators import
fixup_decorator_warning_stack
from airflow.sdk.definitions._internal.node import validate_key
+from airflow.sdk.definitions._internal.setup_teardown import
SetupTeardownContext
from airflow.sdk.definitions._internal.types import NOTSET,
validate_instance_args
from airflow.sdk.definitions.edges import EdgeModifier
from airflow.sdk.definitions.mappedoperator import OperatorPartial,
validate_mapping_kwargs
@@ -65,7 +66,6 @@ from airflow.task.priority_strategy import (
validate_and_load_priority_weight_strategy,
)
from airflow.utils import timezone
-from airflow.utils.setup_teardown import SetupTeardownContext
from airflow.utils.trigger_rule import TriggerRule
from airflow.utils.weight_rule import db_safe_priority
diff --git a/task-sdk/src/airflow/sdk/definitions/_internal/abstractoperator.py
b/task-sdk/src/airflow/sdk/definitions/_internal/abstractoperator.py
index 8934cd0e453..e7ee16af247 100644
--- a/task-sdk/src/airflow/sdk/definitions/_internal/abstractoperator.py
+++ b/task-sdk/src/airflow/sdk/definitions/_internal/abstractoperator.py
@@ -33,9 +33,9 @@ import methodtools
from airflow.configuration import conf
from airflow.sdk.definitions._internal.mixins import DependencyMixin
from airflow.sdk.definitions._internal.node import DAGNode
+from airflow.sdk.definitions._internal.setup_teardown import
SetupTeardownContext
from airflow.sdk.definitions._internal.templater import Templater
from airflow.sdk.definitions.context import Context
-from airflow.utils.setup_teardown import SetupTeardownContext
from airflow.utils.trigger_rule import TriggerRule
from airflow.utils.weight_rule import WeightRule
diff --git a/airflow-core/src/airflow/utils/setup_teardown.py
b/task-sdk/src/airflow/sdk/definitions/_internal/setup_teardown.py
similarity index 99%
rename from airflow-core/src/airflow/utils/setup_teardown.py
rename to task-sdk/src/airflow/sdk/definitions/_internal/setup_teardown.py
index b62e3ab298b..961099c6f97 100644
--- a/airflow-core/src/airflow/utils/setup_teardown.py
+++ b/task-sdk/src/airflow/sdk/definitions/_internal/setup_teardown.py
@@ -1,3 +1,4 @@
+#
# 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
@@ -14,7 +15,6 @@
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
-
from __future__ import annotations
from typing import TYPE_CHECKING, cast
@@ -22,8 +22,8 @@ from typing import TYPE_CHECKING, cast
from airflow.exceptions import AirflowException
if TYPE_CHECKING:
- from airflow.models.taskmixin import DependencyMixin
from airflow.sdk.definitions._internal.abstractoperator import
AbstractOperator
+ from airflow.sdk.definitions._internal.mixins import DependencyMixin
from airflow.sdk.definitions.xcom_arg import PlainXComArg
diff --git a/task-sdk/src/airflow/sdk/definitions/decorators/setup_teardown.py
b/task-sdk/src/airflow/sdk/definitions/decorators/setup_teardown.py
index e5e2bf40ee9..bcb8e6d12e2 100644
--- a/task-sdk/src/airflow/sdk/definitions/decorators/setup_teardown.py
+++ b/task-sdk/src/airflow/sdk/definitions/decorators/setup_teardown.py
@@ -21,17 +21,16 @@ from collections.abc import Callable
from typing import TYPE_CHECKING
from airflow.exceptions import AirflowException
-from airflow.providers.standard.version_compat import AIRFLOW_V_3_0_PLUS
from airflow.sdk.bases.operator import BaseOperator
+from airflow.sdk.definitions._internal.setup_teardown import
SetupTeardownContext
from airflow.sdk.definitions.decorators.task_group import _TaskGroupFactory
-from airflow.utils.setup_teardown import SetupTeardownContext
if TYPE_CHECKING:
from airflow.sdk.definitions.xcom_arg import XComArg
-if AIRFLOW_V_3_0_PLUS:
+try:
from airflow.providers.standard.decorators.python import python_task
-else:
+except (ImportError, AttributeError):
from airflow.decorators import python_task # type: ignore
diff --git a/task-sdk/src/airflow/sdk/definitions/xcom_arg.py
b/task-sdk/src/airflow/sdk/definitions/xcom_arg.py
index 8908429d08a..4b74c085c4c 100644
--- a/task-sdk/src/airflow/sdk/definitions/xcom_arg.py
+++ b/task-sdk/src/airflow/sdk/definitions/xcom_arg.py
@@ -27,10 +27,10 @@ from typing import TYPE_CHECKING, Any, overload
from airflow.exceptions import AirflowException, XComNotFound
from airflow.sdk.definitions._internal.abstractoperator import AbstractOperator
from airflow.sdk.definitions._internal.mixins import DependencyMixin,
ResolveMixin
+from airflow.sdk.definitions._internal.setup_teardown import
SetupTeardownContext
from airflow.sdk.definitions._internal.types import NOTSET, ArgNotSet
from airflow.sdk.execution_time.lazy_sequence import LazyXComSequence
from airflow.sdk.execution_time.xcom import BaseXCom
-from airflow.utils.setup_teardown import SetupTeardownContext
from airflow.utils.trigger_rule import TriggerRule
if TYPE_CHECKING:
