Skip to content

Commit

Permalink
Migrate Papermill example DAGs to new design #22456 (#24146)
Browse files Browse the repository at this point in the history
  • Loading branch information
chethanuk authored Jun 5, 2022
1 parent dbe80c8 commit b4d50d3
Show file tree
Hide file tree
Showing 7 changed files with 76 additions and 22 deletions.
2 changes: 1 addition & 1 deletion docs/apache-airflow-providers-papermill/index.rst
Original file line number Diff line number Diff line change
Expand Up @@ -38,7 +38,7 @@ Content
:maxdepth: 1
:caption: Resources

Example DAGs <https://github.com/apache/airflow/tree/main/airflow/providers/papermill/example_dags>
Example DAGs <https://github.com/apache/airflow/tree/main/tests/system/providers/papermill>
PyPI Repository <https://pypi.org/project/apache-airflow-providers-papermill/>
Installing from sources <installing-providers-from-sources>

Expand Down
9 changes: 8 additions & 1 deletion docs/apache-airflow-providers-papermill/operators.rst
Original file line number Diff line number Diff line change
Expand Up @@ -50,8 +50,15 @@ 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]
: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]
55 changes: 55 additions & 0 deletions tests/system/providers/papermill/example_papermill.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,55 @@
#
# 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-<date>". All fields, including the keys in the parameters, are
templated.
"""
import os
from datetime import datetime, timedelta

from airflow import DAG
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"

with DAG(
dag_id=DAG_ID,
schedule_interval=SCHEDULE_INTERVAL,
start_date=START_DATE,
dagrun_timeout=DAGRUN_TIMEOUT,
tags=['example'],
catchup=False,
) as dag:
# [START howto_operator_papermill]
run_this = PapermillOperator(
task_id="run_example_notebook",
input_nb="/tmp/hello_world.ipynb",
output_nb="/tmp/out-{{ execution_date }}.ipynb",
parameters={"msgs": "Ran from Airflow at {{ execution_date }}!"},
)
# [END howto_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)
Original file line number Diff line number Diff line change
Expand Up @@ -33,25 +33,11 @@
START_DATE = datetime(2021, 1, 1)
SCHEDULE_INTERVAL = '0 0 * * *'
DAGRUN_TIMEOUT = timedelta(minutes=60)

with DAG(
dag_id='example_papermill_operator',
schedule_interval=SCHEDULE_INTERVAL,
start_date=START_DATE,
dagrun_timeout=DAGRUN_TIMEOUT,
tags=['example'],
catchup=False,
) as dag_1:
# [START howto_operator_papermill]
run_this = PapermillOperator(
task_id="run_example_notebook",
input_nb="/tmp/hello_world.ipynb",
output_nb="/tmp/out-{{ execution_date }}.ipynb",
parameters={"msgs": "Ran from Airflow at {{ execution_date }}!"},
)
# [END howto_operator_papermill]
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):
"""
Expand All @@ -68,12 +54,12 @@ def check_notebook(inlets, execution_date):


with DAG(
dag_id='example_papermill_operator_2',
dag_id='example_papermill_operator_verify',
schedule_interval=SCHEDULE_INTERVAL,
start_date=START_DATE,
dagrun_timeout=DAGRUN_TIMEOUT,
catchup=False,
) as dag_2:
) as dag:

run_this = PapermillOperator(
task_id="run_example_notebook",
Expand All @@ -83,3 +69,9 @@ def check_notebook(inlets, 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)
2 changes: 1 addition & 1 deletion tests/www/api/experimental/test_endpoints.py
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand Down

0 comments on commit b4d50d3

Please sign in to comment.