apache-airflow-providers-amazon 9.4.0rc1__py3-none-any.whl → 9.5.0rc2__py3-none-any.whl

airflow/providers/amazon/aws/hooks/mwaa.py

@@ -18,6 +18,7 @@
 
 from __future__ import annotations
 
+import requests
 from botocore.exceptions import ClientError
 
 from airflow.providers.amazon.aws.hooks.base_aws import AwsBaseHook
@@ -29,6 +30,12 @@ class MwaaHook(AwsBaseHook):
 
     Provide thin wrapper around :external+boto3:py:class:`boto3.client("mwaa") <MWAA.Client>`
 
+    If your IAM policy doesn't have `airflow:InvokeRestApi` permission, the hook will use a fallback method
+    that uses the AWS credential to generate a local web login token for the Airflow Web UI and then directly
+    make requests to the Airflow API. This fallback method can be set as the default (and only) method used by
+    setting `generate_local_token` to True.  Learn more here:
+    https://docs.aws.amazon.com/mwaa/latest/userguide/access-mwaa-apache-airflow-rest-api.html#granting-access-MWAA-Enhanced-REST-API
+
     Additional arguments (such as ``aws_conn_id``) may be specified and
     are passed down to the underlying AwsBaseHook.
 
@@ -47,6 +54,7 @@ class MwaaHook(AwsBaseHook):
         method: str,
         body: dict | None = None,
         query_params: dict | None = None,
+        generate_local_token: bool = False,
     ) -> dict:
         """
         Invoke the REST API on the Airflow webserver with the specified inputs.
@@ -56,30 +64,86 @@ class MwaaHook(AwsBaseHook):
 
         :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 method: HTTP method used for making Airflow REST API calls: 'GET'|'PUT'|'POST'|'PATCH'|'DELETE'
         :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
+        :param generate_local_token: If True, only the local web token method is used without trying boto's
+            `invoke_rest_api` first. If False, the local web token method is used as a fallback after trying
+            boto's `invoke_rest_api`
         """
-        body = body or {}
+        # 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} if body else {}
+        query_params = query_params 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 {},
+            "Body": body,
+            "QueryParameters": query_params,
         }
+
+        if generate_local_token:
+            return self._invoke_rest_api_using_local_session_token(**api_kwargs)
+
         try:
-            result = self.conn.invoke_rest_api(**api_kwargs)
+            response = self.conn.invoke_rest_api(**api_kwargs)
             # ResponseMetadata is removed because it contains data that is either very unlikely to be useful
             # in XComs and logs, or redundant given the data already included in the response
-            result.pop("ResponseMetadata", None)
-            return result
+            response.pop("ResponseMetadata", None)
+            return response
+
         except ClientError as e:
-            to_log = e.response
-            # ResponseMetadata and Error are removed because they contain data that is either very unlikely to
-            # be useful in XComs and logs, or redundant given the data already included in the response
-            to_log.pop("ResponseMetadata", None)
-            to_log.pop("Error", None)
-            self.log.error(to_log)
-            raise e
+            if (
+                e.response["Error"]["Code"] == "AccessDeniedException"
+                and "Airflow role" in e.response["Error"]["Message"]
+            ):
+                self.log.info(
+                    "Access Denied due to missing airflow:InvokeRestApi in IAM policy. Trying again by generating local token..."
+                )
+                return self._invoke_rest_api_using_local_session_token(**api_kwargs)
+            else:
+                to_log = e.response
+                # ResponseMetadata is removed because it contains data that is either very unlikely to be
+                # useful in XComs and logs, or redundant given the data already included in the response
+                to_log.pop("ResponseMetadata", None)
+                self.log.error(to_log)
+                raise
+
+    def _invoke_rest_api_using_local_session_token(
+        self,
+        **api_kwargs,
+    ) -> dict:
+        try:
+            session, hostname = self._get_session_conn(api_kwargs["Name"])
+
+            response = session.request(
+                method=api_kwargs["Method"],
+                url=f"https://{hostname}/api/v1{api_kwargs['Path']}",
+                params=api_kwargs["QueryParameters"],
+                json=api_kwargs["Body"],
+                timeout=10,
+            )
+            response.raise_for_status()
+
+        except requests.HTTPError as e:
+            self.log.error(e.response.json())
+            raise
+
+        return {
+            "RestApiStatusCode": response.status_code,
+            "RestApiResponse": response.json(),
+        }
+
+    # Based on: https://docs.aws.amazon.com/mwaa/latest/userguide/access-mwaa-apache-airflow-rest-api.html#create-web-server-session-token
+    def _get_session_conn(self, env_name: str) -> tuple:
+        create_token_response = self.conn.create_web_login_token(Name=env_name)
+        web_server_hostname = create_token_response["WebServerHostname"]
+        web_token = create_token_response["WebToken"]
+
+        login_url = f"https://{web_server_hostname}/aws_mwaa/login"
+        login_payload = {"token": web_token}
+        session = requests.Session()
+        login_response = session.post(login_url, data=login_payload, timeout=10)
+        login_response.raise_for_status()
+
+        return session, web_server_hostname

airflow/providers/amazon/aws/hooks/redshift_cluster.py

@@ -93,7 +93,7 @@ class RedshiftHook(AwsBaseHook):
             return "cluster_not_found"
 
     async def cluster_status_async(self, cluster_identifier: str) -> str:
-        async with self.async_conn as client:
+        async with await self.get_async_conn() as client:
             response = await client.describe_clusters(ClusterIdentifier=cluster_identifier)
             return response["Clusters"][0]["ClusterStatus"] if response else None
 

airflow/providers/amazon/aws/hooks/redshift_data.py

@@ -275,7 +275,7 @@ class RedshiftDataHook(AwsGenericHook["RedshiftDataAPIServiceClient"]):
 
         :param statement_id: the UUID of the statement
         """
-        async with self.async_conn as client:
+        async with await self.get_async_conn() as client:
             desc = await client.describe_statement(Id=statement_id)
             return desc["Status"] in RUNNING_STATES
 
@@ -288,6 +288,6 @@ class RedshiftDataHook(AwsGenericHook["RedshiftDataAPIServiceClient"]):
 
         :param statement_id: the UUID of the statement
         """
-        async with self.async_conn as client:
+        async with await self.get_async_conn() as client:
             resp = await client.describe_statement(Id=statement_id)
             return self.parse_statement_response(resp)

airflow/providers/amazon/aws/hooks/s3.py

@@ -1494,7 +1494,9 @@ class S3Hook(AwsBaseHook):
             get_hook_lineage_collector().add_output_asset(
                 context=self,
                 scheme="file",
-                asset_kwargs={"path": file_path if file_path.is_absolute() else file_path.absolute()},
+                asset_kwargs={
+                    "path": str(file_path) if file_path.is_absolute() else str(file_path.absolute())
+                },
             )
             file = open(file_path, "wb")
         else:

airflow/providers/amazon/aws/hooks/sagemaker.py

@@ -1318,7 +1318,7 @@ class SageMakerHook(AwsBaseHook):
 
         :param job_name: the name of the training job
         """
-        async with self.async_conn as client:
+        async with await self.get_async_conn() as client:
             response: dict[str, Any] = await client.describe_training_job(TrainingJobName=job_name)
             return response
 

airflow/providers/amazon/aws/hooks/sagemaker_unified_studio.py

@@ -0,0 +1,188 @@
+# 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 the Amazon SageMaker Unified Studio Notebook hook."""
+
+from __future__ import annotations
+
+import time
+
+from sagemaker_studio import ClientConfig
+from sagemaker_studio.sagemaker_studio_api import SageMakerStudioAPI
+
+from airflow.exceptions import AirflowException
+from airflow.hooks.base import BaseHook
+from airflow.providers.amazon.aws.utils.sagemaker_unified_studio import is_local_runner
+
+
+class SageMakerNotebookHook(BaseHook):
+    """
+    Interact with Sagemaker Unified Studio Workflows.
+
+    This hook provides a wrapper around the Sagemaker Workflows Notebook Execution API.
+
+    Examples:
+     .. code-block:: python
+
+        from airflow.providers.amazon.aws.hooks.sagemaker_unified_studio import SageMakerNotebookHook
+
+        notebook_hook = SageMakerNotebookHook(
+            input_config={"input_path": "path/to/notebook.ipynb", "input_params": {"param1": "value1"}},
+            output_config={"output_uri": "folder/output/location/prefix", "output_formats": "NOTEBOOK"},
+            execution_name="notebook_execution",
+            waiter_delay=10,
+            waiter_max_attempts=1440,
+        )
+
+    :param execution_name: The name of the notebook job to be executed, this is same as task_id.
+    :param input_config: Configuration for the input file.
+        Example: {'input_path': 'folder/input/notebook.ipynb', 'input_params': {'param1': 'value1'}}
+    :param output_config: Configuration for the output format. It should include an output_formats parameter to specify the output format.
+        Example: {'output_formats': ['NOTEBOOK']}
+    :param compute: compute configuration to use for the notebook execution. This is a required attribute
+        if the execution is on a remote compute.
+        Example: { "instance_type": "ml.m5.large", "volume_size_in_gb": 30, "volume_kms_key_id": "", "image_uri": "string", "container_entrypoint": [ "string" ]}
+    :param termination_condition: conditions to match to terminate the remote execution.
+        Example: { "MaxRuntimeInSeconds": 3600 }
+    :param tags: tags to be associated with the remote execution runs.
+        Example: { "md_analytics": "logs" }
+    :param waiter_delay: Interval in seconds to check the task execution status.
+    :param waiter_max_attempts: Number of attempts to wait before returning FAILED.
+    """
+
+    def __init__(
+        self,
+        execution_name: str,
+        input_config: dict | None = None,
+        output_config: dict | None = None,
+        compute: dict | None = None,
+        termination_condition: dict | None = None,
+        tags: dict | None = None,
+        waiter_delay: int = 10,
+        waiter_max_attempts: int = 1440,
+        *args,
+        **kwargs,
+    ):
+        super().__init__(*args, **kwargs)
+        self._sagemaker_studio = SageMakerStudioAPI(self._get_sagemaker_studio_config())
+        self.execution_name = execution_name
+        self.input_config = input_config or {}
+        self.output_config = output_config or {"output_formats": ["NOTEBOOK"]}
+        self.compute = compute
+        self.termination_condition = termination_condition or {}
+        self.tags = tags or {}
+        self.waiter_delay = waiter_delay
+        self.waiter_max_attempts = waiter_max_attempts
+
+    def _get_sagemaker_studio_config(self):
+        config = ClientConfig()
+        config.overrides["execution"] = {"local": is_local_runner()}
+        return config
+
+    def _format_start_execution_input_config(self):
+        config = {
+            "notebook_config": {
+                "input_path": self.input_config.get("input_path"),
+                "input_parameters": self.input_config.get("input_params"),
+            },
+        }
+
+        return config
+
+    def _format_start_execution_output_config(self):
+        output_formats = self.output_config.get("output_formats")
+        config = {
+            "notebook_config": {
+                "output_formats": output_formats,
+            }
+        }
+        return config
+
+    def start_notebook_execution(self):
+        start_execution_params = {
+            "execution_name": self.execution_name,
+            "execution_type": "NOTEBOOK",
+            "input_config": self._format_start_execution_input_config(),
+            "output_config": self._format_start_execution_output_config(),
+            "termination_condition": self.termination_condition,
+            "tags": self.tags,
+        }
+        if self.compute:
+            start_execution_params["compute"] = self.compute
+        else:
+            start_execution_params["compute"] = {"instance_type": "ml.m4.xlarge"}
+
+        print(start_execution_params)
+        return self._sagemaker_studio.execution_client.start_execution(**start_execution_params)
+
+    def wait_for_execution_completion(self, execution_id, context):
+        wait_attempts = 0
+        while wait_attempts < self.waiter_max_attempts:
+            wait_attempts += 1
+            time.sleep(self.waiter_delay)
+            response = self._sagemaker_studio.execution_client.get_execution(execution_id=execution_id)
+            error_message = response.get("error_details", {}).get("error_message")
+            status = response["status"]
+            if "files" in response:
+                self._set_xcom_files(response["files"], context)
+            if "s3_path" in response:
+                self._set_xcom_s3_path(response["s3_path"], context)
+
+            ret = self._handle_state(execution_id, status, error_message)
+            if ret:
+                return ret
+
+        # If timeout, handle state FAILED with timeout message
+        return self._handle_state(execution_id, "FAILED", "Execution timed out")
+
+    def _set_xcom_files(self, files, context):
+        if not context:
+            error_message = "context is required"
+            raise AirflowException(error_message)
+        for file in files:
+            context["ti"].xcom_push(
+                key=f"{file['display_name']}.{file['file_format']}",
+                value=file["file_path"],
+            )
+
+    def _set_xcom_s3_path(self, s3_path, context):
+        if not context:
+            error_message = "context is required"
+            raise AirflowException(error_message)
+        context["ti"].xcom_push(
+            key="s3_path",
+            value=s3_path,
+        )
+
+    def _handle_state(self, execution_id, status, error_message):
+        finished_states = ["COMPLETED"]
+        in_progress_states = ["IN_PROGRESS", "STOPPING"]
+
+        if status in in_progress_states:
+            info_message = f"Execution {execution_id} is still in progress with state:{status}, will check for a terminal status again in {self.waiter_delay}"
+            self.log.info(info_message)
+            return None
+        execution_message = f"Exiting Execution {execution_id} State: {status}"
+        if status in finished_states:
+            self.log.info(execution_message)
+            return {"Status": status, "ExecutionId": execution_id}
+        else:
+            log_error_message = f"Execution {execution_id} failed with error: {error_message}"
+            self.log.error(log_error_message)
+            if error_message == "":
+                error_message = execution_message
+            raise AirflowException(error_message)

airflow/providers/amazon/aws/links/base_aws.py

@@ -19,14 +19,21 @@ from __future__ import annotations
 
 from typing import TYPE_CHECKING, ClassVar
 
-from airflow.models import BaseOperatorLink, XCom
 from airflow.providers.amazon.aws.utils.suppress import return_on_error
+from airflow.providers.amazon.version_compat import AIRFLOW_V_3_0_PLUS
 
 if TYPE_CHECKING:
     from airflow.models import BaseOperator
     from airflow.models.taskinstancekey import TaskInstanceKey
     from airflow.utils.context import Context
 
+if AIRFLOW_V_3_0_PLUS:
+    from airflow.sdk import BaseOperatorLink
+    from airflow.sdk.execution_time.xcom import XCom
+else:
+    from airflow.models import XCom  # type: ignore[no-redef]
+    from airflow.models.baseoperatorlink import BaseOperatorLink  # type: ignore[no-redef]
+
 
 BASE_AWS_CONSOLE_LINK = "https://console.{aws_domain}"
 

airflow/providers/amazon/aws/links/sagemaker_unified_studio.py

@@ -0,0 +1,27 @@
+# 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 airflow.providers.amazon.aws.links.base_aws import BASE_AWS_CONSOLE_LINK, BaseAwsLink
+
+
+class SageMakerUnifiedStudioLink(BaseAwsLink):
+    """Helper class for constructing Amazon SageMaker Unified Studio Links."""
+
+    name = "Amazon SageMaker Unified Studio"
+    key = "sagemaker_unified_studio"
+    format_str = BASE_AWS_CONSOLE_LINK + "/datazone/home?region={region_name}"

airflow/providers/amazon/aws/log/s3_task_handler.py

@@ -26,6 +26,7 @@ from typing import TYPE_CHECKING
 
 from airflow.configuration import conf
 from airflow.providers.amazon.aws.hooks.s3 import S3Hook
+from airflow.providers.amazon.version_compat import AIRFLOW_V_3_0_PLUS
 from airflow.utils.log.file_task_handler import FileTaskHandler
 from airflow.utils.log.logging_mixin import LoggingMixin
 
@@ -58,7 +59,8 @@ class S3TaskHandler(FileTaskHandler, LoggingMixin):
     def hook(self):
         """Returns S3Hook."""
         return S3Hook(
-            aws_conn_id=conf.get("logging", "REMOTE_LOG_CONN_ID"), transfer_config_args={"use_threads": False}
+            aws_conn_id=conf.get("logging", "REMOTE_LOG_CONN_ID"),
+            transfer_config_args={"use_threads": False},
         )
 
     def set_context(self, ti: TaskInstance, *, identifier: str | None = None) -> None:
@@ -73,7 +75,7 @@ class S3TaskHandler(FileTaskHandler, LoggingMixin):
         is_trigger_log_context = getattr(ti, "is_trigger_log_context", False)
         self.upload_on_close = is_trigger_log_context or not getattr(ti, "raw", None)
         # Clear the file first so that duplicate data is not uploaded
-        # when re-using the same path (e.g. with rescheduled sensors)
+        # when reusing the same path (e.g. with rescheduled sensors)
         if self.upload_on_close:
             with open(self.handler.baseFilename, "w"):
                 pass
@@ -116,12 +118,16 @@ class S3TaskHandler(FileTaskHandler, LoggingMixin):
         keys = self.hook.list_keys(bucket_name=bucket, prefix=prefix)
         if keys:
             keys = sorted(f"s3://{bucket}/{key}" for key in keys)
-            messages.append("Found logs in s3:")
-            messages.extend(f"  * {key}" for key in keys)
+            if AIRFLOW_V_3_0_PLUS:
+                messages = keys
+            else:
+                messages.append("Found logs in s3:")
+                messages.extend(f"  * {key}" for key in keys)
             for key in keys:
                 logs.append(self.s3_read(key, return_error=True))
         else:
-            messages.append(f"No logs found on s3 for ti={ti}")
+            if not AIRFLOW_V_3_0_PLUS:
+                messages.append(f"No logs found on s3 for ti={ti}")
         return messages, logs
 
     def s3_log_exists(self, remote_log_location: str) -> bool:
@@ -152,7 +158,13 @@ class S3TaskHandler(FileTaskHandler, LoggingMixin):
                 return msg
         return ""
 
-    def s3_write(self, log: str, remote_log_location: str, append: bool = True, max_retry: int = 1) -> bool:
+    def s3_write(
+        self,
+        log: str,
+        remote_log_location: str,
+        append: bool = True,
+        max_retry: int = 1,
+    ) -> bool:
         """
         Write the log to the remote_log_location; return `True` or fails silently and return `False`.
 
@@ -185,7 +197,10 @@ class S3TaskHandler(FileTaskHandler, LoggingMixin):
                 break
             except Exception:
                 if try_num < max_retry:
-                    self.log.warning("Failed attempt to write logs to %s, will retry", remote_log_location)
+                    self.log.warning(
+                        "Failed attempt to write logs to %s, will retry",
+                        remote_log_location,
+                    )
                 else:
                     self.log.exception("Could not write logs to %s", remote_log_location)
                     return False

airflow/providers/amazon/aws/notifications/chime.py

@@ -21,12 +21,11 @@ from functools import cached_property
 from typing import TYPE_CHECKING
 
 from airflow.providers.amazon.aws.hooks.chime import ChimeWebhookHook
+from airflow.providers.common.compat.notifier import BaseNotifier
 
 if TYPE_CHECKING:
     from airflow.utils.context import Context
 
-from airflow.notifications.basenotifier import BaseNotifier
-
 
 class ChimeNotifier(BaseNotifier):
     """

airflow/providers/amazon/aws/notifications/sns.py

@@ -20,8 +20,8 @@ from __future__ import annotations
 from collections.abc import Sequence
 from functools import cached_property
 
-from airflow.notifications.basenotifier import BaseNotifier
 from airflow.providers.amazon.aws.hooks.sns import SnsHook
+from airflow.providers.common.compat.notifier import BaseNotifier
 
 
 class SnsNotifier(BaseNotifier):

airflow/providers/amazon/aws/notifications/sqs.py

@@ -20,8 +20,8 @@ from __future__ import annotations
 from collections.abc import Sequence
 from functools import cached_property
 
-from airflow.notifications.basenotifier import BaseNotifier
 from airflow.providers.amazon.aws.hooks.sqs import SqsHook
+from airflow.providers.common.compat.notifier import BaseNotifier
 
 
 class SqsNotifier(BaseNotifier):