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

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

@@ -44,7 +44,6 @@ from airflow.providers.amazon.aws.utils import trim_none_values, validate_execut
 from airflow.providers.amazon.aws.utils.sagemaker import ApprovalStatus
 from airflow.providers.amazon.aws.utils.tags import format_tags
 from airflow.utils.helpers import prune_dict
-from airflow.utils.json import AirflowJsonEncoder
 
 if TYPE_CHECKING:
     from airflow.providers.common.compat.openlineage.facet import Dataset
@@ -56,7 +55,7 @@ CHECK_INTERVAL_SECOND: int = 30
 
 
 def serialize(result: dict) -> dict:
-    return json.loads(json.dumps(result, cls=AirflowJsonEncoder))
+    return json.loads(json.dumps(result, default=repr))
 
 
 class SageMakerBaseOperator(BaseOperator):
@@ -171,7 +170,7 @@ class SageMakerBaseOperator(BaseOperator):
                 timestamp = str(
                     time.time_ns() // 1000000000
                 )  # only keep the relevant datetime (first 10 digits)
-                name = f"{proposed_name[:max_name_len - len(timestamp) - 1]}-{timestamp}"  # we subtract one to make provision for the dash between the truncated name and timestamp
+                name = f"{proposed_name[: max_name_len - len(timestamp) - 1]}-{timestamp}"  # we subtract one to make provision for the dash between the truncated name and timestamp
                 self.log.info("Changed %s name to '%s' to avoid collision.", resource_type, name)
         return name
 
@@ -179,8 +178,7 @@ class SageMakerBaseOperator(BaseOperator):
         """Raise exception if resource type is not 'model' or 'job'."""
         if resource_type not in ("model", "job"):
             raise AirflowException(
-                "Argument resource_type accepts only 'model' and 'job'. "
-                f"Provided value: '{resource_type}'."
+                f"Argument resource_type accepts only 'model' and 'job'. Provided value: '{resource_type}'."
             )
 
     def _check_if_job_exists(self, job_name: str, describe_func: Callable[[str], Any]) -> bool:
@@ -560,8 +558,7 @@ class SageMakerEndpointOperator(SageMakerBaseOperator):
                 self.operation = "update"
                 sagemaker_operation = self.hook.update_endpoint
                 self.log.warning(
-                    "cannot create already existing endpoint %s, "
-                    "updating it with the given config instead",
+                    "cannot create already existing endpoint %s, updating it with the given config instead",
                     endpoint_info["EndpointName"],
                 )
                 if "Tags" in endpoint_info:

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

@@ -0,0 +1,155 @@
+# 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 operator."""
+
+from __future__ import annotations
+
+from functools import cached_property
+from typing import TYPE_CHECKING
+
+from airflow.configuration import conf
+from airflow.exceptions import AirflowException
+from airflow.models import BaseOperator
+from airflow.providers.amazon.aws.hooks.sagemaker_unified_studio import (
+    SageMakerNotebookHook,
+)
+from airflow.providers.amazon.aws.links.sagemaker_unified_studio import (
+    SageMakerUnifiedStudioLink,
+)
+from airflow.providers.amazon.aws.triggers.sagemaker_unified_studio import (
+    SageMakerNotebookJobTrigger,
+)
+
+if TYPE_CHECKING:
+    from airflow.utils.context import Context
+
+
+class SageMakerNotebookOperator(BaseOperator):
+    """
+    Provides Artifact execution functionality for Sagemaker Unified Studio Workflows.
+
+    Examples:
+     .. code-block:: python
+
+        from airflow.providers.amazon.aws.operators.sagemaker_unified_studio import SageMakerNotebookOperator
+
+        notebook_operator = SageMakerNotebookOperator(
+            task_id="notebook_task",
+            input_config={"input_path": "path/to/notebook.ipynb", "input_params": ""},
+            output_config={"output_format": "ipynb"},
+            wait_for_completion=True,
+            waiter_delay=10,
+            waiter_max_attempts=1440,
+        )
+
+    :param task_id: A unique, meaningful id for the task.
+    :param input_config: Configuration for the input file. Input path should be specified as a relative path.
+        The provided relative path will be automatically resolved to an absolute path within
+        the context of the user's home directory in the IDE. Input params should be a dict.
+        Example: {'input_path': 'folder/input/notebook.ipynb', 'input_params':{'key': 'value'}}
+    :param output_config:  Configuration for the output format. It should include an output_format parameter to control
+        the format of the notebook execution output.
+        Example: {"output_formats": ["NOTEBOOK"]}
+    :param compute: compute configuration to use for the artifact execution. This is a required attribute
+        if the execution is on a remote compute.
+        Example: { "InstanceType": "ml.m5.large", "VolumeSizeInGB": 30, "VolumeKmsKeyId": "", "ImageUri": "string", "ContainerEntrypoint": [ "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 wait_for_completion: Indicates whether to wait for the notebook execution to complete. If True, wait for completion; if False, don't wait.
+    :param waiter_delay: Interval in seconds to check the notebook execution status.
+    :param waiter_max_attempts: Number of attempts to wait before returning FAILED.
+    :param deferrable: If True, the operator will wait asynchronously for the job to complete.
+        This implies waiting for completion. This mode requires aiobotocore module to be installed.
+        (default: False)
+
+    .. seealso::
+        For more information on how to use this operator, take a look at the guide:
+        :ref:`howto/operator:SageMakerNotebookOperator`
+    """
+
+    operator_extra_links = (SageMakerUnifiedStudioLink(),)
+
+    def __init__(
+        self,
+        task_id: str,
+        input_config: dict,
+        output_config: dict | None = None,
+        compute: dict | None = None,
+        termination_condition: dict | None = None,
+        tags: dict | None = None,
+        wait_for_completion: bool = True,
+        waiter_delay: int = 10,
+        waiter_max_attempts: int = 1440,
+        deferrable: bool = conf.getboolean("operators", "default_deferrable", fallback=False),
+        **kwargs,
+    ):
+        super().__init__(task_id=task_id, **kwargs)
+        self.execution_name = task_id
+        self.input_config = input_config
+        self.output_config = output_config or {"output_formats": ["NOTEBOOK"]}
+        self.compute = compute or {}
+        self.termination_condition = termination_condition or {}
+        self.tags = tags or {}
+        self.wait_for_completion = wait_for_completion
+        self.waiter_delay = waiter_delay
+        self.waiter_max_attempts = waiter_max_attempts
+        self.deferrable = deferrable
+        self.input_kwargs = kwargs
+
+    @cached_property
+    def notebook_execution_hook(self):
+        if not self.input_config:
+            raise AirflowException("input_config is required")
+
+        if "input_path" not in self.input_config:
+            raise AirflowException("input_path is a required field in the input_config")
+
+        return SageMakerNotebookHook(
+            input_config=self.input_config,
+            output_config=self.output_config,
+            execution_name=self.execution_name,
+            compute=self.compute,
+            termination_condition=self.termination_condition,
+            tags=self.tags,
+            waiter_delay=self.waiter_delay,
+            waiter_max_attempts=self.waiter_max_attempts,
+        )
+
+    def execute(self, context: Context):
+        notebook_execution = self.notebook_execution_hook.start_notebook_execution()
+        execution_id = notebook_execution["execution_id"]
+
+        if self.deferrable:
+            self.defer(
+                trigger=SageMakerNotebookJobTrigger(
+                    execution_id=execution_id,
+                    execution_name=self.execution_name,
+                    waiter_delay=self.waiter_delay,
+                    waiter_max_attempts=self.waiter_max_attempts,
+                ),
+                method_name="execute_complete",
+            )
+        elif self.wait_for_completion:
+            response = self.notebook_execution_hook.wait_for_execution_completion(execution_id, context)
+            status = response["Status"]
+            log_info_message = (
+                f"Notebook Execution: {self.execution_name} Status: {status}. Run Id: {execution_id}"
+            )
+            self.log.info(log_info_message)

airflow/providers/amazon/aws/sensors/ec2.py

@@ -18,21 +18,21 @@
 from __future__ import annotations
 
 from collections.abc import Sequence
-from functools import cached_property
 from typing import TYPE_CHECKING, Any
 
 from airflow.configuration import conf
 from airflow.exceptions import AirflowException
 from airflow.providers.amazon.aws.hooks.ec2 import EC2Hook
+from airflow.providers.amazon.aws.sensors.base_aws import AwsBaseSensor
 from airflow.providers.amazon.aws.triggers.ec2 import EC2StateSensorTrigger
 from airflow.providers.amazon.aws.utils import validate_execute_complete_event
-from airflow.sensors.base import BaseSensorOperator
+from airflow.providers.amazon.aws.utils.mixins import aws_template_fields
 
 if TYPE_CHECKING:
     from airflow.utils.context import Context
 
 
-class EC2InstanceStateSensor(BaseSensorOperator):
+class EC2InstanceStateSensor(AwsBaseSensor[EC2Hook]):
     """
     Poll the state of the AWS EC2 instance until the instance reaches the target state.
 
@@ -46,7 +46,8 @@ class EC2InstanceStateSensor(BaseSensorOperator):
     :param deferrable: if True, the sensor will run in deferrable mode
     """
 
-    template_fields: Sequence[str] = ("target_state", "instance_id", "region_name")
+    aws_hook_class = EC2Hook
+    template_fields: Sequence[str] = aws_template_fields("target_state", "instance_id", "region_name")
     ui_color = "#cc8811"
     ui_fgcolor = "#ffffff"
     valid_states = ["running", "stopped", "terminated"]
@@ -56,8 +57,6 @@ class EC2InstanceStateSensor(BaseSensorOperator):
         *,
         target_state: str,
         instance_id: str,
-        aws_conn_id: str | None = "aws_default",
-        region_name: str | None = None,
         deferrable: bool = conf.getboolean("operators", "default_deferrable", fallback=False),
         **kwargs,
     ):
@@ -66,8 +65,6 @@ class EC2InstanceStateSensor(BaseSensorOperator):
         super().__init__(**kwargs)
         self.target_state = target_state
         self.instance_id = instance_id
-        self.aws_conn_id = aws_conn_id
-        self.region_name = region_name
         self.deferrable = deferrable
 
     def execute(self, context: Context) -> Any:
@@ -85,10 +82,6 @@ class EC2InstanceStateSensor(BaseSensorOperator):
         else:
             super().execute(context=context)
 
-    @cached_property
-    def hook(self):
-        return EC2Hook(aws_conn_id=self.aws_conn_id, region_name=self.region_name)
-
     def poke(self, context: Context):
         instance_state = self.hook.get_instance_state(instance_id=self.instance_id)
         self.log.info("instance state: %s", instance_state)

airflow/providers/amazon/aws/sensors/emr.py

@@ -311,7 +311,7 @@ class EmrContainerSensor(BaseSensorOperator):
         )
 
         if state in self.FAILURE_STATES:
-            raise AirflowException("EMR Containers sensor failed")
+            raise AirflowException(f"EMR Containers sensor failed due to state: {state}")
 
         if state in self.INTERMEDIATE_STATES:
             return False

airflow/providers/amazon/aws/sensors/glacier.py

@@ -95,5 +95,5 @@ class GlacierJobOperationSensor(AwsBaseSensor[GlacierHook]):
             return False
         else:
             raise AirflowException(
-                f'Sensor failed. Job status: {response["Action"]}, code status: {response["StatusCode"]}'
+                f"Sensor failed. Job status: {response['Action']}, code status: {response['StatusCode']}"
             )

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

@@ -0,0 +1,161 @@
+#
+# 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 collections.abc import Collection, Sequence
+from typing import TYPE_CHECKING, Any
+
+from airflow.configuration import conf
+from airflow.exceptions import AirflowException
+from airflow.providers.amazon.aws.hooks.mwaa import MwaaHook
+from airflow.providers.amazon.aws.sensors.base_aws import AwsBaseSensor
+from airflow.providers.amazon.aws.triggers.mwaa import MwaaDagRunCompletedTrigger
+from airflow.providers.amazon.aws.utils import validate_execute_complete_event
+from airflow.providers.amazon.aws.utils.mixins import aws_template_fields
+from airflow.utils.state import DagRunState
+
+if TYPE_CHECKING:
+    from airflow.utils.context import Context
+
+
+class MwaaDagRunSensor(AwsBaseSensor[MwaaHook]):
+    """
+    Waits for a DAG Run in an MWAA Environment to complete.
+
+    If the DAG Run fails, an AirflowException is thrown.
+
+    .. seealso::
+        For more information on how to use this sensor, take a look at the guide:
+        :ref:`howto/sensor:MwaaDagRunSensor`
+
+    :param external_env_name: The external MWAA environment name that contains the DAG Run you want to wait for
+        (templated)
+    :param external_dag_id: The DAG ID in the external MWAA environment that contains the DAG Run you want to wait for
+        (templated)
+    :param external_dag_run_id: The DAG Run ID in the external MWAA environment that you want to wait for (templated)
+    :param success_states: Collection of DAG Run states that would make this task marked as successful, default is
+        ``{airflow.utils.state.DagRunState.SUCCESS}`` (templated)
+    :param failure_states: Collection of DAG Run states that would make this task marked as failed and raise an
+        AirflowException, default is ``{airflow.utils.state.DagRunState.FAILED}`` (templated)
+    :param deferrable: If True, the sensor will operate in deferrable mode. This mode requires aiobotocore
+        module to be installed.
+        (default: False, but can be overridden in config file by setting default_deferrable to True)
+    :param poke_interval: Polling period in seconds to check for the status of the job. (default: 60)
+    :param max_retries: Number of times before returning the current state. (default: 720)
+    :param aws_conn_id: The Airflow connection used for AWS credentials.
+        If this is ``None`` or empty then the default boto3 behaviour is used. If
+        running Airflow in a distributed manner and aws_conn_id is None or
+        empty, then default boto3 configuration would be used (and must be
+        maintained on each worker node).
+    :param region_name: AWS region_name. If not specified then the default boto3 behaviour is used.
+    :param verify: Whether or not to verify SSL certificates. See:
+        https://boto3.amazonaws.com/v1/documentation/api/latest/reference/core/session.html
+    :param botocore_config: Configuration dictionary (key-values) for botocore client. See:
+        https://botocore.amazonaws.com/v1/documentation/api/latest/reference/config.html
+    """
+
+    aws_hook_class = MwaaHook
+    template_fields: Sequence[str] = aws_template_fields(
+        "external_env_name",
+        "external_dag_id",
+        "external_dag_run_id",
+        "success_states",
+        "failure_states",
+        "deferrable",
+        "max_retries",
+        "poke_interval",
+    )
+
+    def __init__(
+        self,
+        *,
+        external_env_name: str,
+        external_dag_id: str,
+        external_dag_run_id: str,
+        success_states: Collection[str] | None = None,
+        failure_states: Collection[str] | None = None,
+        deferrable: bool = conf.getboolean("operators", "default_deferrable", fallback=False),
+        poke_interval: int = 60,
+        max_retries: int = 720,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+
+        self.success_states = set(success_states) if success_states else {DagRunState.SUCCESS.value}
+        self.failure_states = set(failure_states) if failure_states else {DagRunState.FAILED.value}
+
+        if len(self.success_states & self.failure_states):
+            raise ValueError("success_states and failure_states must not have any values in common")
+
+        self.external_env_name = external_env_name
+        self.external_dag_id = external_dag_id
+        self.external_dag_run_id = external_dag_run_id
+        self.deferrable = deferrable
+        self.poke_interval = poke_interval
+        self.max_retries = max_retries
+
+    def poke(self, context: Context) -> bool:
+        self.log.info(
+            "Poking for DAG run %s of DAG %s in MWAA environment %s",
+            self.external_dag_run_id,
+            self.external_dag_id,
+            self.external_env_name,
+        )
+        response = self.hook.invoke_rest_api(
+            env_name=self.external_env_name,
+            path=f"/dags/{self.external_dag_id}/dagRuns/{self.external_dag_run_id}",
+            method="GET",
+        )
+
+        # If RestApiStatusCode == 200, the RestApiResponse must have the "state" key, otherwise something terrible has
+        # happened in the API and KeyError would be raised
+        # If RestApiStatusCode >= 300, a botocore exception would've already been raised during the
+        # self.hook.invoke_rest_api call
+        # The scope of this sensor is going to only be raising AirflowException due to failure of the DAGRun
+
+        state = response["RestApiResponse"]["state"]
+
+        if state in self.failure_states:
+            raise AirflowException(
+                f"The DAG run {self.external_dag_run_id} of DAG {self.external_dag_id} in MWAA environment {self.external_env_name} "
+                f"failed with state: {state}"
+            )
+
+        return state in self.success_states
+
+    def execute_complete(self, context: Context, event: dict[str, Any] | None = None) -> None:
+        validate_execute_complete_event(event)
+
+    def execute(self, context: Context):
+        if self.deferrable:
+            self.defer(
+                trigger=MwaaDagRunCompletedTrigger(
+                    external_env_name=self.external_env_name,
+                    external_dag_id=self.external_dag_id,
+                    external_dag_run_id=self.external_dag_run_id,
+                    success_states=self.success_states,
+                    failure_states=self.failure_states,
+                    # somehow the type of poke_interval is derived as float ??
+                    waiter_delay=self.poke_interval,  # type: ignore[arg-type]
+                    waiter_max_attempts=self.max_retries,
+                    aws_conn_id=self.aws_conn_id,
+                ),
+                method_name="execute_complete",
+            )
+        else:
+            super().execute(context=context)

airflow/providers/amazon/aws/sensors/rds.py

@@ -20,7 +20,7 @@ from collections.abc import Sequence
 from functools import cached_property
 from typing import TYPE_CHECKING
 
-from airflow.exceptions import AirflowNotFoundException
+from airflow.exceptions import AirflowException, AirflowNotFoundException
 from airflow.providers.amazon.aws.hooks.rds import RdsHook
 from airflow.providers.amazon.aws.utils.rds import RdsDbType
 from airflow.sensors.base import BaseSensorOperator
@@ -104,18 +104,17 @@ class RdsExportTaskExistenceSensor(RdsBaseSensor):
 
     :param export_task_identifier: A unique identifier for the snapshot export task.
     :param target_statuses: Target status of export task
+    :param error_statuses: Target error status of export task to fail the sensor
     """
 
-    template_fields: Sequence[str] = (
-        "export_task_identifier",
-        "target_statuses",
-    )
+    template_fields: Sequence[str] = ("export_task_identifier", "target_statuses", "error_statuses")
 
     def __init__(
         self,
         *,
         export_task_identifier: str,
         target_statuses: list[str] | None = None,
+        error_statuses: list[str] | None = None,
         aws_conn_id: str | None = "aws_default",
         **kwargs,
     ):
@@ -129,6 +128,7 @@ class RdsExportTaskExistenceSensor(RdsBaseSensor):
             "canceling",
             "canceled",
         ]
+        self.error_statuses = error_statuses or ["failed"]
 
     def poke(self, context: Context):
         self.log.info(
@@ -136,6 +136,11 @@ class RdsExportTaskExistenceSensor(RdsBaseSensor):
         )
         try:
             state = self.hook.get_export_task_state(self.export_task_identifier)
+            if state in self.error_statuses:
+                raise AirflowException(
+                    f"Export task {self.export_task_identifier} failed with status {state}"
+                )
+
         except AirflowNotFoundException:
             return False
         return state in self.target_statuses