Repository: incubator-airflow Updated Branches: refs/heads/master a279be65b -> 0ed608dca
[AIRFLOW-881] Check if SubDagOperator is in DAG context manager When initializing a SubDagOperator, the `dag` param should not be required if it is within a DAG context manager. So we check if that is the case and use that as the parent DAG if found (and `dag` param is not specified). Closes #2087 from dhuang/AIRFLOW-881 Project: http://git-wip-us.apache.org/repos/asf/incubator-airflow/repo Commit: http://git-wip-us.apache.org/repos/asf/incubator-airflow/commit/0ed608dc Tree: http://git-wip-us.apache.org/repos/asf/incubator-airflow/tree/0ed608dc Diff: http://git-wip-us.apache.org/repos/asf/incubator-airflow/diff/0ed608dc Branch: refs/heads/master Commit: 0ed608dca7b465a4107e7910dfaa838fbd702fdd Parents: a279be6 Author: Daniel Huang <dxhu...@gmail.com> Authored: Sun Feb 19 09:48:21 2017 +0100 Committer: Bolke de Bruin <bo...@xs4all.nl> Committed: Sun Feb 19 09:48:21 2017 +0100 ---------------------------------------------------------------------- airflow/operators/subdag_operator.py | 8 +++++--- tests/operators/subdag_operator.py | 11 +++++++++++ 2 files changed, 16 insertions(+), 3 deletions(-) ---------------------------------------------------------------------- http://git-wip-us.apache.org/repos/asf/incubator-airflow/blob/0ed608dc/airflow/operators/subdag_operator.py ---------------------------------------------------------------------- diff --git a/airflow/operators/subdag_operator.py b/airflow/operators/subdag_operator.py index a44e210..569d0a0 100644 --- a/airflow/operators/subdag_operator.py +++ b/airflow/operators/subdag_operator.py @@ -41,9 +41,11 @@ class SubDagOperator(BaseOperator): :param dag: the parent DAG :type subdag: airflow.DAG """ - if 'dag' not in kwargs: - raise AirflowException("Please pass in the `dag` param") - dag = kwargs['dag'] + import airflow.models + dag = kwargs.get('dag') or airflow.models._CONTEXT_MANAGER_DAG + if not dag: + raise AirflowException('Please pass in the `dag` param or call ' + 'within a DAG context manager') session = kwargs.pop('session') super(SubDagOperator, self).__init__(*args, **kwargs) http://git-wip-us.apache.org/repos/asf/incubator-airflow/blob/0ed608dc/tests/operators/subdag_operator.py ---------------------------------------------------------------------- diff --git a/tests/operators/subdag_operator.py b/tests/operators/subdag_operator.py index 6a25ac3..eeb41ec 100644 --- a/tests/operators/subdag_operator.py +++ b/tests/operators/subdag_operator.py @@ -54,6 +54,17 @@ class SubDagOperatorTests(unittest.TestCase): 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) + + self.assertEqual(op.dag, dag) + self.assertEqual(op.subdag, subdag) + def test_subdag_pools(self): """ Subdags and subdag tasks can't both have a pool with 1 slot