diff --git a/airflow/providers/papermill/operators/papermill.py b/airflow/providers/papermill/operators/papermill.py index 531304d441c7f..9e5d5c2e1b00c 100644 --- a/airflow/providers/papermill/operators/papermill.py +++ b/airflow/providers/papermill/operators/papermill.py @@ -17,7 +17,7 @@ # under the License. from __future__ import annotations -from typing import TYPE_CHECKING, Sequence +from typing import TYPE_CHECKING, ClassVar, Collection, Optional, Sequence import attr import papermill as pm @@ -33,8 +33,17 @@ class NoteBook(File): """Jupyter notebook""" - type_hint: str | None = "jupyter_notebook" - parameters: dict | None = {} + # For compatibility with Airflow 2.3: + # 1. Use predefined set because `File.template_fields` introduced in Airflow 2.4 + # 2. Use old styled annotations because `cattrs` doesn't work well with PEP 604. + + template_fields: ClassVar[Collection[str]] = { + "parameters", + *(File.template_fields if hasattr(File, "template_fields") else {"url"}), + } + + type_hint: Optional[str] = "jupyter_notebook" # noqa: UP007 + parameters: Optional[dict] = {} # noqa: UP007 meta_schema: str = __name__ + ".NoteBook" @@ -43,8 +52,8 @@ class PapermillOperator(BaseOperator): """ Executes a jupyter notebook through papermill that is annotated with parameters - :param input_nb: input notebook (can also be a NoteBook or a File inlet) - :param output_nb: output notebook (can also be a NoteBook or File outlet) + :param input_nb: input notebook, either path or NoteBook inlet. + :param output_nb: output notebook, either path or NoteBook outlet. :param parameters: the notebook parameters to set :param kernel_name: (optional) name of kernel to execute the notebook against (ignores kernel name in the notebook document metadata) @@ -57,36 +66,43 @@ class PapermillOperator(BaseOperator): def __init__( self, *, - input_nb: str | None = None, - output_nb: str | None = None, + input_nb: str | NoteBook | None = None, + output_nb: str | NoteBook | None = None, parameters: dict | None = None, kernel_name: str | None = None, language_name: str | None = None, **kwargs, ) -> None: super().__init__(**kwargs) - - self.input_nb = input_nb - self.output_nb = output_nb self.parameters = parameters + + if not input_nb: + raise ValueError("Input notebook is not specified") + elif not isinstance(input_nb, NoteBook): + self.input_nb = NoteBook(url=input_nb, parameters=self.parameters) + else: + self.input_nb = input_nb + + if not output_nb: + raise ValueError("Output notebook is not specified") + elif not isinstance(output_nb, NoteBook): + self.output_nb = NoteBook(url=output_nb) + else: + self.output_nb = output_nb + self.kernel_name = kernel_name self.language_name = language_name - if input_nb: - self.inlets.append(NoteBook(url=input_nb, parameters=self.parameters)) - if output_nb: - self.outlets.append(NoteBook(url=output_nb)) + + self.inlets.append(self.input_nb) + self.outlets.append(self.output_nb) def execute(self, context: Context): - if not self.inlets or not self.outlets: - raise ValueError("Input notebook or output notebook is not specified") - - for i, item in enumerate(self.inlets): - pm.execute_notebook( - item.url, - self.outlets[i].url, - parameters=item.parameters, - progress_bar=False, - report_mode=True, - kernel_name=self.kernel_name, - language=self.language_name, - ) + pm.execute_notebook( + self.input_nb.url, + self.output_nb.url, + parameters=self.input_nb.parameters, + progress_bar=False, + report_mode=True, + kernel_name=self.kernel_name, + language=self.language_name, + ) diff --git a/tests/providers/papermill/operators/test_papermill.py b/tests/providers/papermill/operators/test_papermill.py index 2ab23280a73be..55569e68bab45 100644 --- a/tests/providers/papermill/operators/test_papermill.py +++ b/tests/providers/papermill/operators/test_papermill.py @@ -19,14 +19,55 @@ from unittest.mock import patch -from airflow.models import DAG, DagRun, TaskInstance -from airflow.providers.papermill.operators.papermill import PapermillOperator +import pytest + +from airflow.providers.papermill.operators.papermill import NoteBook, PapermillOperator from airflow.utils import timezone DEFAULT_DATE = timezone.datetime(2021, 1, 1) +TEST_INPUT_URL = "/foo/bar" +TEST_OUTPUT_URL = "/spam/egg" + + +class TestNoteBook: + """Test NoteBook object.""" + + def test_templated_fields(self): + assert hasattr(NoteBook, "template_fields") + assert "parameters" in NoteBook.template_fields class TestPapermillOperator: + """Test PapermillOperator.""" + + def test_mandatory_attributes(self): + """Test missing Input or Output notebooks.""" + with pytest.raises(ValueError, match="Input notebook is not specified"): + PapermillOperator(task_id="missing_input_nb", output_nb="foo-bar") + + with pytest.raises(ValueError, match="Output notebook is not specified"): + PapermillOperator(task_id="missing_input_nb", input_nb="foo-bar") + + @pytest.mark.parametrize( + "output_nb", + [ + pytest.param(TEST_OUTPUT_URL, id="output-as-string"), + pytest.param(NoteBook(TEST_OUTPUT_URL), id="output-as-notebook-object"), + ], + ) + @pytest.mark.parametrize( + "input_nb", + [ + pytest.param(TEST_INPUT_URL, id="input-as-string"), + pytest.param(NoteBook(TEST_INPUT_URL), id="input-as-notebook-object"), + ], + ) + def test_notebooks_objects(self, input_nb, output_nb): + """Test different type of Input/Output notebooks arguments.""" + op = PapermillOperator(task_id="test_notebooks_objects", input_nb=input_nb, output_nb=output_nb) + assert op.input_nb.url == TEST_INPUT_URL + assert op.output_nb.url == TEST_OUTPUT_URL + @patch("airflow.providers.papermill.operators.papermill.pm") def test_execute(self, mock_papermill): in_nb = "/tmp/does_not_exist" @@ -57,26 +98,41 @@ def test_execute(self, mock_papermill): report_mode=True, ) - def test_render_template(self): - args = {"owner": "airflow", "start_date": DEFAULT_DATE} - dag = DAG("test_render_template", default_args=args) - - operator = PapermillOperator( - task_id="render_dag_test", + def test_render_template(self, create_task_instance_of_operator): + """Test rendering fields.""" + ti = create_task_instance_of_operator( + PapermillOperator, input_nb="/tmp/{{ dag.dag_id }}.ipynb", output_nb="/tmp/out-{{ dag.dag_id }}.ipynb", - parameters={"msgs": "dag id is {{ dag.dag_id }}!"}, - kernel_name="python3", - language_name="python", - dag=dag, + parameters={"msgs": "dag id is {{ dag.dag_id }}!", "test_dt": "{{ ds }}"}, + kernel_name="{{ params.kernel_name }}", + language_name="{{ params.language_name }}", + # Additional parameters for render fields + params={ + "kernel_name": "python3", + "language_name": "python", + }, + # TI Settings + dag_id="test_render_template", + task_id="render_dag_test", + execution_date=DEFAULT_DATE, ) + task = ti.render_templates() + + # Test render Input/Output notebook attributes + assert task.input_nb.url == "/tmp/test_render_template.ipynb" + assert task.input_nb.parameters == { + "msgs": "dag id is test_render_template!", + "test_dt": DEFAULT_DATE.date().isoformat(), + } + assert task.output_nb.url == "/tmp/out-test_render_template.ipynb" + assert task.output_nb.parameters == {} - ti = TaskInstance(operator, run_id="papermill_test") - ti.dag_run = DagRun(execution_date=DEFAULT_DATE) - ti.render_templates() + # Test render other templated attributes + assert task.parameters == task.input_nb.parameters + assert "python3" == task.kernel_name + assert "python" == task.language_name - assert "/tmp/test_render_template.ipynb" == operator.input_nb - assert "/tmp/out-test_render_template.ipynb" == operator.output_nb - assert {"msgs": "dag id is test_render_template!"} == operator.parameters - assert "python3" == operator.kernel_name - assert "python" == operator.language_name + # Test render Lineage inlets/outlets + assert task.inlets[0] == task.input_nb + assert task.outlets[0] == task.output_nb