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

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

@@ -211,7 +211,7 @@ class GlueJobHook(AwsBaseHook):
 
         The async version of get_job_state.
         """
-        async with self.async_conn as client:
+        async with await self.get_async_conn() as client:
             job_run = await client.get_job_run(JobName=job_name, RunId=run_id)
         return job_run["JobRun"]["JobRunState"]
 
@@ -236,6 +236,9 @@ class GlueJobHook(AwsBaseHook):
         """
         log_client = self.logs_hook.get_conn()
         paginator = log_client.get_paginator("filter_log_events")
+        job_run = self.conn.get_job_run(JobName=job_name, RunId=run_id)["JobRun"]
+        # StartTime needs to be an int and is Epoch time in milliseconds
+        start_time = int(job_run["StartedOn"].timestamp() * 1000)
 
         def display_logs_from(log_group: str, continuation_token: str | None) -> str | None:
             """Mutualize iteration over the 2 different log streams glue jobs write to."""
@@ -245,6 +248,7 @@ class GlueJobHook(AwsBaseHook):
                 for response in paginator.paginate(
                     logGroupName=log_group,
                     logStreamNames=[run_id],
+                    startTime=start_time,
                     PaginationConfig={"StartingToken": continuation_token},
                 ):
                     fetched_logs.extend([event["message"] for event in response["events"]])
@@ -270,7 +274,7 @@ class GlueJobHook(AwsBaseHook):
                 self.log.info("No new log from the Glue Job in %s", log_group)
             return next_token
 
-        log_group_prefix = self.conn.get_job_run(JobName=job_name, RunId=run_id)["JobRun"]["LogGroupName"]
+        log_group_prefix = job_run["LogGroupName"]
         log_group_default = f"{log_group_prefix}/{DEFAULT_LOG_SUFFIX}"
         log_group_error = f"{log_group_prefix}/{ERROR_LOG_SUFFIX}"
         # one would think that the error log group would contain only errors, but it actually contains

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

@@ -152,7 +152,7 @@ class AwsLogsHook(AwsBaseHook):
          If the value is LastEventTime , the results are ordered by the event time. The default value is LogStreamName.
         :param count: The maximum number of items returned
         """
-        async with self.async_conn as client:
+        async with await self.get_async_conn() as client:
             try:
                 response: dict[str, Any] = await client.describe_log_streams(
                     logGroupName=log_group,
@@ -194,7 +194,7 @@ class AwsLogsHook(AwsBaseHook):
             else:
                 token_arg = {}
 
-            async with self.async_conn as client:
+            async with await self.get_async_conn() as client:
                 response = await client.get_log_events(
                     logGroupName=log_group,
                     logStreamName=log_stream_name,

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

@@ -67,7 +67,7 @@ class RedshiftHook(AwsBaseHook):
             for the cluster that is being created.
         :param params: Remaining AWS Create cluster API params.
         """
-        response = self.get_conn().create_cluster(
+        response = self.conn.create_cluster(
             ClusterIdentifier=cluster_identifier,
             NodeType=node_type,
             MasterUsername=master_username,
@@ -87,13 +87,13 @@ class RedshiftHook(AwsBaseHook):
         :param cluster_identifier: unique identifier of a cluster
         """
         try:
-            response = self.get_conn().describe_clusters(ClusterIdentifier=cluster_identifier)["Clusters"]
+            response = self.conn.describe_clusters(ClusterIdentifier=cluster_identifier)["Clusters"]
             return response[0]["ClusterStatus"] if response else None
-        except self.get_conn().exceptions.ClusterNotFoundFault:
+        except self.conn.exceptions.ClusterNotFoundFault:
             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
 
@@ -115,7 +115,7 @@ class RedshiftHook(AwsBaseHook):
         """
         final_cluster_snapshot_identifier = final_cluster_snapshot_identifier or ""
 
-        response = self.get_conn().delete_cluster(
+        response = self.conn.delete_cluster(
             ClusterIdentifier=cluster_identifier,
             SkipFinalClusterSnapshot=skip_final_cluster_snapshot,
             FinalClusterSnapshotIdentifier=final_cluster_snapshot_identifier,
@@ -131,7 +131,7 @@ class RedshiftHook(AwsBaseHook):
 
         :param cluster_identifier: unique identifier of a cluster
         """
-        response = self.get_conn().describe_cluster_snapshots(ClusterIdentifier=cluster_identifier)
+        response = self.conn.describe_cluster_snapshots(ClusterIdentifier=cluster_identifier)
         if "Snapshots" not in response:
             return None
         snapshots = response["Snapshots"]
@@ -149,7 +149,7 @@ class RedshiftHook(AwsBaseHook):
         :param cluster_identifier: unique identifier of a cluster
         :param snapshot_identifier: unique identifier for a snapshot of a cluster
         """
-        response = self.get_conn().restore_from_cluster_snapshot(
+        response = self.conn.restore_from_cluster_snapshot(
             ClusterIdentifier=cluster_identifier, SnapshotIdentifier=snapshot_identifier
         )
         return response["Cluster"] if response["Cluster"] else None
@@ -175,7 +175,7 @@ class RedshiftHook(AwsBaseHook):
         """
         if tags is None:
             tags = []
-        response = self.get_conn().create_cluster_snapshot(
+        response = self.conn.create_cluster_snapshot(
             SnapshotIdentifier=snapshot_identifier,
             ClusterIdentifier=cluster_identifier,
             ManualSnapshotRetentionPeriod=retention_period,
@@ -192,11 +192,11 @@ class RedshiftHook(AwsBaseHook):
         :param snapshot_identifier: A unique identifier for the snapshot that you are requesting
         """
         try:
-            response = self.get_conn().describe_cluster_snapshots(
+            response = self.conn.describe_cluster_snapshots(
                 SnapshotIdentifier=snapshot_identifier,
             )
             snapshot = response.get("Snapshots")[0]
             snapshot_status: str = snapshot.get("Status")
             return snapshot_status
-        except self.get_conn().exceptions.ClusterSnapshotNotFoundFault:
+        except self.conn.exceptions.ClusterSnapshotNotFoundFault:
             return None

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

@@ -186,8 +186,7 @@ class RedshiftDataHook(AwsGenericHook["RedshiftDataAPIServiceClient"]):
                 RedshiftDataQueryFailedError if status == FAILED_STATE else RedshiftDataQueryAbortedError
             )
             raise exception_cls(
-                f"Statement {resp['Id']} terminated with status {status}. "
-                f"Response details: {pformat(resp)}"
+                f"Statement {resp['Id']} terminated with status {status}. Response details: {pformat(resp)}"
             )
 
         self.log.info("Query status: %s", status)
@@ -275,7 +274,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 +287,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

@@ -131,7 +131,7 @@ def secondary_training_status_message(
     status_strs = []
     for transition in transitions_to_print:
         message = transition["StatusMessage"]
-        time_utc = timezone.convert_to_utc(cast(datetime, job_description["LastModifiedTime"]))
+        time_utc = timezone.convert_to_utc(cast("datetime", job_description["LastModifiedTime"]))
         status_strs.append(f"{time_utc:%Y-%m-%d %H:%M:%S} {transition['Status']} - {message}")
 
     return "\n".join(status_strs)
@@ -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/athena.py

@@ -25,6 +25,5 @@ class AthenaQueryResultsLink(BaseAwsLink):
     name = "Query Results"
     key = "_athena_query_results"
     format_str = (
-        BASE_AWS_CONSOLE_LINK + "/athena/home?region={region_name}#"
-        "/query-editor/history/{query_execution_id}"
+        BASE_AWS_CONSOLE_LINK + "/athena/home?region={region_name}#/query-editor/history/{query_execution_id}"
     )

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}"