Skip to content

Commit cfdf63e

Browse files
authored
Adding deprecation notice for get_current_context in std provider (#50301)
1 parent bf99522 commit cfdf63e

File tree

2 files changed

+54
-14
lines changed

2 files changed

+54
-14
lines changed

providers/standard/src/airflow/providers/standard/operators/python.py

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,7 @@
2626
import sys
2727
import textwrap
2828
import types
29+
import warnings
2930
from abc import ABCMeta, abstractmethod
3031
from collections.abc import Collection, Container, Iterable, Mapping, Sequence
3132
from functools import cache
@@ -38,6 +39,7 @@
3839
from airflow.exceptions import (
3940
AirflowConfigException,
4041
AirflowException,
42+
AirflowProviderDeprecationWarning,
4143
AirflowSkipException,
4244
DeserializingResultError,
4345
)
@@ -1113,6 +1115,13 @@ def my_task():
11131115
was starting to execute.
11141116
"""
11151117
if AIRFLOW_V_3_0_PLUS:
1118+
warnings.warn(
1119+
"Using get_current_context from standard provider is deprecated and will be removed."
1120+
"Please import `from airflow.sdk import get_current_context` and use it instead.",
1121+
AirflowProviderDeprecationWarning,
1122+
stacklevel=2,
1123+
)
1124+
11161125
from airflow.sdk import get_current_context
11171126

11181127
return get_current_context()

providers/standard/tests/unit/standard/operators/test_python.py

Lines changed: 45 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -44,6 +44,7 @@
4444
from airflow.config_templates.airflow_local_settings import DEFAULT_LOGGING_CONFIG
4545
from airflow.exceptions import (
4646
AirflowException,
47+
AirflowProviderDeprecationWarning,
4748
DeserializingResultError,
4849
)
4950
from airflow.models.baseoperator import BaseOperator
@@ -1810,22 +1811,35 @@ def default_kwargs(*, python_version=DEFAULT_PYTHON_VERSION, **kwargs):
18101811

18111812
class TestCurrentContext:
18121813
def test_current_context_no_context_raise(self):
1813-
with pytest.raises(RuntimeError):
1814-
get_current_context()
1814+
if AIRFLOW_V_3_0_PLUS:
1815+
with pytest.warns(AirflowProviderDeprecationWarning):
1816+
with pytest.raises(RuntimeError):
1817+
get_current_context()
1818+
else:
1819+
with pytest.raises(RuntimeError):
1820+
get_current_context()
18151821

18161822
def test_current_context_roundtrip(self):
18171823
example_context = {"Hello": "World"}
1818-
18191824
with set_current_context(example_context):
1820-
assert get_current_context() == example_context
1825+
if AIRFLOW_V_3_0_PLUS:
1826+
with pytest.warns(AirflowProviderDeprecationWarning):
1827+
assert get_current_context() == example_context
1828+
else:
1829+
assert get_current_context() == example_context
18211830

18221831
def test_context_removed_after_exit(self):
18231832
example_context = {"Hello": "World"}
18241833

18251834
with set_current_context(example_context):
18261835
pass
1827-
with pytest.raises(RuntimeError):
1828-
get_current_context()
1836+
if AIRFLOW_V_3_0_PLUS:
1837+
with pytest.warns(AirflowProviderDeprecationWarning):
1838+
with pytest.raises(RuntimeError):
1839+
get_current_context()
1840+
else:
1841+
with pytest.raises(RuntimeError):
1842+
get_current_context()
18291843

18301844
def test_nested_context(self):
18311845
"""
@@ -1842,12 +1856,21 @@ def test_nested_context(self):
18421856
ctx_obj = set_current_context(new_context)
18431857
ctx_obj.__enter__()
18441858
ctx_list.append(ctx_obj)
1845-
for i in reversed(range(max_stack_depth)):
1846-
# Iterate over contexts in reverse order - stack is LIFO
1847-
ctx = get_current_context()
1848-
assert ctx["ContextId"] == i
1849-
# End of with statement
1850-
ctx_list[i].__exit__(None, None, None)
1859+
if AIRFLOW_V_3_0_PLUS:
1860+
with pytest.warns(AirflowProviderDeprecationWarning):
1861+
for i in reversed(range(max_stack_depth)):
1862+
# Iterate over contexts in reverse order - stack is LIFO
1863+
ctx = get_current_context()
1864+
assert ctx["ContextId"] == i
1865+
# End of with statement
1866+
ctx_list[i].__exit__(None, None, None)
1867+
else:
1868+
for i in reversed(range(max_stack_depth)):
1869+
# Iterate over contexts in reverse order - stack is LIFO
1870+
ctx = get_current_context()
1871+
assert ctx["ContextId"] == i
1872+
# End of with statement
1873+
ctx_list[i].__exit__(None, None, None)
18511874

18521875

18531876
class MyContextAssertOperator(BaseOperator):
@@ -1889,12 +1912,20 @@ class TestCurrentContextRuntime:
18891912
def test_context_in_task(self):
18901913
with DAG(dag_id="assert_context_dag", default_args=DEFAULT_ARGS, schedule="@once"):
18911914
op = MyContextAssertOperator(task_id="assert_context")
1892-
op.run(ignore_first_depends_on_past=True, ignore_ti_state=True)
1915+
if AIRFLOW_V_3_0_PLUS:
1916+
with pytest.warns(AirflowProviderDeprecationWarning):
1917+
op.run(ignore_first_depends_on_past=True, ignore_ti_state=True)
1918+
else:
1919+
op.run(ignore_first_depends_on_past=True, ignore_ti_state=True)
18931920

18941921
def test_get_context_in_old_style_context_task(self):
18951922
with DAG(dag_id="edge_case_context_dag", default_args=DEFAULT_ARGS, schedule="@once"):
18961923
op = PythonOperator(python_callable=get_all_the_context, task_id="get_all_the_context")
1897-
op.run(ignore_first_depends_on_past=True, ignore_ti_state=True)
1924+
if AIRFLOW_V_3_0_PLUS:
1925+
with pytest.warns(AirflowProviderDeprecationWarning):
1926+
op.run(ignore_first_depends_on_past=True, ignore_ti_state=True)
1927+
else:
1928+
op.run(ignore_first_depends_on_past=True, ignore_ti_state=True)
18981929

18991930

19001931
@pytest.mark.need_serialized_dag(False)

0 commit comments

Comments
 (0)