From 99f055b0c370d9dd5c9c14337f09edd1f9394317 Mon Sep 17 00:00:00 2001 From: Kaxil Naik Date: Tue, 8 Jul 2025 23:42:03 +0530 Subject: [PATCH] Remove `set_current_context` from `airflow.models.taskinstance` This function isn't used anymore since Airflow 3.0. Only usage in test has been updated. --- .../src/airflow/models/taskinstance.py | 24 +------------------ .../unit/standard/operators/test_python.py | 5 +++- 2 files changed, 5 insertions(+), 24 deletions(-) diff --git a/airflow-core/src/airflow/models/taskinstance.py b/airflow-core/src/airflow/models/taskinstance.py index c2c840edccfd8..6c8d2680fb51d 100644 --- a/airflow-core/src/airflow/models/taskinstance.py +++ b/airflow-core/src/airflow/models/taskinstance.py @@ -26,7 +26,7 @@ import os import uuid from collections import defaultdict -from collections.abc import Collection, Generator, Iterable +from collections.abc import Collection, Iterable from datetime import timedelta from functools import cache from typing import TYPE_CHECKING, Any @@ -156,28 +156,6 @@ def _add_log( ) -@contextlib.contextmanager -def set_current_context(context: Context) -> Generator[Context, None, None]: - """ - Set the current execution context to the provided context object. - - This method should be called once per Task execution, before calling operator.execute. - """ - from airflow.sdk.definitions._internal.contextmanager import _CURRENT_CONTEXT - - _CURRENT_CONTEXT.append(context) - try: - yield context - finally: - expected_state = _CURRENT_CONTEXT.pop() - if expected_state != context: - log.warning( - "Current context is not equal to the state at context stack. Expected=%s, got=%s", - context, - expected_state, - ) - - def _stop_remaining_tasks(*, task_instance: TaskInstance, task_teardown_map=None, session: Session): """ Stop non-teardown tasks in dag. diff --git a/providers/standard/tests/unit/standard/operators/test_python.py b/providers/standard/tests/unit/standard/operators/test_python.py index ff0cd45a8694a..e9cecea81c34d 100644 --- a/providers/standard/tests/unit/standard/operators/test_python.py +++ b/providers/standard/tests/unit/standard/operators/test_python.py @@ -48,7 +48,7 @@ DeserializingResultError, ) from airflow.models.connection import Connection -from airflow.models.taskinstance import TaskInstance, clear_task_instances, set_current_context +from airflow.models.taskinstance import TaskInstance, clear_task_instances from airflow.providers.standard.operators.empty import EmptyOperator from airflow.providers.standard.operators.python import ( BranchExternalPythonOperator, @@ -74,8 +74,11 @@ if AIRFLOW_V_3_0_PLUS: from airflow.sdk import BaseOperator + from airflow.sdk.execution_time.context import set_current_context else: from airflow.models.baseoperator import BaseOperator # type: ignore[no-redef] + from airflow.models.taskinstance import set_current_context # type: ignore[attr-defined,no-redef] + if TYPE_CHECKING: from airflow.models.dag import DAG