| # |
| # 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 |
| # regarding copyright ownership. The ASF licenses this file |
| # to you under the Apache License, Version 2.0 (the |
| # "License"); you may not use this file except in compliance |
| # with the License. You may obtain a copy of the License at |
| # |
| # http://www.apache.org/licenses/LICENSE-2.0 |
| # |
| # Unless required by applicable law or agreed to in writing, |
| # software distributed under the License is distributed on an |
| # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY |
| # KIND, either express or implied. See the License for the |
| # specific language governing permissions and limitations |
| # under the License. |
| from __future__ import annotations |
| |
| from unittest import mock |
| from unittest.mock import Mock |
| |
| import pytest |
| |
| import airflow |
| from airflow.exceptions import AirflowException |
| from airflow.models.dag import DAG |
| from airflow.models.dagrun import DagRun |
| from airflow.models.taskinstance import TaskInstance |
| from airflow.operators.empty import EmptyOperator |
| from airflow.operators.subdag import SkippedStatePropagationOptions, SubDagOperator |
| from airflow.utils.session import create_session |
| from airflow.utils.state import State |
| from airflow.utils.timezone import datetime |
| from airflow.utils.types import DagRunType |
| from tests.test_utils.db import clear_db_runs |
| |
| pytestmark = pytest.mark.db_test |
| |
| DEFAULT_DATE = datetime(2016, 1, 1) |
| |
| default_args = {"start_date": DEFAULT_DATE} |
| |
| |
| class TestSubDagOperator: |
| def setup_method(self): |
| clear_db_runs() |
| self.dag_run_running = DagRun() |
| self.dag_run_running.state = State.RUNNING |
| self.dag_run_success = DagRun() |
| self.dag_run_success.state = State.SUCCESS |
| self.dag_run_failed = DagRun() |
| self.dag_run_failed.state = State.FAILED |
| |
| def teardown_class(self): |
| clear_db_runs() |
| |
| def test_subdag_name(self): |
| """ |
| Subdag names must be {parent_dag}.{subdag task} |
| """ |
| dag = DAG("parent", default_args=default_args) |
| subdag_good = DAG("parent.test", default_args=default_args) |
| subdag_bad1 = DAG("parent.bad", default_args=default_args) |
| subdag_bad2 = DAG("bad.test", default_args=default_args) |
| subdag_bad3 = DAG("bad.bad", default_args=default_args) |
| |
| SubDagOperator(task_id="test", dag=dag, subdag=subdag_good) |
| with pytest.raises(AirflowException): |
| SubDagOperator(task_id="test", dag=dag, subdag=subdag_bad1) |
| with pytest.raises(AirflowException): |
| SubDagOperator(task_id="test", dag=dag, subdag=subdag_bad2) |
| with pytest.raises(AirflowException): |
| SubDagOperator(task_id="test", dag=dag, subdag=subdag_bad3) |
| |
| def test_subdag_in_context_manager(self): |
| """ |
| Creating a sub DAG within a main DAG's context manager |
| """ |
| with DAG("parent", default_args=default_args) as dag: |
| subdag = DAG("parent.test", default_args=default_args) |
| op = SubDagOperator(task_id="test", subdag=subdag) |
| |
| assert op.dag == dag |
| assert op.subdag == subdag |
| |
| def test_subdag_pools(self): |
| """ |
| Subdags and subdag tasks can't both have a pool with 1 slot |
| """ |
| dag = DAG("parent", default_args=default_args) |
| subdag = DAG("parent.child", default_args=default_args) |
| |
| session = airflow.settings.Session() |
| pool_1 = airflow.models.Pool(pool="test_pool_1", slots=1, include_deferred=False) |
| pool_10 = airflow.models.Pool(pool="test_pool_10", slots=10, include_deferred=False) |
| session.add(pool_1) |
| session.add(pool_10) |
| session.commit() |
| |
| EmptyOperator(task_id="dummy", dag=subdag, pool="test_pool_1") |
| |
| with pytest.raises(AirflowException): |
| SubDagOperator(task_id="child", dag=dag, subdag=subdag, pool="test_pool_1") |
| |
| # recreate dag because failed subdagoperator was already added |
| dag = DAG("parent", default_args=default_args) |
| SubDagOperator(task_id="child", dag=dag, subdag=subdag, pool="test_pool_10") |
| |
| session.delete(pool_1) |
| session.delete(pool_10) |
| session.commit() |
| |
| def test_subdag_pools_no_possible_conflict(self): |
| """ |
| Subdags and subdag tasks with no pool overlap, should not to query |
| pools |
| """ |
| dag = DAG("parent", default_args=default_args) |
| subdag = DAG("parent.child", default_args=default_args) |
| |
| session = airflow.settings.Session() |
| pool_1 = airflow.models.Pool(pool="test_pool_1", slots=1, include_deferred=False) |
| pool_10 = airflow.models.Pool(pool="test_pool_10", slots=10, include_deferred=False) |
| session.add(pool_1) |
| session.add(pool_10) |
| session.commit() |
| |
| EmptyOperator(task_id="dummy", dag=subdag, pool="test_pool_10") |
| |
| mock_session = Mock() |
| SubDagOperator(task_id="child", dag=dag, subdag=subdag, pool="test_pool_1", session=mock_session) |
| assert not mock_session.query.called |
| |
| session.delete(pool_1) |
| session.delete(pool_10) |
| session.commit() |
| |
| def test_execute_create_dagrun_wait_until_success(self): |
| """ |
| When SubDagOperator executes, it creates a DagRun if there is no existing one |
| and wait until the DagRun succeeds. |
| """ |
| dag = DAG("parent", default_args=default_args) |
| subdag = DAG("parent.test", default_args=default_args) |
| subdag_task = SubDagOperator(task_id="test", subdag=subdag, dag=dag, poke_interval=1) |
| |
| subdag.create_dagrun = Mock() |
| subdag.create_dagrun.return_value = self.dag_run_running |
| |
| subdag_task._get_dagrun = Mock() |
| subdag_task._get_dagrun.side_effect = [None, self.dag_run_success, self.dag_run_success] |
| |
| context = { |
| "data_interval_start": None, |
| "data_interval_end": None, |
| "execution_date": DEFAULT_DATE, |
| } |
| |
| subdag_task.pre_execute(context=context) |
| subdag_task.execute(context=context) |
| subdag_task.post_execute(context=context) |
| |
| subdag.create_dagrun.assert_called_once_with( |
| run_type=DagRunType.SCHEDULED, |
| execution_date=DEFAULT_DATE, |
| data_interval=None, |
| conf=None, |
| state=State.RUNNING, |
| external_trigger=True, |
| ) |
| |
| assert 3 == subdag_task._get_dagrun.call_count |
| |
| def test_execute_create_dagrun_with_conf(self): |
| """ |
| When SubDagOperator executes, it creates a DagRun if there is no existing one |
| and wait until the DagRun succeeds. |
| """ |
| conf = {"key": "value"} |
| dag = DAG("parent", default_args=default_args) |
| subdag = DAG("parent.test", default_args=default_args) |
| subdag_task = SubDagOperator(task_id="test", subdag=subdag, dag=dag, poke_interval=1, conf=conf) |
| |
| subdag.create_dagrun = Mock() |
| subdag.create_dagrun.return_value = self.dag_run_running |
| |
| subdag_task._get_dagrun = Mock() |
| subdag_task._get_dagrun.side_effect = [None, self.dag_run_success, self.dag_run_success] |
| |
| context = { |
| "data_interval_start": None, |
| "data_interval_end": None, |
| "execution_date": DEFAULT_DATE, |
| } |
| |
| subdag_task.pre_execute(context=context) |
| subdag_task.execute(context=context) |
| subdag_task.post_execute(context=context) |
| |
| subdag.create_dagrun.assert_called_once_with( |
| run_type=DagRunType.SCHEDULED, |
| execution_date=DEFAULT_DATE, |
| data_interval=None, |
| conf=conf, |
| state=State.RUNNING, |
| external_trigger=True, |
| ) |
| |
| assert 3 == subdag_task._get_dagrun.call_count |
| |
| def test_execute_dagrun_failed(self): |
| """ |
| When the DagRun failed during the execution, it raises an Airflow Exception. |
| """ |
| dag = DAG("parent", default_args=default_args) |
| subdag = DAG("parent.test", default_args=default_args) |
| subdag_task = SubDagOperator(task_id="test", subdag=subdag, dag=dag, poke_interval=1) |
| |
| subdag.create_dagrun = Mock() |
| subdag.create_dagrun.return_value = self.dag_run_running |
| |
| subdag_task._get_dagrun = Mock() |
| subdag_task._get_dagrun.side_effect = [None, self.dag_run_failed, self.dag_run_failed] |
| |
| context = { |
| "data_interval_start": None, |
| "data_interval_end": None, |
| "execution_date": DEFAULT_DATE, |
| } |
| |
| subdag_task.pre_execute(context=context) |
| subdag_task.execute(context=context) |
| with pytest.raises(AirflowException): |
| subdag_task.post_execute(context=context) |
| |
| def test_execute_skip_if_dagrun_success(self): |
| """ |
| When there is an existing DagRun in SUCCESS state, skip the execution. |
| """ |
| dag = DAG("parent", default_args=default_args) |
| subdag = DAG("parent.test", default_args=default_args) |
| |
| subdag.create_dagrun = Mock() |
| subdag_task = SubDagOperator(task_id="test", subdag=subdag, dag=dag, poke_interval=1) |
| subdag_task._get_dagrun = Mock() |
| subdag_task._get_dagrun.return_value = self.dag_run_success |
| |
| context = { |
| "data_interval_start": None, |
| "data_interval_end": None, |
| "execution_date": DEFAULT_DATE, |
| } |
| |
| subdag_task.pre_execute(context=context) |
| subdag_task.execute(context=context) |
| subdag_task.post_execute(context=context) |
| |
| subdag.create_dagrun.assert_not_called() |
| assert 3 == subdag_task._get_dagrun.call_count |
| |
| def test_rerun_failed_subdag(self, dag_maker): |
| """ |
| When there is an existing DagRun with failed state, reset the DagRun and the |
| corresponding TaskInstances |
| """ |
| with create_session() as session: |
| with dag_maker("parent.test", default_args=default_args, session=session) as subdag: |
| dummy_task = EmptyOperator(task_id="dummy") |
| sub_dagrun = dag_maker.create_dagrun( |
| run_type=DagRunType.SCHEDULED, |
| execution_date=DEFAULT_DATE, |
| state=State.FAILED, |
| external_trigger=True, |
| ) |
| |
| (dummy_task_instance,) = sub_dagrun.task_instances |
| dummy_task_instance.refresh_from_task(dummy_task) |
| dummy_task_instance.state == State.FAILED |
| |
| with dag_maker("parent", default_args=default_args, session=session): |
| subdag_task = SubDagOperator(task_id="test", subdag=subdag, poke_interval=1) |
| dag_maker.create_dagrun(execution_date=DEFAULT_DATE, run_type=DagRunType.SCHEDULED) |
| |
| subdag_task._reset_dag_run_and_task_instances(sub_dagrun, execution_date=DEFAULT_DATE) |
| |
| dummy_task_instance.refresh_from_db() |
| assert dummy_task_instance.state == State.NONE |
| |
| sub_dagrun.refresh_from_db() |
| assert sub_dagrun.state == State.RUNNING |
| |
| @pytest.mark.parametrize( |
| "propagate_option, states, skip_parent", |
| [ |
| (SkippedStatePropagationOptions.ALL_LEAVES, [State.SKIPPED, State.SKIPPED], True), |
| (SkippedStatePropagationOptions.ALL_LEAVES, [State.SKIPPED, State.SUCCESS], False), |
| (SkippedStatePropagationOptions.ANY_LEAF, [State.SKIPPED, State.SUCCESS], True), |
| (SkippedStatePropagationOptions.ANY_LEAF, [State.FAILED, State.SKIPPED], True), |
| (None, [State.SKIPPED, State.SKIPPED], False), |
| ], |
| ) |
| @mock.patch("airflow.operators.subdag.SubDagOperator.skip") |
| @mock.patch("airflow.operators.subdag.get_task_instance") |
| def test_subdag_with_propagate_skipped_state( |
| self, |
| mock_get_task_instance, |
| mock_skip, |
| dag_maker, |
| propagate_option, |
| states, |
| skip_parent, |
| ): |
| """ |
| Tests that skipped state of leaf tasks propagates to the parent dag. |
| Note that the skipped state propagation only takes affect when the dagrun's state is SUCCESS. |
| """ |
| with dag_maker("parent.test", default_args=default_args) as subdag: |
| dummy_subdag_tasks = [EmptyOperator(task_id=f"dummy_subdag_{i}") for i in range(len(states))] |
| dag_maker.create_dagrun(execution_date=DEFAULT_DATE) |
| |
| with dag_maker("parent", default_args=default_args): |
| subdag_task = SubDagOperator( |
| task_id="test", |
| subdag=subdag, |
| poke_interval=1, |
| propagate_skipped_state=propagate_option, |
| ) |
| dummy_dag_task = EmptyOperator(task_id="dummy_dag") |
| subdag_task >> dummy_dag_task |
| dag_run = dag_maker.create_dagrun(execution_date=DEFAULT_DATE) |
| |
| subdag_task._get_dagrun = Mock(return_value=self.dag_run_success) |
| |
| mock_get_task_instance.side_effect = [ |
| TaskInstance(task=task, run_id=dag_run.run_id, state=state) |
| for task, state in zip(dummy_subdag_tasks, states) |
| ] |
| |
| context = { |
| "execution_date": DEFAULT_DATE, |
| "dag_run": dag_run, |
| "task": subdag_task, |
| "ti": mock.MagicMock(map_index=-1), |
| } |
| subdag_task.post_execute(context) |
| |
| if skip_parent: |
| mock_skip.assert_called_once_with( |
| context["dag_run"], context["execution_date"], [dummy_dag_task], map_index=-1 |
| ) |
| else: |
| mock_skip.assert_not_called() |
| |
| def test_deprecation_warning(self): |
| dag = DAG("parent", default_args=default_args) |
| subdag = DAG("parent.test", default_args=default_args) |
| warning_message = """This class is deprecated. Please use `airflow.utils.task_group.TaskGroup`.""" |
| |
| with pytest.warns(DeprecationWarning) as warnings: |
| SubDagOperator(task_id="test", subdag=subdag, dag=dag) |
| assert warning_message == str(warnings[0].message) |
