Skip to content

Commit 630eb6f

Browse files
potiukpotiuk
committed
Test standard provider with Airflow 2.8 and 2.9
The standard provider has now min version of Airflow = 2.8 since apache#43553, but we have not tested it for Airflow 2.8 and 2.9.
1 parent ff6038b commit 630eb6f

7 files changed

Lines changed: 119 additions & 41 deletions

File tree

‎dev/breeze/src/airflow_breeze/global_constants.py‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -574,13 +574,13 @@ def get_airflow_extras():
574574
{
575575
"python-version": "3.9",
576576
"airflow-version": "2.8.4",
577-
"remove-providers": "cloudant fab edge standard",
577+
"remove-providers": "cloudant fab edge",
578578
"run-tests": "true",
579579
},
580580
{
581581
"python-version": "3.9",
582582
"airflow-version": "2.9.3",
583-
"remove-providers": "cloudant edge standard",
583+
"remove-providers": "cloudant edge",
584584
"run-tests": "true",
585585
},
586586
{

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

Lines changed: 41 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -54,15 +54,16 @@
5454
from airflow.settings import _ENABLE_AIP_44
5555
from airflow.typing_compat import Literal
5656
from airflow.utils import hashlib_wrapper
57-
from airflow.utils.context import context_copy_partial, context_get_outlet_events, context_merge
57+
from airflow.utils.context import context_copy_partial, context_merge
5858
from airflow.utils.file import get_unique_dag_module_name
59-
from airflow.utils.operator_helpers import ExecutionCallableRunner, KeywordParameters
60-
from airflow.utils.process_utils import execute_in_subprocess
59+
from airflow.utils.operator_helpers import KeywordParameters
60+
from airflow.utils.process_utils import execute_in_subprocess, execute_in_subprocess_with_kwargs
6161
from airflow.utils.session import create_session
6262

6363
log = logging.getLogger(__name__)
6464

6565
AIRFLOW_VERSION = Version(airflow_version)
66+
AIRFLOW_V_2_10_PLUS = Version(AIRFLOW_VERSION.base_version) >= Version("2.10.0")
6667
AIRFLOW_V_3_0_PLUS = Version(AIRFLOW_VERSION.base_version) >= Version("3.0.0")
6768

6869
if TYPE_CHECKING:
@@ -187,7 +188,15 @@ def __init__(
187188
def execute(self, context: Context) -> Any:
188189
context_merge(context, self.op_kwargs, templates_dict=self.templates_dict)
189190
self.op_kwargs = self.determine_kwargs(context)
190-
self._asset_events = context_get_outlet_events(context)
191+
192+
if AIRFLOW_V_3_0_PLUS:
193+
from airflow.utils.context import context_get_outlet_events
194+
195+
self._asset_events = context_get_outlet_events(context)
196+
elif AIRFLOW_V_2_10_PLUS:
197+
from airflow.utils.context import context_get_outlet_events
198+
199+
self._dataset_events = context_get_outlet_events(context)
191200

192201
return_value = self.execute_callable()
193202
if self.show_return_value_in_logs:
@@ -206,7 +215,15 @@ def execute_callable(self) -> Any:
206215
207216
:return: the return value of the call.
208217
"""
209-
runner = ExecutionCallableRunner(self.python_callable, self._asset_events, logger=self.log)
218+
try:
219+
from airflow.utils.operator_helpers import ExecutionCallableRunner
220+
221+
asset_events = self._asset_events if AIRFLOW_V_3_0_PLUS else self._dataset_events
222+
223+
runner = ExecutionCallableRunner(self.python_callable, asset_events, logger=self.log)
224+
except ImportError:
225+
# Handle Pre Airflow 3.10 case where ExecutionCallableRunner was not available
226+
return self.python_callable(*self.op_args, **self.op_kwargs)
210227
return runner.run(*self.op_args, **self.op_kwargs)
211228

212229

@@ -551,18 +568,25 @@ def _execute_python_callable_in_subprocess(self, python_path: Path):
551568
env_vars.update(self.env_vars)
552569

553570
try:
554-
execute_in_subprocess(
555-
cmd=[
556-
os.fspath(python_path),
557-
os.fspath(script_path),
558-
os.fspath(input_path),
559-
os.fspath(output_path),
560-
os.fspath(string_args_path),
561-
os.fspath(termination_log_path),
562-
os.fspath(airflow_context_path),
563-
],
564-
env=env_vars,
565-
)
571+
cmd: list[str] = [
572+
os.fspath(python_path),
573+
os.fspath(script_path),
574+
os.fspath(input_path),
575+
os.fspath(output_path),
576+
os.fspath(string_args_path),
577+
os.fspath(termination_log_path),
578+
os.fspath(airflow_context_path),
579+
]
580+
if AIRFLOW_V_2_10_PLUS:
581+
execute_in_subprocess(
582+
cmd=cmd,
583+
env=env_vars,
584+
)
585+
else:
586+
execute_in_subprocess_with_kwargs(
587+
cmd=cmd,
588+
env=env_vars,
589+
)
566590
except subprocess.CalledProcessError as e:
567591
if e.returncode in self.skip_on_exit_code:
568592
raise AirflowSkipException(f"Process exited with code {e.returncode}. Skipping.")

‎providers/src/airflow/providers/standard/sensors/date_time.py‎

Lines changed: 21 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -18,10 +18,27 @@
1818
from __future__ import annotations
1919

2020
import datetime
21+
from dataclasses import dataclass
2122
from typing import TYPE_CHECKING, Any, NoReturn, Sequence
2223

24+
from airflow.providers.standard.operators.python import AIRFLOW_V_3_0_PLUS
2325
from airflow.sensors.base import BaseSensorOperator
24-
from airflow.triggers.base import StartTriggerArgs
26+
27+
try:
28+
from airflow.triggers.base import StartTriggerArgs
29+
except ImportError:
30+
# TODO: Remove this when min airflow version is 2.10.0 for standard provider
31+
@dataclass
32+
class StartTriggerArgs: # type: ignore[no-redef]
33+
"""Arguments required for start task execution from triggerer."""
34+
35+
trigger_cls: str
36+
next_method: str
37+
trigger_kwargs: dict[str, Any] | None = None
38+
next_kwargs: dict[str, Any] | None = None
39+
timeout: datetime.timedelta | None = None
40+
41+
2542
from airflow.triggers.temporal import DateTimeTrigger
2643
from airflow.utils import timezone
2744

@@ -125,7 +142,9 @@ def execute(self, context: Context) -> NoReturn:
125142
trigger=DateTimeTrigger(
126143
moment=timezone.parse(self.target_time),
127144
end_from_trigger=self.end_from_trigger,
128-
),
145+
)
146+
if AIRFLOW_V_3_0_PLUS
147+
else DateTimeTrigger(moment=timezone.parse(self.target_time)),
129148
)
130149

131150
def execute_complete(self, context: Context, event: Any = None) -> None:

‎providers/src/airflow/providers/standard/sensors/time.py‎

Lines changed: 21 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -18,10 +18,27 @@
1818
from __future__ import annotations
1919

2020
import datetime
21+
from dataclasses import dataclass
2122
from typing import TYPE_CHECKING, Any, NoReturn
2223

24+
from airflow.providers.standard.operators.python import AIRFLOW_V_3_0_PLUS
2325
from airflow.sensors.base import BaseSensorOperator
24-
from airflow.triggers.base import StartTriggerArgs
26+
27+
try:
28+
from airflow.triggers.base import StartTriggerArgs
29+
except ImportError:
30+
# TODO: Remove this when min airflow version is 2.10.0 for standard provider
31+
@dataclass
32+
class StartTriggerArgs: # type: ignore[no-redef]
33+
"""Arguments required for start task execution from triggerer."""
34+
35+
trigger_cls: str
36+
next_method: str
37+
trigger_kwargs: dict[str, Any] | None = None
38+
next_kwargs: dict[str, Any] | None = None
39+
timeout: datetime.timedelta | None = None
40+
41+
2542
from airflow.triggers.temporal import DateTimeTrigger
2643
from airflow.utils import timezone
2744

@@ -102,7 +119,9 @@ def __init__(
102119

103120
def execute(self, context: Context) -> NoReturn:
104121
self.defer(
105-
trigger=DateTimeTrigger(moment=self.target_datetime, end_from_trigger=self.end_from_trigger),
122+
trigger=DateTimeTrigger(moment=self.target_datetime, end_from_trigger=self.end_from_trigger)
123+
if AIRFLOW_V_3_0_PLUS
124+
else DateTimeTrigger(moment=self.target_datetime),
106125
method_name="execute_complete",
107126
)
108127

‎providers/src/airflow/providers/standard/sensors/time_delta.py‎

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -27,6 +27,8 @@
2727
from airflow.triggers.temporal import DateTimeTrigger, TimeDeltaTrigger
2828
from airflow.utils import timezone
2929

30+
from tests_common.test_utils.compat import AIRFLOW_V_3_0_PLUS
31+
3032
if TYPE_CHECKING:
3133
from airflow.utils.context import Context
3234

@@ -81,7 +83,10 @@ def execute(self, context: Context) -> bool | NoReturn:
8183
# If the target datetime is in the past, return immediately
8284
return True
8385
try:
84-
trigger = DateTimeTrigger(moment=target_dttm, end_from_trigger=self.end_from_trigger)
86+
if AIRFLOW_V_3_0_PLUS:
87+
trigger = DateTimeTrigger(moment=target_dttm, end_from_trigger=self.end_from_trigger)
88+
else:
89+
trigger = DateTimeTrigger(moment=target_dttm)
8590
except (TypeError, ValueError) as e:
8691
if self.soft_fail:
8792
raise AirflowSkipException("Skipping due to soft_fail is set to True.") from e

‎providers/tests/openlineage/utils/test_utils.py‎

Lines changed: 1 addition & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -43,14 +43,11 @@
4343
from airflow.utils.task_group import TaskGroup
4444
from airflow.utils.types import DagRunType
4545

46-
from tests_common.test_utils.compat import AIRFLOW_V_2_10_PLUS, BashOperator, PythonOperator
46+
from tests_common.test_utils.compat import BashOperator, PythonOperator
4747
from tests_common.test_utils.mock_operators import MockOperator
4848

4949
BASH_OPERATOR_PATH = "airflow.providers.standard.operators.bash"
5050
PYTHON_OPERATOR_PATH = "airflow.providers.standard.operators.python"
51-
if not AIRFLOW_V_2_10_PLUS:
52-
BASH_OPERATOR_PATH = "airflow.operators.bash"
53-
PYTHON_OPERATOR_PATH = "airflow.operators.python"
5451

5552

5653
class CustomOperatorForTest(BashOperator):

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

Lines changed: 27 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -51,6 +51,7 @@
5151
from airflow.models.taskinstance import TaskInstance, clear_task_instances, set_current_context
5252
from airflow.operators.empty import EmptyOperator
5353
from airflow.providers.standard.operators.python import (
54+
AIRFLOW_V_2_10_PLUS,
5455
BranchExternalPythonOperator,
5556
BranchPythonOperator,
5657
BranchPythonVirtualenvOperator,
@@ -509,7 +510,7 @@ def f():
509510
ti = self.create_ti(f)
510511
with pytest.raises(
511512
AirflowException,
512-
match="'branch_task_ids' expected all task IDs are strings.",
513+
match=r"'branch_task_ids'.*task.*",
513514
):
514515
ti.run()
515516

@@ -518,7 +519,9 @@ def f():
518519
return "some_task_id"
519520

520521
ti = self.create_ti(f)
521-
with pytest.raises(AirflowException, match="Invalid tasks found: {'some_task_id'}"):
522+
with pytest.raises(
523+
AirflowException, match=r"Invalid tasks found: {\(False, 'bool'\)}.|'branch_task_ids'.*task.*"
524+
):
522525
ti.run()
523526

524527
@pytest.mark.skip_if_database_isolation_mode # tests pure logic with run() method, can not run in isolation mode
@@ -903,9 +906,12 @@ def test_virtualenv_serializable_context_fields(self, create_task_instance):
903906
"ti",
904907
"var", # Accessor for Variable; var->json and var->value.
905908
"conn", # Accessor for Connection.
906-
"inlet_events", # Accessor for inlet AssetEvent.
907-
"outlet_events", # Accessor for outlet AssetEvent.
908909
]
910+
if AIRFLOW_V_2_10_PLUS:
911+
intentionally_excluded_context_keys.extend(
912+
# Accessors for inlet_events and outlet_events
913+
["inlet_events", "outlet_events"]
914+
)
909915

910916
ti = create_task_instance(dag_id=self.dag_id, task_id=self.task_id, schedule=None)
911917
context = ti.get_template_context()
@@ -1627,21 +1633,25 @@ def f(a, b, c=False, d=False):
16271633
else:
16281634
raise RuntimeError
16291635

1630-
with pytest.raises(AirflowException, match=r"Invalid tasks found: {\((True|False), 'bool'\)}"):
1636+
with pytest.raises(
1637+
AirflowException, match=r"Invalid tasks found: {\(False, 'bool'\)}.|'branch_task_ids'.*task.*"
1638+
):
16311639
self.run_as_task(f, op_args=[0, 1], op_kwargs={"c": True})
16321640

16331641
def test_return_false(self):
16341642
def f():
16351643
return False
16361644

1637-
with pytest.raises(AirflowException, match=r"Invalid tasks found: {\(False, 'bool'\)}."):
1645+
with pytest.raises(
1646+
AirflowException, match=r"Invalid tasks found: {\(False, 'bool'\)}.|'branch_task_ids'.*task.*"
1647+
):
16381648
self.run_as_task(f)
16391649

16401650
def test_context(self):
16411651
def f(templates_dict):
16421652
return templates_dict["ds"]
16431653

1644-
with pytest.raises(AirflowException, match="Invalid tasks found:"):
1654+
with pytest.raises(AirflowException, match="Invalid tasks found:|'branch_task_ids'.*task.*"):
16451655
self.run_as_task(f, templates_dict={"ds": "{{ ds }}"})
16461656

16471657
def test_environment_variables(self):
@@ -1652,7 +1662,7 @@ def f():
16521662

16531663
with pytest.raises(
16541664
AirflowException,
1655-
match=r"'branch_task_ids' must contain only valid task_ids. Invalid tasks found: {'ABCDE'}",
1665+
match=r"'branch_task_ids'.*task.*",
16561666
):
16571667
self.run_as_task(f, env_vars={"MY_ENV_VAR": "ABCDE"})
16581668

@@ -1666,7 +1676,7 @@ def f():
16661676

16671677
with pytest.raises(
16681678
AirflowException,
1669-
match=r"'branch_task_ids' must contain only valid task_ids. Invalid tasks found: {'QWERT'}",
1679+
match=r"'branch_task_ids'.*task.*",
16701680
):
16711681
self.run_as_task(f, inherit_env=True)
16721682

@@ -1691,7 +1701,7 @@ def f():
16911701

16921702
with pytest.raises(
16931703
AirflowException,
1694-
match=r"'branch_task_ids' must contain only valid task_ids. Invalid tasks found: {'EFGHI'}",
1704+
match=r"'branch_task_ids'.*task.*",
16951705
):
16961706
self.run_as_task(f, env_vars={"MY_ENV_VAR": "EFGHI"}, inherit_env=True)
16971707

@@ -1706,7 +1716,9 @@ def test_with_no_caching(self):
17061716
def f():
17071717
return False
17081718

1709-
with pytest.raises(AirflowException, match=r"Invalid tasks found: {\(False, 'bool'\)}."):
1719+
with pytest.raises(
1720+
AirflowException, match=r"Invalid tasks found: {\(False, 'bool'\)}.|'branch_task_ids'.*task.*"
1721+
):
17101722
self.run_as_task(f, do_not_use_caching=True)
17111723

17121724
def test_with_dag_run(self):
@@ -1827,7 +1839,7 @@ def f():
18271839
ti = self.create_ti(f)
18281840
with pytest.raises(
18291841
AirflowException,
1830-
match="'branch_task_ids' expected all task IDs are strings.",
1842+
match=r"'branch_task_ids'.*task.*",
18311843
):
18321844
ti.run()
18331845

@@ -1836,7 +1848,9 @@ def f():
18361848
return "some_task_id"
18371849

18381850
ti = self.create_ti(f)
1839-
with pytest.raises(AirflowException, match="Invalid tasks found: {'some_task_id'}"):
1851+
with pytest.raises(
1852+
AirflowException, match=r"Invalid tasks found: {\(False, 'bool'\)}.|'branch_task_ids'.*task.*"
1853+
):
18401854
ti.run()
18411855

18421856

0 commit comments

Comments
 (0)