o-nikolas commented on code in PR #46579: URL: https://github.com/apache/airflow/pull/46579#discussion_r1949653704
########## providers/src/airflow/providers/amazon/aws/hooks/mwaa.py: ########## @@ -0,0 +1,81 @@ +# 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 module contains AWS MWAA hook.""" + +from __future__ import annotations + +from botocore.exceptions import ClientError + +from airflow.providers.amazon.aws.hooks.base_aws import AwsBaseHook + + +class MwaaHook(AwsBaseHook): + """ + Interact with AWS Manager Workflows for Apache Airflow. + + Provide thin wrapper around :external+boto3:py:class:`boto3.client("mwaa") <MWAA.Client>` + + Additional arguments (such as ``aws_conn_id``) may be specified and + are passed down to the underlying AwsBaseHook. + + .. seealso:: + - :class:`airflow.providers.amazon.aws.hooks.base_aws.AwsBaseHook` + """ + + def __init__(self, *args, **kwargs) -> None: + kwargs["client_type"] = "mwaa" + super().__init__(*args, **kwargs) + + def invoke_rest_api( + self, + env_name: str, + path: str, + method: str, + body: dict | None = None, + query_params: dict | None = None, + ) -> dict: + """ + Invoke the REST API on the Airflow webserver with the specified inputs. + + .. seealso:: + - :external+boto3:py:meth:`MWAA.Client.invoke_rest_api` + + :param env_name: name of the MWAA environment + :param path: Apache Airflow REST API endpoint path to be called + :param method: HTTP method used for making Airflow REST API calls + :param body: Request body for the Apache Airflow REST API call + :param query_params: Query parameters to be included in the Apache Airflow REST API call + """ + body = body or {} + api_kwargs = { + "Name": env_name, + "Path": path, + "Method": method, + # Filter out keys with None values because Airflow REST API doesn't accept requests otherwise + "Body": {k: v for k, v in body.items() if v is not None}, + "QueryParameters": query_params if query_params else {}, + } + try: + result = self.get_conn().invoke_rest_api(**api_kwargs) + result.pop("ResponseMetadata", None) Review Comment: Can you add a comment describing why ResponseMetadata is being removed here, for posterity. Same in the exception handling block. ########## providers/src/airflow/providers/amazon/aws/operators/mwaa.py: ########## @@ -0,0 +1,111 @@ +# 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 module contains AWS MWAA operators.""" + +from __future__ import annotations + +from collections.abc import Sequence +from typing import TYPE_CHECKING + +from airflow.providers.amazon.aws.hooks.mwaa import MwaaHook +from airflow.providers.amazon.aws.operators.base_aws import AwsBaseOperator +from airflow.providers.amazon.aws.utils.mixins import aws_template_fields + +if TYPE_CHECKING: + from airflow.utils.context import Context + + +class MwaaTriggerDagRunOperator(AwsBaseOperator[MwaaHook]): + """ + Trigger a Dag Run for a Dag in an Amazon MWAA environment. + + .. seealso:: + For more information on how to use this operator, take a look at the guide: + :ref:`howto/operator:MwaaTriggerDagRunOperator` + + :param env_name: The MWAA environment name (templated) + :param trigger_dag_id: The ID of the DAG to be triggered (templated) + :param trigger_run_id: The Run ID. The value of this field can be set only when creating the object. This + together with trigger_dag_id are a unique key. (templated) + :param logical_date: The logical date (previously called execution date). This is the time or interval + covered by this DAG run, according to the DAG definition. The value of this field can be set only when Review Comment: > The value of this field can be set only when creating the object You keep repeating this, I have not seen it before in any operator. Is it lifted from somewhere? It seems out of place here, I'm not sure what you're trying to indicate. ########## providers/tests/amazon/aws/hooks/test_mwaa.py: ########## @@ -0,0 +1,130 @@ +# 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 + +import pytest +from botocore.exceptions import ClientError +from moto import mock_aws + +from airflow.providers.amazon.aws.hooks.mwaa import MwaaHook + +ENV_NAME = "test_env" +PATH = "/dags/test_dag/dagRuns" +METHOD = "POST" +QUERY_PARAMS = {"limit": 30} + + +class TestMwaaHook: + def setup_method(self): + self.hook = MwaaHook() + + def test_init(self): + assert self.hook.client_type == "mwaa" + + @mock_aws + def test_get_conn(self): + assert self.hook.get_conn() is not None + + @pytest.mark.parametrize( + "body", + [ + None, # test case: empty body + {"conf": {}}, # test case: non-empty body + ], + ) + @mock.patch.object(MwaaHook, "get_conn") + def test_invoke_rest_api_success(self, mock_conn, body) -> None: + boto_invoke_mock = mock.MagicMock(return_value=self.example_responses["success"]) + mock_conn.return_value.invoke_rest_api = boto_invoke_mock + + retval = self.hook.invoke_rest_api(ENV_NAME, PATH, METHOD, body, QUERY_PARAMS) + kwargs_to_assert = { + "Name": ENV_NAME, + "Path": PATH, + "Method": METHOD, + "Body": body if body else {}, + "QueryParameters": QUERY_PARAMS, + } + boto_invoke_mock.assert_called_once_with(**kwargs_to_assert) + assert retval == { + k: v for k, v in self.example_responses["success"].items() if k != "ResponseMetadata" + } + + @mock.patch.object(MwaaHook, "get_conn") + def test_invoke_rest_api_failure(self, mock_conn) -> None: + error = ClientError( + error_response=self.example_responses["failure"], operation_name="invoke_rest_api" + ) + boto_invoke_mock = mock.MagicMock(side_effect=error) + mock_conn.return_value.invoke_rest_api = boto_invoke_mock + mock_log = mock.MagicMock() + self.hook.log.error = mock_log + + with pytest.raises(ClientError) as caught_error: + self.hook.invoke_rest_api(ENV_NAME, PATH, METHOD) + + assert caught_error.value == error + expected_log = { + k: v + for k, v in self.example_responses["failure"].items() + if k != "ResponseMetadata" and k != "Error" + } + mock_log.assert_called_once_with(expected_log) + + @pytest.fixture(autouse=True) + def _setup_test_cases(self): Review Comment: Also it's worth throwing a comment somewhere around this (wherever it ends up) mentioning that it is mutated by the code under test, which is why it's recreated before each test. ########## providers/tests/system/amazon/aws/example_mwaa.py: ########## @@ -0,0 +1,87 @@ +# 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 datetime import datetime + +from airflow.models.baseoperator import chain +from airflow.models.dag import DAG +from airflow.providers.amazon.aws.operators.mwaa import MwaaTriggerDagRunOperator + +from providers.tests.system.amazon.aws.utils import SystemTestContextBuilder + +DAG_ID = "example_mwaa" + +# Externally fetched variables: +EXISTING_ENVIRONMENT_NAME_KEY = "ENVIRONMENT_NAME" +EXISTING_DAG_ID_KEY = "DAG_ID" + + +sys_test_context_task = ( + SystemTestContextBuilder() + # NOTE: Creating a functional MWAA environment is time-consuming and requires + # manually creating and configuring an S3 bucket for DAG storage and a VPC with + # private subnets which is out of scope for this demo. To simplify this demo and + # make it run in a reasonable time, follow these steps in the AWS Console to create + # a new MWAA environment with default configuration: + # 1. Create an S3 bucket and upload your DAGs to a 'dags' directory + # 2. Navigate to the MWAA console + # 3. Create an environment, making sure to use the S3 bucket from the previous step + # 4. Use the default VPC/network settings (or create a new one using this guide: + # https://docs.aws.amazon.com/mwaa/latest/userguide/vpc-create.html) + # 5. Select Public network for web server access and click through to creation Review Comment: Love this comment! One nit: A lot of these steps are either included or completely automated by the console wizard, and they may change over time. There is also a quick start doc guide that deploys everything in one CFN template. I would just link to those resources rather than type out the steps here: https://docs.aws.amazon.com/mwaa/latest/userguide/quick-start.html ########## providers/tests/amazon/aws/hooks/test_mwaa.py: ########## @@ -0,0 +1,130 @@ +# 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 + +import pytest +from botocore.exceptions import ClientError +from moto import mock_aws + +from airflow.providers.amazon.aws.hooks.mwaa import MwaaHook + +ENV_NAME = "test_env" +PATH = "/dags/test_dag/dagRuns" +METHOD = "POST" +QUERY_PARAMS = {"limit": 30} + + +class TestMwaaHook: + def setup_method(self): + self.hook = MwaaHook() + + def test_init(self): + assert self.hook.client_type == "mwaa" + + @mock_aws + def test_get_conn(self): + assert self.hook.get_conn() is not None Review Comment: Nice, glad `mock_aws` worked here! ########## providers/tests/amazon/aws/hooks/test_mwaa.py: ########## @@ -0,0 +1,130 @@ +# 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 + +import pytest +from botocore.exceptions import ClientError +from moto import mock_aws + +from airflow.providers.amazon.aws.hooks.mwaa import MwaaHook + +ENV_NAME = "test_env" +PATH = "/dags/test_dag/dagRuns" +METHOD = "POST" +QUERY_PARAMS = {"limit": 30} + + +class TestMwaaHook: + def setup_method(self): + self.hook = MwaaHook() + + def test_init(self): + assert self.hook.client_type == "mwaa" + + @mock_aws + def test_get_conn(self): + assert self.hook.get_conn() is not None + + @pytest.mark.parametrize( + "body", + [ + None, # test case: empty body + {"conf": {}}, # test case: non-empty body + ], + ) + @mock.patch.object(MwaaHook, "get_conn") + def test_invoke_rest_api_success(self, mock_conn, body) -> None: + boto_invoke_mock = mock.MagicMock(return_value=self.example_responses["success"]) + mock_conn.return_value.invoke_rest_api = boto_invoke_mock + + retval = self.hook.invoke_rest_api(ENV_NAME, PATH, METHOD, body, QUERY_PARAMS) + kwargs_to_assert = { + "Name": ENV_NAME, + "Path": PATH, + "Method": METHOD, + "Body": body if body else {}, + "QueryParameters": QUERY_PARAMS, + } + boto_invoke_mock.assert_called_once_with(**kwargs_to_assert) + assert retval == { + k: v for k, v in self.example_responses["success"].items() if k != "ResponseMetadata" + } + + @mock.patch.object(MwaaHook, "get_conn") + def test_invoke_rest_api_failure(self, mock_conn) -> None: + error = ClientError( + error_response=self.example_responses["failure"], operation_name="invoke_rest_api" + ) + boto_invoke_mock = mock.MagicMock(side_effect=error) + mock_conn.return_value.invoke_rest_api = boto_invoke_mock + mock_log = mock.MagicMock() + self.hook.log.error = mock_log + + with pytest.raises(ClientError) as caught_error: + self.hook.invoke_rest_api(ENV_NAME, PATH, METHOD) + + assert caught_error.value == error + expected_log = { + k: v + for k, v in self.example_responses["failure"].items() + if k != "ResponseMetadata" and k != "Error" + } + mock_log.assert_called_once_with(expected_log) + + @pytest.fixture(autouse=True) + def _setup_test_cases(self): Review Comment: Why not do this in the `setup_method` you already have instead of creating an autouse fixture? -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: commits-unsubscr...@airflow.apache.org For queries about this service, please contact Infrastructure at: us...@infra.apache.org
