apache-airflow-providers-standard 1.2.0__py3-none-any.whl → 1.3.0rc1__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (30) hide show
  1. airflow/providers/standard/__init__.py +1 -1
  2. airflow/providers/standard/example_dags/__init__.py +16 -0
  3. airflow/providers/standard/example_dags/example_bash_decorator.py +114 -0
  4. airflow/providers/standard/example_dags/example_bash_operator.py +74 -0
  5. airflow/providers/standard/example_dags/example_branch_datetime_operator.py +105 -0
  6. airflow/providers/standard/example_dags/example_branch_day_of_week_operator.py +61 -0
  7. airflow/providers/standard/example_dags/example_branch_operator.py +166 -0
  8. airflow/providers/standard/example_dags/example_branch_operator_decorator.py +142 -0
  9. airflow/providers/standard/example_dags/example_external_task_child_deferrable.py +34 -0
  10. airflow/providers/standard/example_dags/example_external_task_marker_dag.py +98 -0
  11. airflow/providers/standard/example_dags/example_external_task_parent_deferrable.py +64 -0
  12. airflow/providers/standard/example_dags/example_latest_only.py +40 -0
  13. airflow/providers/standard/example_dags/example_python_decorator.py +132 -0
  14. airflow/providers/standard/example_dags/example_python_operator.py +147 -0
  15. airflow/providers/standard/example_dags/example_sensor_decorator.py +66 -0
  16. airflow/providers/standard/example_dags/example_sensors.py +135 -0
  17. airflow/providers/standard/example_dags/example_short_circuit_decorator.py +60 -0
  18. airflow/providers/standard/example_dags/example_short_circuit_operator.py +66 -0
  19. airflow/providers/standard/example_dags/example_trigger_controller_dag.py +46 -0
  20. airflow/providers/standard/example_dags/sql/__init__.py +16 -0
  21. airflow/providers/standard/example_dags/sql/sample.sql +24 -0
  22. airflow/providers/standard/operators/python.py +15 -9
  23. airflow/providers/standard/sensors/date_time.py +10 -4
  24. airflow/providers/standard/sensors/external_task.py +7 -6
  25. airflow/providers/standard/sensors/time.py +52 -40
  26. airflow/providers/standard/sensors/time_delta.py +47 -20
  27. {apache_airflow_providers_standard-1.2.0.dist-info → apache_airflow_providers_standard-1.3.0rc1.dist-info}/METADATA +7 -7
  28. {apache_airflow_providers_standard-1.2.0.dist-info → apache_airflow_providers_standard-1.3.0rc1.dist-info}/RECORD +30 -10
  29. {apache_airflow_providers_standard-1.2.0.dist-info → apache_airflow_providers_standard-1.3.0rc1.dist-info}/WHEEL +0 -0
  30. {apache_airflow_providers_standard-1.2.0.dist-info → apache_airflow_providers_standard-1.3.0rc1.dist-info}/entry_points.txt +0 -0
@@ -29,7 +29,7 @@ from airflow import __version__ as airflow_version
29
29
 
30
30
  __all__ = ["__version__"]
31
31
 
32
- __version__ = "1.2.0"
32
+ __version__ = "1.3.0"
33
33
 
34
34
  if packaging.version.parse(packaging.version.parse(airflow_version).base_version) < packaging.version.parse(
35
35
  "2.10.0"
@@ -0,0 +1,16 @@
1
+ # Licensed to the Apache Software Foundation (ASF) under one
2
+ # or more contributor license agreements. See the NOTICE file
3
+ # distributed with this work for additional information
4
+ # regarding copyright ownership. The ASF licenses this file
5
+ # to you under the Apache License, Version 2.0 (the
6
+ # "License"); you may not use this file except in compliance
7
+ # with the License. You may obtain a copy of the License at
8
+ #
9
+ # http://www.apache.org/licenses/LICENSE-2.0
10
+ #
11
+ # Unless required by applicable law or agreed to in writing,
12
+ # software distributed under the License is distributed on an
13
+ # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14
+ # KIND, either express or implied. See the License for the
15
+ # specific language governing permissions and limitations
16
+ # under the License.
@@ -0,0 +1,114 @@
1
+ # Licensed to the Apache Software Foundation (ASF) under one
2
+ # or more contributor license agreements. See the NOTICE file
3
+ # distributed with this work for additional information
4
+ # regarding copyright ownership. The ASF licenses this file
5
+ # to you under the Apache License, Version 2.0 (the
6
+ # "License"); you may not use this file except in compliance
7
+ # with the License. You may obtain a copy of the License at
8
+ #
9
+ # http://www.apache.org/licenses/LICENSE-2.0
10
+ #
11
+ # Unless required by applicable law or agreed to in writing,
12
+ # software distributed under the License is distributed on an
13
+ # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14
+ # KIND, either express or implied. See the License for the
15
+ # specific language governing permissions and limitations
16
+ # under the License.
17
+
18
+ from __future__ import annotations
19
+
20
+ import pendulum
21
+
22
+ from airflow.exceptions import AirflowSkipException
23
+ from airflow.providers.standard.operators.empty import EmptyOperator
24
+ from airflow.providers.standard.utils.weekday import WeekDay
25
+ from airflow.sdk import chain, dag, task
26
+ from airflow.utils.trigger_rule import TriggerRule
27
+
28
+
29
+ @dag(schedule=None, start_date=pendulum.datetime(2023, 1, 1, tz="UTC"), catchup=False)
30
+ def example_bash_decorator():
31
+ @task.bash
32
+ def run_me(sleep_seconds: int, task_instance_key_str: str) -> str:
33
+ return f"echo {task_instance_key_str} && sleep {sleep_seconds}"
34
+
35
+ run_me_loop = [run_me.override(task_id=f"runme_{i}")(sleep_seconds=i) for i in range(3)]
36
+
37
+ # [START howto_decorator_bash]
38
+ @task.bash
39
+ def run_after_loop() -> str:
40
+ return "echo https://airflow.apache.org/"
41
+
42
+ run_this = run_after_loop()
43
+ # [END howto_decorator_bash]
44
+
45
+ # [START howto_decorator_bash_template]
46
+ @task.bash
47
+ def also_run_this() -> str:
48
+ return 'echo "ti_key={{ task_instance_key_str }}"'
49
+
50
+ also_this = also_run_this()
51
+ # [END howto_decorator_bash_template]
52
+
53
+ # [START howto_decorator_bash_context_vars]
54
+ @task.bash
55
+ def also_run_this_again(task_instance_key_str) -> str:
56
+ return f'echo "ti_key={task_instance_key_str}"'
57
+
58
+ also_this_again = also_run_this_again()
59
+ # [END howto_decorator_bash_context_vars]
60
+
61
+ # [START howto_decorator_bash_skip]
62
+ @task.bash
63
+ def this_will_skip() -> str:
64
+ return 'echo "hello world"; exit 99;'
65
+
66
+ this_skips = this_will_skip()
67
+ # [END howto_decorator_bash_skip]
68
+
69
+ run_this_last = EmptyOperator(task_id="run_this_last", trigger_rule=TriggerRule.ALL_DONE)
70
+
71
+ # [START howto_decorator_bash_conditional]
72
+ @task.bash
73
+ def sleep_in(day: str) -> str:
74
+ if day in (WeekDay.SATURDAY, WeekDay.SUNDAY):
75
+ return f"sleep {60 * 60}"
76
+ raise AirflowSkipException("No sleeping in today!")
77
+
78
+ sleep_in(day="{{ dag_run.logical_date.strftime('%A').lower() }}")
79
+ # [END howto_decorator_bash_conditional]
80
+
81
+ # [START howto_decorator_bash_parametrize]
82
+ @task.bash(env={"BASE_DIR": "{{ dag_run.logical_date.strftime('%Y/%m/%d') }}"}, append_env=True)
83
+ def make_dynamic_dirs(new_dirs: str) -> str:
84
+ return f"mkdir -p $AIRFLOW_HOME/$BASE_DIR/{new_dirs}"
85
+
86
+ make_dynamic_dirs(new_dirs="foo/bar/baz")
87
+ # [END howto_decorator_bash_parametrize]
88
+
89
+ # [START howto_decorator_bash_build_cmd]
90
+ def _get_files_in_cwd() -> list[str]:
91
+ from pathlib import Path
92
+
93
+ dir_contents = Path.cwd().glob("airflow-core/src/airflow/example_dags/*.py")
94
+ files = [str(elem) for elem in dir_contents if elem.is_file()]
95
+
96
+ return files
97
+
98
+ @task.bash
99
+ def get_file_stats() -> str:
100
+ from shlex import join
101
+
102
+ files = _get_files_in_cwd()
103
+ cmd = join(["stat", *files])
104
+
105
+ return cmd
106
+
107
+ get_file_stats()
108
+ # [END howto_decorator_bash_build_cmd]
109
+
110
+ chain(run_me_loop, run_this)
111
+ chain([also_this, also_this_again, this_skips, run_this], run_this_last)
112
+
113
+
114
+ example_bash_decorator()
@@ -0,0 +1,74 @@
1
+ #
2
+ # Licensed to the Apache Software Foundation (ASF) under one
3
+ # or more contributor license agreements. See the NOTICE file
4
+ # distributed with this work for additional information
5
+ # regarding copyright ownership. The ASF licenses this file
6
+ # to you under the Apache License, Version 2.0 (the
7
+ # "License"); you may not use this file except in compliance
8
+ # with the License. You may obtain a copy of the License at
9
+ #
10
+ # http://www.apache.org/licenses/LICENSE-2.0
11
+ #
12
+ # Unless required by applicable law or agreed to in writing,
13
+ # software distributed under the License is distributed on an
14
+ # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
15
+ # KIND, either express or implied. See the License for the
16
+ # specific language governing permissions and limitations
17
+ # under the License.
18
+ """Example DAG demonstrating the usage of the BashOperator."""
19
+
20
+ from __future__ import annotations
21
+
22
+ import datetime
23
+
24
+ import pendulum
25
+
26
+ from airflow.providers.standard.operators.bash import BashOperator
27
+ from airflow.providers.standard.operators.empty import EmptyOperator
28
+ from airflow.sdk import DAG
29
+
30
+ with DAG(
31
+ dag_id="example_bash_operator",
32
+ schedule="0 0 * * *",
33
+ start_date=pendulum.datetime(2021, 1, 1, tz="UTC"),
34
+ catchup=False,
35
+ dagrun_timeout=datetime.timedelta(minutes=60),
36
+ tags=["example", "example2"],
37
+ params={"example_key": "example_value"},
38
+ ) as dag:
39
+ run_this_last = EmptyOperator(
40
+ task_id="run_this_last",
41
+ )
42
+
43
+ # [START howto_operator_bash]
44
+ run_this = BashOperator(
45
+ task_id="run_after_loop",
46
+ bash_command="echo https://airflow.apache.org/",
47
+ )
48
+ # [END howto_operator_bash]
49
+
50
+ run_this >> run_this_last
51
+
52
+ for i in range(3):
53
+ task = BashOperator(
54
+ task_id=f"runme_{i}",
55
+ bash_command='echo "{{ task_instance_key_str }}" && sleep 1',
56
+ )
57
+ task >> run_this
58
+
59
+ # [START howto_operator_bash_template]
60
+ also_run_this = BashOperator(
61
+ task_id="also_run_this",
62
+ bash_command='echo "ti_key={{ task_instance_key_str }}"',
63
+ )
64
+ # [END howto_operator_bash_template]
65
+ also_run_this >> run_this_last
66
+
67
+ # [START howto_operator_bash_skip]
68
+ this_will_skip = BashOperator(
69
+ task_id="this_will_skip",
70
+ bash_command='echo "hello world"; exit 99;',
71
+ dag=dag,
72
+ )
73
+ # [END howto_operator_bash_skip]
74
+ this_will_skip >> run_this_last
@@ -0,0 +1,105 @@
1
+ #
2
+ # Licensed to the Apache Software Foundation (ASF) under one
3
+ # or more contributor license agreements. See the NOTICE file
4
+ # distributed with this work for additional information
5
+ # regarding copyright ownership. The ASF licenses this file
6
+ # to you under the Apache License, Version 2.0 (the
7
+ # "License"); you may not use this file except in compliance
8
+ # with the License. You may obtain a copy of the License at
9
+ #
10
+ # http://www.apache.org/licenses/LICENSE-2.0
11
+ #
12
+ # Unless required by applicable law or agreed to in writing,
13
+ # software distributed under the License is distributed on an
14
+ # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
15
+ # KIND, either express or implied. See the License for the
16
+ # specific language governing permissions and limitations
17
+ # under the License.
18
+ """
19
+ Example DAG demonstrating the usage of DateTimeBranchOperator with datetime as well as time objects as
20
+ targets.
21
+ """
22
+
23
+ from __future__ import annotations
24
+
25
+ import pendulum
26
+
27
+ from airflow.providers.standard.operators.datetime import BranchDateTimeOperator
28
+ from airflow.providers.standard.operators.empty import EmptyOperator
29
+ from airflow.sdk import DAG
30
+
31
+ dag1 = DAG(
32
+ dag_id="example_branch_datetime_operator",
33
+ start_date=pendulum.datetime(2021, 1, 1, tz="UTC"),
34
+ catchup=False,
35
+ tags=["example"],
36
+ schedule="@daily",
37
+ )
38
+
39
+ # [START howto_branch_datetime_operator]
40
+ empty_task_11 = EmptyOperator(task_id="date_in_range", dag=dag1)
41
+ empty_task_21 = EmptyOperator(task_id="date_outside_range", dag=dag1)
42
+
43
+ cond1 = BranchDateTimeOperator(
44
+ task_id="datetime_branch",
45
+ follow_task_ids_if_true=["date_in_range"],
46
+ follow_task_ids_if_false=["date_outside_range"],
47
+ target_upper=pendulum.datetime(2020, 10, 10, 15, 0, 0),
48
+ target_lower=pendulum.datetime(2020, 10, 10, 14, 0, 0),
49
+ dag=dag1,
50
+ )
51
+
52
+ # Run empty_task_11 if cond1 executes between 2020-10-10 14:00:00 and 2020-10-10 15:00:00
53
+ cond1 >> [empty_task_11, empty_task_21]
54
+ # [END howto_branch_datetime_operator]
55
+
56
+
57
+ dag2 = DAG(
58
+ dag_id="example_branch_datetime_operator_2",
59
+ start_date=pendulum.datetime(2021, 1, 1, tz="UTC"),
60
+ catchup=False,
61
+ tags=["example"],
62
+ schedule="@daily",
63
+ )
64
+ # [START howto_branch_datetime_operator_next_day]
65
+ empty_task_12 = EmptyOperator(task_id="date_in_range", dag=dag2)
66
+ empty_task_22 = EmptyOperator(task_id="date_outside_range", dag=dag2)
67
+
68
+ cond2 = BranchDateTimeOperator(
69
+ task_id="datetime_branch",
70
+ follow_task_ids_if_true=["date_in_range"],
71
+ follow_task_ids_if_false=["date_outside_range"],
72
+ target_upper=pendulum.time(0, 0, 0),
73
+ target_lower=pendulum.time(15, 0, 0),
74
+ dag=dag2,
75
+ )
76
+
77
+ # Since target_lower happens after target_upper, target_upper will be moved to the following day
78
+ # Run empty_task_12 if cond2 executes between 15:00:00, and 00:00:00 of the following day
79
+ cond2 >> [empty_task_12, empty_task_22]
80
+ # [END howto_branch_datetime_operator_next_day]
81
+
82
+ dag3 = DAG(
83
+ dag_id="example_branch_datetime_operator_3",
84
+ start_date=pendulum.datetime(2021, 1, 1, tz="UTC"),
85
+ catchup=False,
86
+ tags=["example"],
87
+ schedule="@daily",
88
+ )
89
+ # [START howto_branch_datetime_operator_logical_date]
90
+ empty_task_13 = EmptyOperator(task_id="date_in_range", dag=dag3)
91
+ empty_task_23 = EmptyOperator(task_id="date_outside_range", dag=dag3)
92
+
93
+ cond3 = BranchDateTimeOperator(
94
+ task_id="datetime_branch",
95
+ use_task_logical_date=True,
96
+ follow_task_ids_if_true=["date_in_range"],
97
+ follow_task_ids_if_false=["date_outside_range"],
98
+ target_upper=pendulum.datetime(2020, 10, 10, 15, 0, 0),
99
+ target_lower=pendulum.datetime(2020, 10, 10, 14, 0, 0),
100
+ dag=dag3,
101
+ )
102
+
103
+ # Run empty_task_13 if cond3 executes between 2020-10-10 14:00:00 and 2020-10-10 15:00:00
104
+ cond3 >> [empty_task_13, empty_task_23]
105
+ # [END howto_branch_datetime_operator_logical_date]
@@ -0,0 +1,61 @@
1
+ #
2
+ # Licensed to the Apache Software Foundation (ASF) under one
3
+ # or more contributor license agreements. See the NOTICE file
4
+ # distributed with this work for additional information
5
+ # regarding copyright ownership. The ASF licenses this file
6
+ # to you under the Apache License, Version 2.0 (the
7
+ # "License"); you may not use this file except in compliance
8
+ # with the License. You may obtain a copy of the License at
9
+ #
10
+ # http://www.apache.org/licenses/LICENSE-2.0
11
+ #
12
+ # Unless required by applicable law or agreed to in writing,
13
+ # software distributed under the License is distributed on an
14
+ # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
15
+ # KIND, either express or implied. See the License for the
16
+ # specific language governing permissions and limitations
17
+ # under the License.
18
+ """
19
+ Example DAG demonstrating the usage of BranchDayOfWeekOperator.
20
+ """
21
+
22
+ from __future__ import annotations
23
+
24
+ import pendulum
25
+
26
+ from airflow.providers.standard.operators.empty import EmptyOperator
27
+ from airflow.providers.standard.operators.weekday import BranchDayOfWeekOperator
28
+ from airflow.providers.standard.utils.weekday import WeekDay
29
+ from airflow.sdk import DAG
30
+
31
+ with DAG(
32
+ dag_id="example_weekday_branch_operator",
33
+ start_date=pendulum.datetime(2021, 1, 1, tz="UTC"),
34
+ catchup=False,
35
+ tags=["example"],
36
+ schedule="@daily",
37
+ ) as dag:
38
+ # [START howto_operator_day_of_week_branch]
39
+ empty_task_1 = EmptyOperator(task_id="branch_true")
40
+ empty_task_2 = EmptyOperator(task_id="branch_false")
41
+ empty_task_3 = EmptyOperator(task_id="branch_weekend")
42
+ empty_task_4 = EmptyOperator(task_id="branch_mid_week")
43
+
44
+ branch = BranchDayOfWeekOperator(
45
+ task_id="make_choice",
46
+ follow_task_ids_if_true="branch_true",
47
+ follow_task_ids_if_false="branch_false",
48
+ week_day="Monday",
49
+ )
50
+ branch_weekend = BranchDayOfWeekOperator(
51
+ task_id="make_weekend_choice",
52
+ follow_task_ids_if_true="branch_weekend",
53
+ follow_task_ids_if_false="branch_mid_week",
54
+ week_day={WeekDay.SATURDAY, WeekDay.SUNDAY},
55
+ )
56
+
57
+ # Run empty_task_1 if branch executes on Monday, empty_task_2 otherwise
58
+ branch >> [empty_task_1, empty_task_2]
59
+ # Run empty_task_3 if it's a weekend, empty_task_4 otherwise
60
+ empty_task_2 >> branch_weekend >> [empty_task_3, empty_task_4]
61
+ # [END howto_operator_day_of_week_branch]
@@ -0,0 +1,166 @@
1
+ #
2
+ # Licensed to the Apache Software Foundation (ASF) under one
3
+ # or more contributor license agreements. See the NOTICE file
4
+ # distributed with this work for additional information
5
+ # regarding copyright ownership. The ASF licenses this file
6
+ # to you under the Apache License, Version 2.0 (the
7
+ # "License"); you may not use this file except in compliance
8
+ # with the License. You may obtain a copy of the License at
9
+ #
10
+ # http://www.apache.org/licenses/LICENSE-2.0
11
+ #
12
+ # Unless required by applicable law or agreed to in writing,
13
+ # software distributed under the License is distributed on an
14
+ # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
15
+ # KIND, either express or implied. See the License for the
16
+ # specific language governing permissions and limitations
17
+ # under the License.
18
+ """Example DAG demonstrating the usage of the Classic branching Python operators.
19
+
20
+ It is showcasing the basic BranchPythonOperator and its sisters BranchExternalPythonOperator
21
+ and BranchPythonVirtualenvOperator."""
22
+
23
+ from __future__ import annotations
24
+
25
+ import random
26
+ import sys
27
+ import tempfile
28
+ from pathlib import Path
29
+
30
+ import pendulum
31
+
32
+ from airflow.providers.standard.operators.empty import EmptyOperator
33
+ from airflow.providers.standard.operators.python import (
34
+ BranchExternalPythonOperator,
35
+ BranchPythonOperator,
36
+ BranchPythonVirtualenvOperator,
37
+ ExternalPythonOperator,
38
+ PythonOperator,
39
+ PythonVirtualenvOperator,
40
+ )
41
+ from airflow.sdk import DAG, Label
42
+ from airflow.utils.trigger_rule import TriggerRule
43
+
44
+ PATH_TO_PYTHON_BINARY = sys.executable
45
+
46
+ with DAG(
47
+ dag_id="example_branch_operator",
48
+ start_date=pendulum.datetime(2021, 1, 1, tz="UTC"),
49
+ catchup=False,
50
+ schedule="@daily",
51
+ tags=["example", "example2"],
52
+ ) as dag:
53
+ run_this_first = EmptyOperator(
54
+ task_id="run_this_first",
55
+ )
56
+
57
+ options = ["a", "b", "c", "d"]
58
+
59
+ # Example branching on standard Python tasks
60
+
61
+ # [START howto_operator_branch_python]
62
+ branching = BranchPythonOperator(
63
+ task_id="branching",
64
+ python_callable=lambda: f"branch_{random.choice(options)}",
65
+ )
66
+ # [END howto_operator_branch_python]
67
+ run_this_first >> branching
68
+
69
+ join = EmptyOperator(
70
+ task_id="join",
71
+ trigger_rule=TriggerRule.NONE_FAILED_MIN_ONE_SUCCESS,
72
+ )
73
+
74
+ for option in options:
75
+ t = PythonOperator(
76
+ task_id=f"branch_{option}",
77
+ python_callable=lambda: print("Hello World"),
78
+ )
79
+
80
+ empty_follow = EmptyOperator(
81
+ task_id="follow_" + option,
82
+ )
83
+
84
+ # Label is optional here, but it can help identify more complex branches
85
+ branching >> Label(option) >> t >> empty_follow >> join
86
+
87
+ # Example the same with external Python calls
88
+
89
+ # [START howto_operator_branch_ext_py]
90
+ def branch_with_external_python(choices):
91
+ import random
92
+
93
+ return f"ext_py_{random.choice(choices)}"
94
+
95
+ branching_ext_py = BranchExternalPythonOperator(
96
+ task_id="branching_ext_python",
97
+ python=PATH_TO_PYTHON_BINARY,
98
+ python_callable=branch_with_external_python,
99
+ op_args=[options],
100
+ )
101
+ # [END howto_operator_branch_ext_py]
102
+ join >> branching_ext_py
103
+
104
+ join_ext_py = EmptyOperator(
105
+ task_id="join_ext_python",
106
+ trigger_rule=TriggerRule.NONE_FAILED_MIN_ONE_SUCCESS,
107
+ )
108
+
109
+ def hello_world_with_external_python():
110
+ print("Hello World from external Python")
111
+
112
+ for option in options:
113
+ t = ExternalPythonOperator(
114
+ task_id=f"ext_py_{option}",
115
+ python=PATH_TO_PYTHON_BINARY,
116
+ python_callable=hello_world_with_external_python,
117
+ )
118
+
119
+ # Label is optional here, but it can help identify more complex branches
120
+ branching_ext_py >> Label(option) >> t >> join_ext_py
121
+
122
+ # Example the same with Python virtual environments
123
+
124
+ # [START howto_operator_branch_virtualenv]
125
+ # Note: Passing a caching dir allows to keep the virtual environment over multiple runs
126
+ # Run the example a second time and see that it reuses it and is faster.
127
+ VENV_CACHE_PATH = Path(tempfile.gettempdir())
128
+
129
+ def branch_with_venv(choices):
130
+ import random
131
+
132
+ import numpy as np
133
+
134
+ print(f"Some numpy stuff: {np.arange(6)}")
135
+ return f"venv_{random.choice(choices)}"
136
+
137
+ branching_venv = BranchPythonVirtualenvOperator(
138
+ task_id="branching_venv",
139
+ requirements=["numpy~=1.26.0"],
140
+ venv_cache_path=VENV_CACHE_PATH,
141
+ python_callable=branch_with_venv,
142
+ op_args=[options],
143
+ )
144
+ # [END howto_operator_branch_virtualenv]
145
+ join_ext_py >> branching_venv
146
+
147
+ join_venv = EmptyOperator(
148
+ task_id="join_venv",
149
+ trigger_rule=TriggerRule.NONE_FAILED_MIN_ONE_SUCCESS,
150
+ )
151
+
152
+ def hello_world_with_venv():
153
+ import numpy as np
154
+
155
+ print(f"Hello World with some numpy stuff: {np.arange(6)}")
156
+
157
+ for option in options:
158
+ t = PythonVirtualenvOperator(
159
+ task_id=f"venv_{option}",
160
+ requirements=["numpy~=1.26.0"],
161
+ venv_cache_path=VENV_CACHE_PATH,
162
+ python_callable=hello_world_with_venv,
163
+ )
164
+
165
+ # Label is optional here, but it can help identify more complex branches
166
+ branching_venv >> Label(option) >> t >> join_venv
@@ -0,0 +1,142 @@
1
+ #
2
+ # Licensed to the Apache Software Foundation (ASF) under one
3
+ # or more contributor license agreements. See the NOTICE file
4
+ # distributed with this work for additional information
5
+ # regarding copyright ownership. The ASF licenses this file
6
+ # to you under the Apache License, Version 2.0 (the
7
+ # "License"); you may not use this file except in compliance
8
+ # with the License. You may obtain a copy of the License at
9
+ #
10
+ # http://www.apache.org/licenses/LICENSE-2.0
11
+ #
12
+ # Unless required by applicable law or agreed to in writing,
13
+ # software distributed under the License is distributed on an
14
+ # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
15
+ # KIND, either express or implied. See the License for the
16
+ # specific language governing permissions and limitations
17
+ # under the License.
18
+ """Example DAG demonstrating the usage of the branching TaskFlow API decorators.
19
+
20
+ It shows how to use standard Python ``@task.branch`` as well as the external Python
21
+ version ``@task.branch_external_python`` which calls an external Python interpreter and
22
+ the ``@task.branch_virtualenv`` which builds a temporary Python virtual environment.
23
+ """
24
+
25
+ from __future__ import annotations
26
+
27
+ import random
28
+ import sys
29
+ import tempfile
30
+
31
+ import pendulum
32
+
33
+ from airflow.providers.standard.operators.empty import EmptyOperator
34
+ from airflow.sdk import DAG, Label, task
35
+ from airflow.utils.trigger_rule import TriggerRule
36
+
37
+ PATH_TO_PYTHON_BINARY = sys.executable
38
+
39
+ with DAG(
40
+ dag_id="example_branch_python_operator_decorator",
41
+ start_date=pendulum.datetime(2021, 1, 1, tz="UTC"),
42
+ catchup=False,
43
+ schedule="@daily",
44
+ tags=["example", "example2"],
45
+ ) as dag:
46
+ run_this_first = EmptyOperator(task_id="run_this_first")
47
+
48
+ options = ["a", "b", "c", "d"]
49
+
50
+ # Example branching on standard Python tasks
51
+
52
+ # [START howto_operator_branch_python]
53
+ @task.branch()
54
+ def branching(choices: list[str]) -> str:
55
+ return f"branch_{random.choice(choices)}"
56
+
57
+ # [END howto_operator_branch_python]
58
+
59
+ random_choice_instance = branching(choices=options)
60
+
61
+ run_this_first >> random_choice_instance
62
+
63
+ join = EmptyOperator(task_id="join", trigger_rule=TriggerRule.NONE_FAILED_MIN_ONE_SUCCESS)
64
+
65
+ for option in options:
66
+
67
+ @task(task_id=f"branch_{option}")
68
+ def some_task():
69
+ print("doing something in Python")
70
+
71
+ t = some_task()
72
+ empty = EmptyOperator(task_id=f"follow_{option}")
73
+
74
+ # Label is optional here, but it can help identify more complex branches
75
+ random_choice_instance >> Label(option) >> t >> empty >> join
76
+
77
+ # Example the same with external Python calls
78
+
79
+ # [START howto_operator_branch_ext_py]
80
+ @task.branch_external_python(python=PATH_TO_PYTHON_BINARY)
81
+ def branching_ext_python(choices) -> str:
82
+ import random
83
+
84
+ return f"ext_py_{random.choice(choices)}"
85
+
86
+ # [END howto_operator_branch_ext_py]
87
+
88
+ random_choice_ext_py = branching_ext_python(choices=options)
89
+
90
+ join >> random_choice_ext_py
91
+
92
+ join_ext_py = EmptyOperator(task_id="join_ext_py", trigger_rule=TriggerRule.NONE_FAILED_MIN_ONE_SUCCESS)
93
+
94
+ for option in options:
95
+
96
+ @task.external_python(task_id=f"ext_py_{option}", python=PATH_TO_PYTHON_BINARY)
97
+ def some_ext_py_task():
98
+ print("doing something in external Python")
99
+
100
+ t = some_ext_py_task()
101
+
102
+ # Label is optional here, but it can help identify more complex branches
103
+ random_choice_ext_py >> Label(option) >> t >> join_ext_py
104
+
105
+ # Example the same with Python virtual environments
106
+
107
+ # [START howto_operator_branch_virtualenv]
108
+ # Note: Passing a caching dir allows to keep the virtual environment over multiple runs
109
+ # Run the example a second time and see that it reuses it and is faster.
110
+ VENV_CACHE_PATH = tempfile.gettempdir()
111
+
112
+ @task.branch_virtualenv(requirements=["numpy~=1.26.0"], venv_cache_path=VENV_CACHE_PATH)
113
+ def branching_virtualenv(choices) -> str:
114
+ import random
115
+
116
+ import numpy as np
117
+
118
+ print(f"Some numpy stuff: {np.arange(6)}")
119
+ return f"venv_{random.choice(choices)}"
120
+
121
+ # [END howto_operator_branch_virtualenv]
122
+
123
+ random_choice_venv = branching_virtualenv(choices=options)
124
+
125
+ join_ext_py >> random_choice_venv
126
+
127
+ join_venv = EmptyOperator(task_id="join_venv", trigger_rule=TriggerRule.NONE_FAILED_MIN_ONE_SUCCESS)
128
+
129
+ for option in options:
130
+
131
+ @task.virtualenv(
132
+ task_id=f"venv_{option}", requirements=["numpy~=1.26.0"], venv_cache_path=VENV_CACHE_PATH
133
+ )
134
+ def some_venv_task():
135
+ import numpy as np
136
+
137
+ print(f"Some numpy stuff: {np.arange(6)}")
138
+
139
+ t = some_venv_task()
140
+
141
+ # Label is optional here, but it can help identify more complex branches
142
+ random_choice_venv >> Label(option) >> t >> join_venv