From 2baa3485da5e4dc77e12f06d1678ee2c162080ad Mon Sep 17 00:00:00 2001 From: Chethan UK Date: Fri, 3 Jun 2022 01:12:32 +0100 Subject: [PATCH 1/3] Migrate Papermill example DAGs to new design #22456 --- .../papermill/example_dags/__init__.py | 17 ----------------- .../index.rst | 2 +- .../operators.rst | 2 +- .../providers/papermill}/example_papermill.py | 10 +++++++++- .../providers/papermill}/input_notebook.ipynb | 0 5 files changed, 11 insertions(+), 20 deletions(-) delete mode 100644 airflow/providers/papermill/example_dags/__init__.py rename {airflow/providers/papermill/example_dags => tests/system/providers/papermill}/example_papermill.py (90%) rename {airflow/providers/papermill/example_dags => tests/system/providers/papermill}/input_notebook.ipynb (100%) diff --git a/airflow/providers/papermill/example_dags/__init__.py b/airflow/providers/papermill/example_dags/__init__.py deleted file mode 100644 index 217e5db960782..0000000000000 --- a/airflow/providers/papermill/example_dags/__init__.py +++ /dev/null @@ -1,17 +0,0 @@ -# -# 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. diff --git a/docs/apache-airflow-providers-papermill/index.rst b/docs/apache-airflow-providers-papermill/index.rst index 7effd4b35ca76..54fe848e7014e 100644 --- a/docs/apache-airflow-providers-papermill/index.rst +++ b/docs/apache-airflow-providers-papermill/index.rst @@ -38,7 +38,7 @@ Content :maxdepth: 1 :caption: Resources - Example DAGs + Example DAGs PyPI Repository Installing from sources diff --git a/docs/apache-airflow-providers-papermill/operators.rst b/docs/apache-airflow-providers-papermill/operators.rst index 274d83adb919c..a92876ab85db3 100644 --- a/docs/apache-airflow-providers-papermill/operators.rst +++ b/docs/apache-airflow-providers-papermill/operators.rst @@ -50,7 +50,7 @@ Example DAG Use the :class:`~airflow.providers.papermill.operators.papermill.PapermillOperator` to execute a jupyter notebook: -.. exampleinclude:: /../../airflow/providers/papermill/example_dags/example_papermill.py +.. exampleinclude:: /../../tests/system/providers/papermill/example_papermill.py :language: python :dedent: 4 :start-after: [START howto_operator_papermill] diff --git a/airflow/providers/papermill/example_dags/example_papermill.py b/tests/system/providers/papermill/example_papermill.py similarity index 90% rename from airflow/providers/papermill/example_dags/example_papermill.py rename to tests/system/providers/papermill/example_papermill.py index c49b7715794ad..92ca8588fe9ad 100644 --- a/airflow/providers/papermill/example_dags/example_papermill.py +++ b/tests/system/providers/papermill/example_papermill.py @@ -33,9 +33,11 @@ START_DATE = datetime(2021, 1, 1) SCHEDULE_INTERVAL = '0 0 * * *' DAGRUN_TIMEOUT = timedelta(minutes=60) +ENV_ID = os.environ.get("SYSTEM_TESTS_ENV_ID") +DAG_ID = "example_papermill_operator" with DAG( - dag_id='example_papermill_operator', + dag_id=DAG_ID, schedule_interval=SCHEDULE_INTERVAL, start_date=START_DATE, dagrun_timeout=DAGRUN_TIMEOUT, @@ -83,3 +85,9 @@ def check_notebook(inlets, execution_date): ) run_this >> check_notebook(inlets=AUTO, execution_date="{{ execution_date }}") + + +from tests.system.utils import get_test_run # noqa: E402 + +# Needed to run the example DAG with pytest (see: tests/system/README.md#run_via_pytest) +test_run = get_test_run(dag) diff --git a/airflow/providers/papermill/example_dags/input_notebook.ipynb b/tests/system/providers/papermill/input_notebook.ipynb similarity index 100% rename from airflow/providers/papermill/example_dags/input_notebook.ipynb rename to tests/system/providers/papermill/input_notebook.ipynb From e214932b01ce2fd990b53365bb06c96959b8edf8 Mon Sep 17 00:00:00 2001 From: Chethan UK Date: Fri, 3 Jun 2022 12:32:05 +0100 Subject: [PATCH 2/3] Fix static checks --- tests/system/providers/papermill/example_papermill.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/tests/system/providers/papermill/example_papermill.py b/tests/system/providers/papermill/example_papermill.py index 92ca8588fe9ad..22580278d39fa 100644 --- a/tests/system/providers/papermill/example_papermill.py +++ b/tests/system/providers/papermill/example_papermill.py @@ -90,4 +90,5 @@ def check_notebook(inlets, execution_date): from tests.system.utils import get_test_run # noqa: E402 # Needed to run the example DAG with pytest (see: tests/system/README.md#run_via_pytest) -test_run = get_test_run(dag) +test_run = get_test_run(dag_1) +test_run_1 = get_test_run(dag_2) From de005d98f7e716353fbc98851dc90fb67bd1c274 Mon Sep 17 00:00:00 2001 From: Chethan UK Date: Fri, 3 Jun 2022 13:50:44 +0100 Subject: [PATCH 3/3] Fix static checks --- .../operators.rst | 7 ++ tests/system/__init__.py | 17 ++++ tests/system/providers/__init__.py | 17 ++++ tests/system/providers/papermill/__init__.py | 17 ++++ .../providers/papermill/example_papermill.py | 43 +---------- .../papermill/example_papermill_verify.py | 77 +++++++++++++++++++ tests/www/api/experimental/test_endpoints.py | 2 +- 7 files changed, 138 insertions(+), 42 deletions(-) create mode 100644 tests/system/__init__.py create mode 100644 tests/system/providers/__init__.py create mode 100644 tests/system/providers/papermill/__init__.py create mode 100644 tests/system/providers/papermill/example_papermill_verify.py diff --git a/docs/apache-airflow-providers-papermill/operators.rst b/docs/apache-airflow-providers-papermill/operators.rst index a92876ab85db3..c760eecaf0164 100644 --- a/docs/apache-airflow-providers-papermill/operators.rst +++ b/docs/apache-airflow-providers-papermill/operators.rst @@ -55,3 +55,10 @@ to execute a jupyter notebook: :dedent: 4 :start-after: [START howto_operator_papermill] :end-before: [END howto_operator_papermill] + +Example DAG to Verify the message in the notebook: + +.. exampleinclude:: /../../tests/system/providers/papermill/example_papermill_verify.py + :language: python + :start-after: [START howto_verify_operator_papermill] + :end-before: [END howto_verify_operator_papermill] diff --git a/tests/system/__init__.py b/tests/system/__init__.py new file mode 100644 index 0000000000000..217e5db960782 --- /dev/null +++ b/tests/system/__init__.py @@ -0,0 +1,17 @@ +# +# 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. diff --git a/tests/system/providers/__init__.py b/tests/system/providers/__init__.py new file mode 100644 index 0000000000000..217e5db960782 --- /dev/null +++ b/tests/system/providers/__init__.py @@ -0,0 +1,17 @@ +# +# 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. diff --git a/tests/system/providers/papermill/__init__.py b/tests/system/providers/papermill/__init__.py new file mode 100644 index 0000000000000..217e5db960782 --- /dev/null +++ b/tests/system/providers/papermill/__init__.py @@ -0,0 +1,17 @@ +# +# 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. diff --git a/tests/system/providers/papermill/example_papermill.py b/tests/system/providers/papermill/example_papermill.py index 22580278d39fa..8f828e4ed11f8 100644 --- a/tests/system/providers/papermill/example_papermill.py +++ b/tests/system/providers/papermill/example_papermill.py @@ -23,11 +23,7 @@ import os from datetime import datetime, timedelta -import scrapbook as sb - from airflow import DAG -from airflow.decorators import task -from airflow.lineage import AUTO from airflow.providers.papermill.operators.papermill import PapermillOperator START_DATE = datetime(2021, 1, 1) @@ -43,7 +39,7 @@ dagrun_timeout=DAGRUN_TIMEOUT, tags=['example'], catchup=False, -) as dag_1: +) as dag: # [START howto_operator_papermill] run_this = PapermillOperator( task_id="run_example_notebook", @@ -53,42 +49,7 @@ ) # [END howto_operator_papermill] - -@task -def check_notebook(inlets, execution_date): - """ - Verify the message in the notebook - """ - notebook = sb.read_notebook(inlets[0].url) - message = notebook.scraps['message'] - print(f"Message in notebook {message} for {execution_date}") - - if message.data != f"Ran from Airflow at {execution_date}!": - return False - - return True - - -with DAG( - dag_id='example_papermill_operator_2', - schedule_interval=SCHEDULE_INTERVAL, - start_date=START_DATE, - dagrun_timeout=DAGRUN_TIMEOUT, - catchup=False, -) as dag_2: - - run_this = PapermillOperator( - task_id="run_example_notebook", - input_nb=os.path.join(os.path.dirname(os.path.realpath(__file__)), "input_notebook.ipynb"), - output_nb="/tmp/out-{{ execution_date }}.ipynb", - parameters={"msgs": "Ran from Airflow at {{ execution_date }}!"}, - ) - - run_this >> check_notebook(inlets=AUTO, execution_date="{{ execution_date }}") - - from tests.system.utils import get_test_run # noqa: E402 # Needed to run the example DAG with pytest (see: tests/system/README.md#run_via_pytest) -test_run = get_test_run(dag_1) -test_run_1 = get_test_run(dag_2) +test_run = get_test_run(dag) diff --git a/tests/system/providers/papermill/example_papermill_verify.py b/tests/system/providers/papermill/example_papermill_verify.py new file mode 100644 index 0000000000000..dd24fc51f6363 --- /dev/null +++ b/tests/system/providers/papermill/example_papermill_verify.py @@ -0,0 +1,77 @@ +# +# 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. +""" +This DAG will use Papermill to run the notebook "hello_world", based on the execution date +it will create an output notebook "out-". All fields, including the keys in the parameters, are +templated. +""" +import os +from datetime import datetime, timedelta + +import scrapbook as sb + +from airflow import DAG +from airflow.decorators import task +from airflow.lineage import AUTO +from airflow.providers.papermill.operators.papermill import PapermillOperator + +START_DATE = datetime(2021, 1, 1) +SCHEDULE_INTERVAL = '0 0 * * *' +DAGRUN_TIMEOUT = timedelta(minutes=60) +ENV_ID = os.environ.get("SYSTEM_TESTS_ENV_ID") +DAG_ID = "example_papermill_operator_verify" + + +# [START howto_verify_operator_papermill] +@task +def check_notebook(inlets, execution_date): + """ + Verify the message in the notebook + """ + notebook = sb.read_notebook(inlets[0].url) + message = notebook.scraps['message'] + print(f"Message in notebook {message} for {execution_date}") + + if message.data != f"Ran from Airflow at {execution_date}!": + return False + + return True + + +with DAG( + dag_id='example_papermill_operator_verify', + schedule_interval=SCHEDULE_INTERVAL, + start_date=START_DATE, + dagrun_timeout=DAGRUN_TIMEOUT, + catchup=False, +) as dag: + + run_this = PapermillOperator( + task_id="run_example_notebook", + input_nb=os.path.join(os.path.dirname(os.path.realpath(__file__)), "input_notebook.ipynb"), + output_nb="/tmp/out-{{ execution_date }}.ipynb", + parameters={"msgs": "Ran from Airflow at {{ execution_date }}!"}, + ) + + run_this >> check_notebook(inlets=AUTO, execution_date="{{ execution_date }}") +# [END howto_verify_operator_papermill] + +from tests.system.utils import get_test_run # noqa: E402 + +# Needed to run the example DAG with pytest (see: tests/system/README.md#run_via_pytest) +test_run = get_test_run(dag) diff --git a/tests/www/api/experimental/test_endpoints.py b/tests/www/api/experimental/test_endpoints.py index 9710b38c68954..8d5516a409f66 100644 --- a/tests/www/api/experimental/test_endpoints.py +++ b/tests/www/api/experimental/test_endpoints.py @@ -312,7 +312,7 @@ def test_dagrun_status(self): class TestLineageApiExperimental(TestBase): - PAPERMILL_EXAMPLE_DAGS = os.path.join(ROOT_FOLDER, "airflow", "providers", "papermill", "example_dags") + PAPERMILL_EXAMPLE_DAGS = os.path.join(ROOT_FOLDER, "tests", "system", "providers", "papermill") @pytest.fixture(scope="class", autouse=True) def _populate_db(self):