ramitkataria commented on code in PR #46579: URL: https://github.com/apache/airflow/pull/46579#discussion_r1951800498
########## providers/src/airflow/providers/amazon/aws/hooks/mwaa.py: ########## @@ -0,0 +1,85 @@ +# 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) Review Comment: Just looked at `hooks/base_aws.py` and yes it does look like get_conn() exists just for compatibility now. So fixed, thanks! -- 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
