apache-airflow-providers-amazon 9.19.0rc1__py3-none-any.whl → 9.21.0__py3-none-any.whl

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

@@ -35,12 +35,7 @@ from airflow.providers.amazon.aws.hooks.base_aws import AwsBaseHook
 from airflow.providers.amazon.aws.hooks.logs import AwsLogsHook
 from airflow.providers.amazon.aws.hooks.s3 import S3Hook
 from airflow.providers.amazon.aws.utils.tags import format_tags
-from airflow.providers.common.compat.sdk import AirflowException
-
-try:
-    from airflow.sdk import timezone
-except ImportError:
-    from airflow.utils import timezone  # type: ignore[attr-defined,no-redef]
+from airflow.providers.common.compat.sdk import AirflowException, timezone
 
 
 class LogState:

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

@@ -52,13 +52,12 @@ class SageMakerNotebookHook(BaseHook):
         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 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_details": {"ecr_uri": "string"}, "container_entrypoint": ["string"]}
     :param termination_condition: conditions to match to terminate the remote execution.
-        Example: { "MaxRuntimeInSeconds": 3600 }
+        Example: {"MaxRuntimeInSeconds": 3600}
     :param tags: tags to be associated with the remote execution runs.
-        Example: { "md_analytics": "logs" }
+        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.
     """

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

@@ -99,3 +99,25 @@ class SsmHook(AwsBaseHook):
         :return: Response from SSM list_command_invocations API.
         """
         return self.conn.list_command_invocations(CommandId=command_id)
+
+    @staticmethod
+    def is_aws_level_failure(status: str) -> bool:
+        """
+        Check if a command status represents an AWS-level failure.
+
+        AWS-level failures are service-level issues that should always raise exceptions,
+        as opposed to command-level failures (non-zero exit codes) which may be tolerated
+        depending on the fail_on_nonzero_exit parameter.
+
+        According to AWS SSM documentation, the possible statuses are:
+        Pending, InProgress, Delayed, Success, Cancelled, TimedOut, Failed, Cancelling
+
+        AWS-level failures are:
+        - Cancelled: Command was cancelled before completion
+        - TimedOut: Command exceeded the timeout period
+        - Cancelling: Command is in the process of being cancelled
+
+        :param status: The command invocation status from SSM.
+        :return: True if the status represents an AWS-level failure, False otherwise.
+        """
+        return status in ("Cancelled", "TimedOut", "Cancelling")

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

@@ -24,7 +24,7 @@ from airflow.providers.common.compat.sdk import BaseOperatorLink, XCom
 
 if TYPE_CHECKING:
     from airflow.models import BaseOperator
-    from airflow.models.taskinstancekey import TaskInstanceKey
+    from airflow.providers.common.compat.sdk import TaskInstanceKey
     from airflow.sdk import Context
 
 

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

@@ -31,9 +31,9 @@ from typing import TYPE_CHECKING, Any
 import attrs
 import watchtower
 
-from airflow.configuration import conf
 from airflow.providers.amazon.aws.hooks.logs import AwsLogsHook
 from airflow.providers.amazon.aws.utils import datetime_to_epoch_utc_ms
+from airflow.providers.common.compat.sdk import conf
 from airflow.utils.log.file_task_handler import FileTaskHandler
 from airflow.utils.log.logging_mixin import LoggingMixin
 

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

@@ -26,9 +26,9 @@ from typing import TYPE_CHECKING
 
 import attrs
 
-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.providers.common.compat.sdk import conf
 from airflow.utils.log.file_task_handler import FileTaskHandler
 from airflow.utils.log.logging_mixin import LoggingMixin
 

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

@@ -21,14 +21,13 @@ from collections.abc import Sequence
 from typing import TYPE_CHECKING, Any
 from urllib.parse import urlparse
 
-from airflow.configuration import conf
 from airflow.providers.amazon.aws.hooks.athena import AthenaHook
 from airflow.providers.amazon.aws.links.athena import AthenaQueryResultsLink
 from airflow.providers.amazon.aws.operators.base_aws import AwsBaseOperator
 from airflow.providers.amazon.aws.triggers.athena import AthenaTrigger
 from airflow.providers.amazon.aws.utils import validate_execute_complete_event
 from airflow.providers.amazon.aws.utils.mixins import aws_template_fields
-from airflow.providers.common.compat.sdk import AirflowException
+from airflow.providers.common.compat.sdk import AirflowException, conf
 
 if TYPE_CHECKING:
     from airflow.providers.common.compat.openlineage.facet import BaseFacet, Dataset, DatasetFacet

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

@@ -30,7 +30,6 @@ from collections.abc import Sequence
 from datetime import timedelta
 from typing import TYPE_CHECKING, Any
 
-from airflow.configuration import conf
 from airflow.providers.amazon.aws.hooks.batch_client import BatchClientHook
 from airflow.providers.amazon.aws.links.batch import (
     BatchJobDefinitionLink,
@@ -46,7 +45,7 @@ from airflow.providers.amazon.aws.triggers.batch import (
 from airflow.providers.amazon.aws.utils import trim_none_values, validate_execute_complete_event
 from airflow.providers.amazon.aws.utils.mixins import aws_template_fields
 from airflow.providers.amazon.aws.utils.task_log_fetcher import AwsTaskLogFetcher
-from airflow.providers.common.compat.sdk import AirflowException
+from airflow.providers.common.compat.sdk import AirflowException, conf
 
 if TYPE_CHECKING:
     from airflow.sdk import Context

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

@@ -23,7 +23,6 @@ from typing import TYPE_CHECKING, Any
 
 from botocore.exceptions import ClientError
 
-from airflow.configuration import conf
 from airflow.providers.amazon.aws.hooks.bedrock import (
     BedrockAgentHook,
     BedrockAgentRuntimeHook,
@@ -40,7 +39,7 @@ from airflow.providers.amazon.aws.triggers.bedrock import (
 )
 from airflow.providers.amazon.aws.utils import validate_execute_complete_event
 from airflow.providers.amazon.aws.utils.mixins import aws_template_fields
-from airflow.providers.common.compat.sdk import AirflowException, timezone
+from airflow.providers.common.compat.sdk import AirflowException, conf, timezone
 from airflow.utils.helpers import prune_dict
 
 if TYPE_CHECKING:

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

@@ -20,7 +20,6 @@ from collections.abc import Sequence
 from functools import cached_property
 from typing import TYPE_CHECKING, Any, ClassVar
 
-from airflow.configuration import conf
 from airflow.providers.amazon.aws.hooks.comprehend import ComprehendHook
 from airflow.providers.amazon.aws.links.comprehend import (
     ComprehendDocumentClassifierLink,
@@ -33,7 +32,7 @@ from airflow.providers.amazon.aws.triggers.comprehend import (
 )
 from airflow.providers.amazon.aws.utils import validate_execute_complete_event
 from airflow.providers.amazon.aws.utils.mixins import aws_template_fields
-from airflow.providers.common.compat.sdk import AirflowException, timezone
+from airflow.providers.common.compat.sdk import AirflowException, conf, timezone
 
 if TYPE_CHECKING:
     import boto3

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

@@ -21,7 +21,6 @@ from collections.abc import Sequence
 from datetime import datetime
 from typing import Any, ClassVar
 
-from airflow.configuration import conf
 from airflow.providers.amazon.aws.hooks.dms import DmsHook
 from airflow.providers.amazon.aws.operators.base_aws import AwsBaseOperator
 from airflow.providers.amazon.aws.triggers.dms import (
@@ -32,7 +31,7 @@ from airflow.providers.amazon.aws.triggers.dms import (
     DmsReplicationTerminalStatusTrigger,
 )
 from airflow.providers.amazon.aws.utils.mixins import aws_template_fields
-from airflow.providers.common.compat.sdk import AirflowException, Context
+from airflow.providers.common.compat.sdk import AirflowException, Context, conf
 
 
 class DmsCreateTaskOperator(AwsBaseOperator[DmsHook]):

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

@@ -218,31 +218,46 @@ class EC2CreateInstanceOperator(AwsBaseOperator[EC2Hook]):
             MaxCount=self.max_count,
             **self.config,
         )["Instances"]
-
-        instance_ids = self._on_kill_instance_ids = [instance["InstanceId"] for instance in instances]
-        # Console link is for EC2 dashboard list, not individual instances when more than 1 instance
-
-        EC2InstanceDashboardLink.persist(
-            context=context,
-            operator=self,
-            region_name=self.hook.conn_region_name,
-            aws_partition=self.hook.conn_partition,
-            instance_ids=EC2InstanceDashboardLink.format_instance_id_filter(instance_ids),
-        )
-        for instance_id in instance_ids:
-            self.log.info("Created EC2 instance %s", instance_id)
-
-            if self.wait_for_completion:
-                self.hook.get_waiter("instance_running").wait(
-                    InstanceIds=[instance_id],
-                    WaiterConfig={
-                        "Delay": self.poll_interval,
-                        "MaxAttempts": self.max_attempts,
-                    },
+        try:
+            instance_ids = self._on_kill_instance_ids = [instance["InstanceId"] for instance in instances]
+            # Console link is for EC2 dashboard list, not individual instances when more than 1 instance
+
+            EC2InstanceDashboardLink.persist(
+                context=context,
+                operator=self,
+                region_name=self.hook.conn_region_name,
+                aws_partition=self.hook.conn_partition,
+                instance_ids=EC2InstanceDashboardLink.format_instance_id_filter(instance_ids),
+            )
+            for instance_id in instance_ids:
+                self.log.info("Created EC2 instance %s", instance_id)
+
+                if self.wait_for_completion:
+                    self.hook.get_waiter("instance_running").wait(
+                        InstanceIds=[instance_id],
+                        WaiterConfig={
+                            "Delay": self.poll_interval,
+                            "MaxAttempts": self.max_attempts,
+                        },
+                    )
+
+            # leave "_on_kill_instance_ids" in place for finishing post-processing
+            return instance_ids
+
+        # Best-effort cleanup when post-creation steps fail (e.g. IAM/permission errors).
+        except Exception:
+            self.log.exception(
+                "Exception after EC2 instance creation; attempting cleanup for instances %s",
+                instance_ids,
+            )
+            try:
+                self.hook.terminate_instances(instance_ids=instance_ids)
+            except Exception:
+                self.log.exception(
+                    "Failed to cleanup EC2 instances %s after task failure",
+                    instance_ids,
                 )
-
-        # leave "_on_kill_instance_ids" in place for finishing post-processing
-        return instance_ids
+            raise
 
     def on_kill(self) -> None:
         instance_ids = getattr(self, "_on_kill_instance_ids", [])

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

@@ -24,7 +24,6 @@ from functools import cached_property
 from time import sleep
 from typing import TYPE_CHECKING, Any
 
-from airflow.configuration import conf
 from airflow.providers.amazon.aws.exceptions import EcsOperatorError, EcsTaskFailToStart
 from airflow.providers.amazon.aws.hooks.ecs import EcsClusterStates, EcsHook
 from airflow.providers.amazon.aws.hooks.logs import AwsLogsHook
@@ -38,7 +37,7 @@ from airflow.providers.amazon.aws.utils import validate_execute_complete_event
 from airflow.providers.amazon.aws.utils.identifiers import generate_uuid
 from airflow.providers.amazon.aws.utils.mixins import aws_template_fields
 from airflow.providers.amazon.aws.utils.task_log_fetcher import AwsTaskLogFetcher
-from airflow.providers.common.compat.sdk import AirflowException
+from airflow.providers.common.compat.sdk import AirflowException, conf
 from airflow.utils.helpers import prune_dict
 
 if TYPE_CHECKING:

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

@@ -27,7 +27,6 @@ from typing import TYPE_CHECKING, Any, cast
 
 from botocore.exceptions import ClientError, WaiterError
 
-from airflow.configuration import conf
 from airflow.exceptions import AirflowProviderDeprecationWarning
 from airflow.providers.amazon.aws.hooks.eks import EksHook
 from airflow.providers.amazon.aws.operators.base_aws import AwsBaseOperator
@@ -43,7 +42,7 @@ from airflow.providers.amazon.aws.utils import validate_execute_complete_event
 from airflow.providers.amazon.aws.utils.mixins import aws_template_fields
 from airflow.providers.amazon.aws.utils.waiter_with_logging import wait
 from airflow.providers.cncf.kubernetes.utils.pod_manager import OnFinishAction
-from airflow.providers.common.compat.sdk import AirflowException
+from airflow.providers.common.compat.sdk import AirflowException, conf
 
 try:
     from airflow.providers.cncf.kubernetes.operators.pod import KubernetesPodOperator

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

@@ -24,7 +24,6 @@ from datetime import timedelta
 from typing import TYPE_CHECKING, Any
 from uuid import uuid4
 
-from airflow.configuration import conf
 from airflow.exceptions import AirflowProviderDeprecationWarning
 from airflow.providers.amazon.aws.hooks.emr import EmrContainerHook, EmrHook, EmrServerlessHook
 from airflow.providers.amazon.aws.links.emr import (
@@ -58,7 +57,7 @@ from airflow.providers.amazon.aws.utils.waiter import (
 )
 from airflow.providers.amazon.aws.utils.waiter_with_logging import wait
 from airflow.providers.amazon.version_compat import NOTSET, ArgNotSet
-from airflow.providers.common.compat.sdk import AirflowException
+from airflow.providers.common.compat.sdk import AirflowException, conf
 from airflow.utils.helpers import exactly_one, prune_dict
 
 if TYPE_CHECKING:
@@ -1275,16 +1274,32 @@ class EmrServerlessStartJobOperator(AwsBaseOperator[EmrServerlessHook]):
                     timeout=timedelta(seconds=self.waiter_max_attempts * self.waiter_delay),
                 )
             else:
-                waiter = self.hook.get_waiter("serverless_job_completed")
-                wait(
-                    waiter=waiter,
-                    waiter_max_attempts=self.waiter_max_attempts,
-                    waiter_delay=self.waiter_delay,
-                    args={"applicationId": self.application_id, "jobRunId": self.job_id},
-                    failure_message="Serverless Job failed",
-                    status_message="Serverless Job status is",
-                    status_args=["jobRun.state", "jobRun.stateDetails"],
-                )
+                try:
+                    waiter = self.hook.get_waiter("serverless_job_completed")
+                    wait(
+                        waiter=waiter,
+                        waiter_max_attempts=self.waiter_max_attempts,
+                        waiter_delay=self.waiter_delay,
+                        args={"applicationId": self.application_id, "jobRunId": self.job_id},
+                        failure_message="Serverless Job failed",
+                        status_message="Serverless Job status is",
+                        status_args=["jobRun.state", "jobRun.stateDetails"],
+                    )
+                except AirflowException as e:
+                    if "Waiter error: max attempts reached" in str(e):
+                        self.log.info(
+                            "Cancelling EMR Serverless job %s due to max waiter attempts reached", self.job_id
+                        )
+                        try:
+                            self.hook.conn.cancel_job_run(
+                                applicationId=self.application_id, jobRunId=self.job_id
+                            )
+                        except Exception:
+                            self.log.exception(
+                                "Failed to cancel EMR Serverless job %s after waiter timeout",
+                                self.job_id,
+                            )
+                    raise
 
         return self.job_id
 
@@ -1293,7 +1308,13 @@ class EmrServerlessStartJobOperator(AwsBaseOperator[EmrServerlessHook]):
 
         if validated_event["status"] == "success":
             self.log.info("Serverless job completed")
-            return validated_event["job_id"]
+            return validated_event["job_details"]["job_id"]
+        self.log.info("Cancelling EMR Serverless job %s", self.job_id)
+        self.hook.conn.cancel_job_run(
+            applicationId=validated_event["job_details"]["application_id"],
+            jobRunId=validated_event["job_details"]["job_id"],
+        )
+        raise AirflowException("EMR Serverless job failed or timed out in deferrable mode")
 
     def on_kill(self) -> None:
         """

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

@@ -24,7 +24,6 @@ from typing import TYPE_CHECKING, Any
 
 from botocore.exceptions import ClientError
 
-from airflow.configuration import conf
 from airflow.providers.amazon.aws.hooks.glue import GlueDataQualityHook, GlueJobHook
 from airflow.providers.amazon.aws.hooks.s3 import S3Hook
 from airflow.providers.amazon.aws.links.glue import GlueJobRunDetailsLink
@@ -36,7 +35,7 @@ from airflow.providers.amazon.aws.triggers.glue import (
 )
 from airflow.providers.amazon.aws.utils import validate_execute_complete_event
 from airflow.providers.amazon.aws.utils.mixins import aws_template_fields
-from airflow.providers.common.compat.sdk import AirflowException
+from airflow.providers.common.compat.sdk import AirflowException, conf
 
 if TYPE_CHECKING:
     from airflow.sdk import Context

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

@@ -20,13 +20,12 @@ from __future__ import annotations
 from collections.abc import Sequence
 from typing import TYPE_CHECKING, Any
 
-from airflow.configuration import conf
 from airflow.providers.amazon.aws.hooks.glue_crawler import GlueCrawlerHook
 from airflow.providers.amazon.aws.operators.base_aws import AwsBaseOperator
 from airflow.providers.amazon.aws.triggers.glue_crawler import GlueCrawlerCompleteTrigger
 from airflow.providers.amazon.aws.utils import validate_execute_complete_event
 from airflow.providers.amazon.aws.utils.mixins import aws_template_fields
-from airflow.providers.common.compat.sdk import AirflowException
+from airflow.providers.common.compat.sdk import AirflowException, conf
 
 if TYPE_CHECKING:
     from airflow.sdk import Context

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

@@ -20,13 +20,12 @@ from __future__ import annotations
 from collections.abc import Sequence
 from typing import TYPE_CHECKING, Any
 
-from airflow.configuration import conf
 from airflow.providers.amazon.aws.hooks.glue_databrew import GlueDataBrewHook
 from airflow.providers.amazon.aws.operators.base_aws import AwsBaseOperator
 from airflow.providers.amazon.aws.triggers.glue_databrew import GlueDataBrewJobCompleteTrigger
 from airflow.providers.amazon.aws.utils import validate_execute_complete_event
 from airflow.providers.amazon.aws.utils.mixins import aws_template_fields
-from airflow.providers.common.compat.sdk import AirflowException
+from airflow.providers.common.compat.sdk import AirflowException, conf
 
 if TYPE_CHECKING:
     from airflow.sdk import Context

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

@@ -21,7 +21,6 @@ from typing import TYPE_CHECKING, Any, ClassVar
 
 from botocore.exceptions import ClientError
 
-from airflow.configuration import conf
 from airflow.providers.amazon.aws.hooks.kinesis_analytics import KinesisAnalyticsV2Hook
 from airflow.providers.amazon.aws.operators.base_aws import AwsBaseOperator
 from airflow.providers.amazon.aws.triggers.kinesis_analytics import (
@@ -29,7 +28,7 @@ from airflow.providers.amazon.aws.triggers.kinesis_analytics import (
 )
 from airflow.providers.amazon.aws.utils import validate_execute_complete_event
 from airflow.providers.amazon.aws.utils.mixins import aws_template_fields
-from airflow.providers.common.compat.sdk import AirflowException
+from airflow.providers.common.compat.sdk import AirflowException, conf
 
 if TYPE_CHECKING:
     from airflow.sdk import Context

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

@@ -22,13 +22,12 @@ from collections.abc import Sequence
 from datetime import timedelta
 from typing import TYPE_CHECKING, Any
 
-from airflow.configuration import conf
 from airflow.providers.amazon.aws.hooks.lambda_function import LambdaHook
 from airflow.providers.amazon.aws.operators.base_aws import AwsBaseOperator
 from airflow.providers.amazon.aws.triggers.lambda_function import LambdaCreateFunctionCompleteTrigger
 from airflow.providers.amazon.aws.utils import validate_execute_complete_event
 from airflow.providers.amazon.aws.utils.mixins import aws_template_fields
-from airflow.providers.common.compat.sdk import AirflowException
+from airflow.providers.common.compat.sdk import AirflowException, conf
 
 if TYPE_CHECKING:
     from airflow.sdk import Context

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

@@ -21,13 +21,12 @@ from __future__ import annotations
 from collections.abc import Sequence
 from typing import TYPE_CHECKING, Any, Literal
 
-from airflow.configuration import conf
 from airflow.providers.amazon.aws.hooks.mwaa import MwaaHook
 from airflow.providers.amazon.aws.operators.base_aws import AwsBaseOperator
 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.providers.common.compat.sdk import AirflowException
+from airflow.providers.common.compat.sdk import AirflowException, conf
 
 if TYPE_CHECKING:
     from airflow.sdk import Context

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

@@ -22,7 +22,6 @@ from typing import TYPE_CHECKING, Any
 
 from botocore.exceptions import ClientError
 
-from airflow.configuration import conf
 from airflow.providers.amazon.aws.hooks.neptune import NeptuneHook
 from airflow.providers.amazon.aws.operators.base_aws import AwsBaseOperator
 from airflow.providers.amazon.aws.triggers.neptune import (
@@ -31,7 +30,7 @@ from airflow.providers.amazon.aws.triggers.neptune import (
     NeptuneClusterStoppedTrigger,
 )
 from airflow.providers.amazon.aws.utils.mixins import aws_template_fields
-from airflow.providers.common.compat.sdk import AirflowException
+from airflow.providers.common.compat.sdk import AirflowException, conf
 
 if TYPE_CHECKING:
     from airflow.sdk import Context

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

@@ -22,7 +22,6 @@ from collections.abc import Sequence
 from datetime import timedelta
 from typing import TYPE_CHECKING, Any
 
-from airflow.configuration import conf
 from airflow.providers.amazon.aws.hooks.rds import RdsHook
 from airflow.providers.amazon.aws.operators.base_aws import AwsBaseOperator
 from airflow.providers.amazon.aws.triggers.rds import (
@@ -35,7 +34,7 @@ from airflow.providers.amazon.aws.utils.mixins import aws_template_fields
 from airflow.providers.amazon.aws.utils.rds import RdsDbType
 from airflow.providers.amazon.aws.utils.tags import format_tags
 from airflow.providers.amazon.aws.utils.waiter_with_logging import wait
-from airflow.providers.common.compat.sdk import AirflowException
+from airflow.providers.common.compat.sdk import AirflowException, conf
 from airflow.utils.helpers import prune_dict
 
 if TYPE_CHECKING:

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

@@ -21,7 +21,6 @@ from collections.abc import Sequence
 from datetime import timedelta
 from typing import TYPE_CHECKING, Any
 
-from airflow.configuration import conf
 from airflow.providers.amazon.aws.hooks.redshift_cluster import RedshiftHook
 from airflow.providers.amazon.aws.operators.base_aws import AwsBaseOperator
 from airflow.providers.amazon.aws.triggers.redshift_cluster import (
@@ -33,7 +32,7 @@ from airflow.providers.amazon.aws.triggers.redshift_cluster import (
 )
 from airflow.providers.amazon.aws.utils import validate_execute_complete_event
 from airflow.providers.amazon.aws.utils.mixins import aws_template_fields
-from airflow.providers.common.compat.sdk import AirflowException
+from airflow.providers.common.compat.sdk import AirflowException, conf
 from airflow.utils.helpers import prune_dict
 
 if TYPE_CHECKING:

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

@@ -19,13 +19,12 @@ from __future__ import annotations
 
 from typing import TYPE_CHECKING, Any
 
-from airflow.configuration import conf
 from airflow.providers.amazon.aws.hooks.redshift_data import RedshiftDataHook
 from airflow.providers.amazon.aws.operators.base_aws import AwsBaseOperator
 from airflow.providers.amazon.aws.triggers.redshift_data import RedshiftDataTrigger
 from airflow.providers.amazon.aws.utils import validate_execute_complete_event
 from airflow.providers.amazon.aws.utils.mixins import aws_template_fields
-from airflow.providers.common.compat.sdk import AirflowException
+from airflow.providers.common.compat.sdk import AirflowException, conf
 
 if TYPE_CHECKING:
     from mypy_boto3_redshift_data.type_defs import (

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

@@ -25,7 +25,6 @@ from typing import TYPE_CHECKING, Any, ClassVar
 
 from botocore.exceptions import ClientError
 
-from airflow.configuration import conf
 from airflow.providers.amazon.aws.hooks.base_aws import AwsBaseHook
 from airflow.providers.amazon.aws.hooks.sagemaker import (
     LogState,
@@ -42,7 +41,7 @@ from airflow.providers.amazon.aws.utils import trim_none_values, validate_execut
 from airflow.providers.amazon.aws.utils.mixins import aws_template_fields
 from airflow.providers.amazon.aws.utils.sagemaker import ApprovalStatus
 from airflow.providers.amazon.aws.utils.tags import format_tags
-from airflow.providers.common.compat.sdk import AirflowException
+from airflow.providers.common.compat.sdk import AirflowException, conf
 from airflow.utils.helpers import prune_dict
 
 if TYPE_CHECKING:

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

@@ -22,7 +22,6 @@ from __future__ import annotations
 from functools import cached_property
 from typing import TYPE_CHECKING
 
-from airflow.configuration import conf
 from airflow.providers.amazon.aws.hooks.sagemaker_unified_studio import (
     SageMakerNotebookHook,
 )
@@ -32,7 +31,7 @@ from airflow.providers.amazon.aws.links.sagemaker_unified_studio import (
 from airflow.providers.amazon.aws.triggers.sagemaker_unified_studio import (
     SageMakerNotebookJobTrigger,
 )
-from airflow.providers.common.compat.sdk import AirflowException, BaseOperator
+from airflow.providers.common.compat.sdk import AirflowException, BaseOperator, conf
 
 if TYPE_CHECKING:
     from airflow.sdk import Context
@@ -64,9 +63,8 @@ class SageMakerNotebookOperator(BaseOperator):
     :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 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_details": {"ecr_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.