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

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

@@ -338,7 +338,7 @@ class EksCreateClusterOperator(BaseOperator):
             fargate_pod_execution_role_arn=self.fargate_pod_execution_role_arn,
             fargate_selectors=self.fargate_selectors,
             create_fargate_profile_kwargs=self.create_fargate_profile_kwargs,
-            subnets=cast(list[str], self.resources_vpc_config.get("subnetIds")),
+            subnets=cast("list[str]", self.resources_vpc_config.get("subnetIds")),
         )
 
     def deferrable_create_cluster_next(self, context: Context, event: dict[str, Any] | None = None) -> None:
@@ -377,7 +377,7 @@ class EksCreateClusterOperator(BaseOperator):
                 fargate_pod_execution_role_arn=self.fargate_pod_execution_role_arn,
                 fargate_selectors=self.fargate_selectors,
                 create_fargate_profile_kwargs=self.create_fargate_profile_kwargs,
-                subnets=cast(list[str], self.resources_vpc_config.get("subnetIds")),
+                subnets=cast("list[str]", self.resources_vpc_config.get("subnetIds")),
             )
             if self.compute == "fargate":
                 self.defer(
@@ -503,7 +503,7 @@ class EksCreateNodegroupOperator(BaseOperator):
             nodegroup_subnets_list: list[str] = []
             if self.nodegroup_subnets != "":
                 try:
-                    nodegroup_subnets_list = cast(list, literal_eval(self.nodegroup_subnets))
+                    nodegroup_subnets_list = cast("list", literal_eval(self.nodegroup_subnets))
                 except ValueError:
                     self.log.warning(
                         "The nodegroup_subnets should be List or string representing "

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

@@ -20,19 +20,19 @@ from __future__ import annotations
 import json
 from collections.abc import Sequence
 from datetime import timedelta
-from functools import cached_property
 from typing import TYPE_CHECKING, Any
 
 from airflow.configuration import conf
 from airflow.exceptions import AirflowException
-from airflow.models import BaseOperator
 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 (
     RdsDbAvailableTrigger,
     RdsDbDeletedTrigger,
     RdsDbStoppedTrigger,
 )
 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.rds import RdsDbType
 from airflow.providers.amazon.aws.utils.tags import format_tags
 from airflow.providers.amazon.aws.utils.waiter_with_logging import wait
@@ -44,9 +44,10 @@ if TYPE_CHECKING:
     from airflow.utils.context import Context
 
 
-class RdsBaseOperator(BaseOperator):
+class RdsBaseOperator(AwsBaseOperator[RdsHook]):
     """Base operator that implements common functions for all operators."""
 
+    aws_hook_class = RdsHook
     ui_color = "#eeaa88"
     ui_fgcolor = "#ffffff"
 
@@ -63,10 +64,6 @@ class RdsBaseOperator(BaseOperator):
 
         self._await_interval = 60  # seconds
 
-    @cached_property
-    def hook(self) -> RdsHook:
-        return RdsHook(aws_conn_id=self.aws_conn_id, region_name=self.region_name)
-
     def execute(self, context: Context) -> str:
         """Different implementations for snapshots, tasks and events."""
         raise NotImplementedError
@@ -92,9 +89,19 @@ class RdsCreateDbSnapshotOperator(RdsBaseOperator):
     :param tags: A dictionary of tags or a list of tags in format `[{"Key": "...", "Value": "..."},]`
         `USER Tagging <https://docs.aws.amazon.com/AmazonRDS/latest/UserGuide/USER_Tagging.html>`__
     :param wait_for_completion:  If True, waits for creation of the DB snapshot to complete. (default: True)
+    :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
     """
 
-    template_fields = ("db_snapshot_identifier", "db_identifier", "tags")
+    template_fields = aws_template_fields("db_snapshot_identifier", "db_identifier", "tags")
 
     def __init__(
         self,
@@ -167,9 +174,14 @@ class RdsCopyDbSnapshotOperator(RdsBaseOperator):
         Only when db_type='instance'
     :param source_region: The ID of the region that contains the snapshot to be copied
     :param wait_for_completion:  If True, waits for snapshot copy to complete. (default: True)
+    :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
     """
 
-    template_fields = (
+    template_fields = aws_template_fields(
         "source_db_snapshot_identifier",
         "target_db_snapshot_identifier",
         "tags",
@@ -260,9 +272,16 @@ class RdsDeleteDbSnapshotOperator(RdsBaseOperator):
 
     :param db_type: Type of the DB - either "instance" or "cluster"
     :param db_snapshot_identifier: The identifier for the DB instance or DB cluster snapshot
+    :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
     """
 
-    template_fields = ("db_snapshot_identifier",)
+    template_fields = aws_template_fields(
+        "db_snapshot_identifier",
+    )
 
     def __init__(
         self,
@@ -319,9 +338,14 @@ class RdsStartExportTaskOperator(RdsBaseOperator):
     :param wait_for_completion:  If True, waits for the DB snapshot export to complete. (default: True)
     :param waiter_interval: The number of seconds to wait before checking the export status. (default: 30)
     :param waiter_max_attempts: The number of attempts to make before failing. (default: 40)
+    :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
     """
 
-    template_fields = (
+    template_fields = aws_template_fields(
         "export_task_identifier",
         "source_arn",
         "s3_bucket_name",
@@ -394,9 +418,16 @@ class RdsCancelExportTaskOperator(RdsBaseOperator):
     :param wait_for_completion:  If True, waits for DB snapshot export to cancel. (default: True)
     :param check_interval: The amount of time in seconds to wait between attempts
     :param max_attempts: The maximum number of attempts to be made
+    :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
     """
 
-    template_fields = ("export_task_identifier",)
+    template_fields = aws_template_fields(
+        "export_task_identifier",
+    )
 
     def __init__(
         self,
@@ -450,9 +481,14 @@ class RdsCreateEventSubscriptionOperator(RdsBaseOperator):
     :param tags: A dictionary of tags or a list of tags in format `[{"Key": "...", "Value": "..."},]`
         `USER Tagging <https://docs.aws.amazon.com/AmazonRDS/latest/UserGuide/USER_Tagging.html>`__
     :param wait_for_completion:  If True, waits for creation of the subscription to complete. (default: True)
+    :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
     """
 
-    template_fields = (
+    template_fields = aws_template_fields(
         "subscription_name",
         "sns_topic_arn",
         "source_type",
@@ -513,9 +549,16 @@ class RdsDeleteEventSubscriptionOperator(RdsBaseOperator):
         :ref:`howto/operator:RdsDeleteEventSubscriptionOperator`
 
     :param subscription_name: The name of the RDS event notification subscription you want to delete
+    :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
     """
 
-    template_fields = ("subscription_name",)
+    template_fields = aws_template_fields(
+        "subscription_name",
+    )
 
     def __init__(
         self,
@@ -560,9 +603,16 @@ class RdsCreateDbInstanceOperator(RdsBaseOperator):
     :param deferrable: If True, the operator will wait asynchronously for the DB instance to be created.
         This implies waiting for completion. This mode requires aiobotocore module to be installed.
         (default: False)
+    :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
     """
 
-    template_fields = ("db_instance_identifier", "db_instance_class", "engine", "rds_kwargs")
+    template_fields = aws_template_fields(
+        "db_instance_identifier", "db_instance_class", "engine", "rds_kwargs"
+    )
 
     def __init__(
         self,
@@ -652,9 +702,14 @@ class RdsDeleteDbInstanceOperator(RdsBaseOperator):
     :param deferrable: If True, the operator will wait asynchronously for the DB instance to be created.
         This implies waiting for completion. This mode requires aiobotocore module to be installed.
         (default: False)
+    :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
     """
 
-    template_fields = ("db_instance_identifier", "rds_kwargs")
+    template_fields = aws_template_fields("db_instance_identifier", "rds_kwargs")
 
     def __init__(
         self,
@@ -735,9 +790,14 @@ class RdsStartDbOperator(RdsBaseOperator):
     :param waiter_max_attempts: The maximum number of attempts to check DB instance state
     :param deferrable: If True, the operator will wait asynchronously for the DB instance to be created.
         This implies waiting for completion. This mode requires aiobotocore module to be installed.
+    :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
     """
 
-    template_fields = ("db_identifier", "db_type")
+    template_fields = aws_template_fields("db_identifier", "db_type")
 
     def __init__(
         self,
@@ -832,9 +892,14 @@ class RdsStopDbOperator(RdsBaseOperator):
     :param waiter_max_attempts: The maximum number of attempts to check DB instance state
     :param deferrable: If True, the operator will wait asynchronously for the DB instance to be created.
         This implies waiting for completion. This mode requires aiobotocore module to be installed.
+    :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
     """
 
-    template_fields = ("db_identifier", "db_snapshot_identifier", "db_type")
+    template_fields = aws_template_fields("db_identifier", "db_snapshot_identifier", "db_type")
 
     def __init__(
         self,

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

@@ -755,11 +755,18 @@ class RedshiftDeleteClusterOperator(BaseOperator):
                     final_cluster_snapshot_identifier=self.final_cluster_snapshot_identifier,
                 )
                 break
-            except self.redshift_hook.get_conn().exceptions.InvalidClusterStateFault:
+            except self.redshift_hook.conn.exceptions.InvalidClusterStateFault:
                 self._attempts -= 1
 
                 if self._attempts:
-                    self.log.error("Unable to delete cluster. %d attempts remaining.", self._attempts)
+                    current_state = self.redshift_hook.conn.describe_clusters(
+                        ClusterIdentifier=self.cluster_identifier
+                    )["Clusters"][0]["ClusterStatus"]
+                    self.log.error(
+                        "Cluster in %s state, unable to delete. %d attempts remaining.",
+                        current_state,
+                        self._attempts,
+                    )
                     time.sleep(self._attempt_interval)
                 else:
                     raise
@@ -785,7 +792,7 @@ class RedshiftDeleteClusterOperator(BaseOperator):
                 )
 
         elif self.wait_for_completion:
-            waiter = self.redshift_hook.get_conn().get_waiter("cluster_deleted")
+            waiter = self.redshift_hook.conn.get_waiter("cluster_deleted")
             waiter.wait(
                 ClusterIdentifier=self.cluster_identifier,
                 WaiterConfig={"Delay": self.poll_interval, "MaxAttempts": self.max_attempts},

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

@@ -170,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
 
@@ -178,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:
@@ -559,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/sensors/bedrock.py

@@ -26,6 +26,8 @@ from airflow.exceptions import AirflowException
 from airflow.providers.amazon.aws.hooks.bedrock import BedrockAgentHook, BedrockHook
 from airflow.providers.amazon.aws.sensors.base_aws import AwsBaseSensor
 from airflow.providers.amazon.aws.triggers.bedrock import (
+    BedrockBatchInferenceCompletedTrigger,
+    BedrockBatchInferenceScheduledTrigger,
     BedrockCustomizeModelCompletedTrigger,
     BedrockIngestionJobTrigger,
     BedrockKnowledgeBaseActiveTrigger,
@@ -34,6 +36,7 @@ from airflow.providers.amazon.aws.triggers.bedrock import (
 from airflow.providers.amazon.aws.utils.mixins import aws_template_fields
 
 if TYPE_CHECKING:
+    from airflow.providers.amazon.aws.triggers.bedrock import BedrockBaseBatchInferenceTrigger
     from airflow.utils.context import Context
 
 
@@ -368,3 +371,110 @@ class BedrockIngestionJobSensor(BedrockBaseSensor[BedrockAgentHook]):
             )
         else:
             super().execute(context=context)
+
+
+class BedrockBatchInferenceSensor(BedrockBaseSensor[BedrockHook]):
+    """
+    Poll the batch inference job status until it reaches a terminal state; fails if creation fails.
+
+    .. seealso::
+        For more information on how to use this sensor, take a look at the guide:
+        :ref:`howto/sensor:BedrockBatchInferenceSensor`
+
+    :param job_arn: The Amazon Resource Name (ARN) of the batch inference job. (templated)
+    :param success_state: A BedrockBatchInferenceSensor.TargetState; defaults to 'SCHEDULED' (templated)
+
+    :param deferrable: If True, the sensor will operate in deferrable more. 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: 5)
+    :param max_retries: Number of times before returning the current state (default: 24)
+    :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
+    """
+
+    class SuccessState:
+        """
+        Target state for the BedrockBatchInferenceSensor.
+
+        Bedrock adds batch inference jobs to a queue, and they may take some time to complete.
+        If you want to wait for the job to complete, use TargetState.COMPLETED, but if you only want
+        to wait until the service confirms that the job is in the queue, use TargetState.SCHEDULED.
+
+        The normal successful progression of states is:
+            Submitted > Validating > Scheduled > InProgress > PartiallyCompleted > Completed
+        """
+
+        SCHEDULED = "scheduled"
+        COMPLETED = "completed"
+
+    INTERMEDIATE_STATES: tuple[str, ...]  # Defined in __init__ based on target state
+    FAILURE_STATES: tuple[str, ...] = ("Failed", "Stopped", "PartiallyCompleted", "Expired")
+    SUCCESS_STATES: tuple[str, ...]  # Defined in __init__ based on target state
+    FAILURE_MESSAGE = "Bedrock batch inference job sensor failed."
+    INVALID_SUCCESS_STATE_MESSAGE = "success_state must be an instance of TargetState."
+
+    aws_hook_class = BedrockHook
+
+    template_fields: Sequence[str] = aws_template_fields("job_arn", "success_state")
+
+    def __init__(
+        self,
+        *,
+        job_arn: str,
+        success_state: SuccessState | str = SuccessState.SCHEDULED,
+        poke_interval: int = 120,
+        max_retries: int = 75,
+        **kwargs,
+    ) -> None:
+        super().__init__(**kwargs)
+        self.poke_interval = poke_interval
+        self.max_retries = max_retries
+        self.job_arn = job_arn
+        self.success_state = success_state
+
+        base_success_states: tuple[str, ...] = ("Completed",)
+        base_intermediate_states: tuple[str, ...] = ("Submitted", "InProgress", "Stopping", "Validating")
+        scheduled_state = ("Scheduled",)
+        self.trigger_class: type[BedrockBaseBatchInferenceTrigger]
+
+        if self.success_state == BedrockBatchInferenceSensor.SuccessState.COMPLETED:
+            intermediate_states = base_intermediate_states + scheduled_state
+            success_states = base_success_states
+            self.trigger_class = BedrockBatchInferenceCompletedTrigger
+        elif self.success_state == BedrockBatchInferenceSensor.SuccessState.SCHEDULED:
+            intermediate_states = base_intermediate_states
+            success_states = base_success_states + scheduled_state
+            self.trigger_class = BedrockBatchInferenceScheduledTrigger
+        else:
+            raise ValueError(
+                "Success states for BedrockBatchInferenceSensor must be set using a BedrockBatchInferenceSensor.SuccessState"
+            )
+
+        BedrockBatchInferenceSensor.INTERMEDIATE_STATES = intermediate_states or base_intermediate_states
+        BedrockBatchInferenceSensor.SUCCESS_STATES = success_states or base_success_states
+
+    def get_state(self) -> str:
+        return self.hook.conn.get_model_invocation_job(jobIdentifier=self.job_arn)["status"]
+
+    def execute(self, context: Context) -> Any:
+        if self.deferrable:
+            self.defer(
+                trigger=self.trigger_class(
+                    job_arn=self.job_arn,
+                    waiter_delay=int(self.poke_interval),
+                    waiter_max_attempts=self.max_retries,
+                    aws_conn_id=self.aws_conn_id,
+                ),
+                method_name="poke",
+            )
+        else:
+            super().execute(context=context)

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

@@ -150,7 +150,8 @@ class MwaaDagRunSensor(AwsBaseSensor[MwaaHook]):
                     external_dag_run_id=self.external_dag_run_id,
                     success_states=self.success_states,
                     failure_states=self.failure_states,
-                    waiter_delay=self.poke_interval,
+                    # 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,
                 ),

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

@@ -17,36 +17,30 @@
 from __future__ import annotations
 
 from collections.abc import Sequence
-from functools import cached_property
 from typing import TYPE_CHECKING
 
 from airflow.exceptions import AirflowException, AirflowNotFoundException
 from airflow.providers.amazon.aws.hooks.rds import RdsHook
+from airflow.providers.amazon.aws.sensors.base_aws import AwsBaseSensor
+from airflow.providers.amazon.aws.utils.mixins import aws_template_fields
 from airflow.providers.amazon.aws.utils.rds import RdsDbType
-from airflow.sensors.base import BaseSensorOperator
 
 if TYPE_CHECKING:
     from airflow.utils.context import Context
 
 
-class RdsBaseSensor(BaseSensorOperator):
+class RdsBaseSensor(AwsBaseSensor[RdsHook]):
     """Base operator that implements common functions for all sensors."""
 
+    aws_hook_class = RdsHook
     ui_color = "#ddbb77"
     ui_fgcolor = "#ffffff"
 
-    def __init__(
-        self, *args, aws_conn_id: str | None = "aws_conn_id", hook_params: dict | None = None, **kwargs
-    ):
+    def __init__(self, *args, hook_params: dict | None = None, **kwargs):
         self.hook_params = hook_params or {}
-        self.aws_conn_id = aws_conn_id
         self.target_statuses: list[str] = []
         super().__init__(*args, **kwargs)
 
-    @cached_property
-    def hook(self):
-        return RdsHook(aws_conn_id=self.aws_conn_id, **self.hook_params)
-
 
 class RdsSnapshotExistenceSensor(RdsBaseSensor):
     """
@@ -59,9 +53,19 @@ class RdsSnapshotExistenceSensor(RdsBaseSensor):
     :param db_type: Type of the DB - either "instance" or "cluster"
     :param db_snapshot_identifier: The identifier for the DB snapshot
     :param target_statuses: Target status of snapshot
+    :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
     """
 
-    template_fields: Sequence[str] = (
+    template_fields: Sequence[str] = aws_template_fields(
         "db_snapshot_identifier",
         "target_statuses",
     )
@@ -72,10 +76,9 @@ class RdsSnapshotExistenceSensor(RdsBaseSensor):
         db_type: str,
         db_snapshot_identifier: str,
         target_statuses: list[str] | None = None,
-        aws_conn_id: str | None = "aws_conn_id",
         **kwargs,
     ):
-        super().__init__(aws_conn_id=aws_conn_id, **kwargs)
+        super().__init__(**kwargs)
         self.db_type = RdsDbType(db_type)
         self.db_snapshot_identifier = db_snapshot_identifier
         self.target_statuses = target_statuses or ["available"]
@@ -107,7 +110,9 @@ class RdsExportTaskExistenceSensor(RdsBaseSensor):
     :param error_statuses: Target error status of export task to fail the sensor
     """
 
-    template_fields: Sequence[str] = ("export_task_identifier", "target_statuses", "error_statuses")
+    template_fields: Sequence[str] = aws_template_fields(
+        "export_task_identifier", "target_statuses", "error_statuses"
+    )
 
     def __init__(
         self,
@@ -115,10 +120,9 @@ class RdsExportTaskExistenceSensor(RdsBaseSensor):
         export_task_identifier: str,
         target_statuses: list[str] | None = None,
         error_statuses: list[str] | None = None,
-        aws_conn_id: str | None = "aws_default",
         **kwargs,
     ):
-        super().__init__(aws_conn_id=aws_conn_id, **kwargs)
+        super().__init__(**kwargs)
 
         self.export_task_identifier = export_task_identifier
         self.target_statuses = target_statuses or [
@@ -159,7 +163,7 @@ class RdsDbSensor(RdsBaseSensor):
     :param target_statuses: Target status of DB
     """
 
-    template_fields: Sequence[str] = (
+    template_fields: Sequence[str] = aws_template_fields(
         "db_identifier",
         "db_type",
         "target_statuses",
@@ -171,10 +175,9 @@ class RdsDbSensor(RdsBaseSensor):
         db_identifier: str,
         db_type: RdsDbType | str = RdsDbType.INSTANCE,
         target_statuses: list[str] | None = None,
-        aws_conn_id: str | None = "aws_default",
         **kwargs,
     ):
-        super().__init__(aws_conn_id=aws_conn_id, **kwargs)
+        super().__init__(**kwargs)
         self.db_identifier = db_identifier
         self.target_statuses = target_statuses or ["available"]
         self.db_type = db_type

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

@@ -192,7 +192,7 @@ class S3KeySensor(AwsBaseSensor[S3Hook]):
         self.defer(
             timeout=timedelta(seconds=self.timeout),
             trigger=S3KeyTrigger(
-                bucket_name=cast(str, self.bucket_name),
+                bucket_name=cast("str", self.bucket_name),
                 bucket_key=self.bucket_key,
                 wildcard_match=self.wildcard_match,
                 aws_conn_id=self.aws_conn_id,

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

@@ -81,5 +81,6 @@ class StepFunctionExecutionSensor(AwsBaseSensor[StepFunctionHook]):
             return False
 
         self.log.info("Doing xcom_push of output")
-        self.xcom_push(context, "output", output)
+
+        context["ti"].xcom_push(key="output", value=output)
         return True

airflow/providers/amazon/aws/transfers/mongo_to_s3.py

@@ -103,7 +103,7 @@ class MongoToS3Operator(BaseOperator):
         if self.is_pipeline:
             results: CommandCursor[Any] | Cursor = MongoHook(self.mongo_conn_id).aggregate(
                 mongo_collection=self.mongo_collection,
-                aggregate_query=cast(list, self.mongo_query),
+                aggregate_query=cast("list", self.mongo_query),
                 mongo_db=self.mongo_db,
                 allowDiskUse=self.allow_disk_use,
             )
@@ -111,7 +111,7 @@ class MongoToS3Operator(BaseOperator):
         else:
             results = MongoHook(self.mongo_conn_id).find(
                 mongo_collection=self.mongo_collection,
-                query=cast(dict, self.mongo_query),
+                query=cast("dict", self.mongo_query),
                 projection=self.mongo_projection,
                 mongo_db=self.mongo_db,
                 find_one=False,

airflow/providers/amazon/aws/transfers/sql_to_s3.py

@@ -223,7 +223,7 @@ class SqlToS3Operator(BaseOperator):
             return
         for group_label in (grouped_df := df.groupby(**self.groupby_kwargs)).groups:
             yield (
-                cast(str, group_label),
+                cast("str", group_label),
                 grouped_df.get_group(group_label)
                 .drop(random_column_name, axis=1, errors="ignore")
                 .reset_index(drop=True),