apache-airflow-providers-amazon 8.16.0__py3-none-any.whl → 8.17.0__py3-none-any.whl

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

@@ -18,16 +18,15 @@ from __future__ import annotations
 
 from typing import TYPE_CHECKING, Sequence
 
-from airflow.models import BaseOperator
 from airflow.providers.amazon.aws.hooks.quicksight import QuickSightHook
+from airflow.providers.amazon.aws.operators.base_aws import AwsBaseOperator
+from airflow.providers.amazon.aws.utils.mixins import aws_template_fields
 
 if TYPE_CHECKING:
     from airflow.utils.context import Context
 
-DEFAULT_CONN_ID = "aws_default"
 
-
-class QuickSightCreateIngestionOperator(BaseOperator):
+class QuickSightCreateIngestionOperator(AwsBaseOperator[QuickSightHook]):
     """
     Creates and starts a new SPICE ingestion for a dataset;  also helps to Refresh existing SPICE datasets.
 
@@ -43,23 +42,25 @@ class QuickSightCreateIngestionOperator(BaseOperator):
         that the operation waits to check the status of the Amazon QuickSight Ingestion.
     :param check_interval: if wait is set to be true, this is the time interval
         in seconds which the operator will check the status of the Amazon QuickSight Ingestion
-    :param aws_conn_id: The Airflow connection used for AWS credentials. (templated)
-         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 the default boto3 configuration would be used (and must be
-         maintained on each worker node).
-    :param region: Which AWS region the connection should use. (templated)
-         If this is None or empty then the default boto3 behaviour is used.
+    :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] = (
+    aws_hook_class = QuickSightHook
+    template_fields: Sequence[str] = aws_template_fields(
         "data_set_id",
         "ingestion_id",
         "ingestion_type",
         "wait_for_completion",
         "check_interval",
-        "aws_conn_id",
-        "region",
     )
     ui_color = "#ffd700"
 
@@ -70,26 +71,18 @@ class QuickSightCreateIngestionOperator(BaseOperator):
         ingestion_type: str = "FULL_REFRESH",
         wait_for_completion: bool = True,
         check_interval: int = 30,
-        aws_conn_id: str = DEFAULT_CONN_ID,
-        region: str | None = None,
         **kwargs,
     ):
+        super().__init__(**kwargs)
         self.data_set_id = data_set_id
         self.ingestion_id = ingestion_id
         self.ingestion_type = ingestion_type
         self.wait_for_completion = wait_for_completion
         self.check_interval = check_interval
-        self.aws_conn_id = aws_conn_id
-        self.region = region
-        super().__init__(**kwargs)
 
     def execute(self, context: Context):
-        hook = QuickSightHook(
-            aws_conn_id=self.aws_conn_id,
-            region_name=self.region,
-        )
         self.log.info("Running the Amazon QuickSight SPICE Ingestion on Dataset ID: %s", self.data_set_id)
-        return hook.create_ingestion(
+        return self.hook.create_ingestion(
             data_set_id=self.data_set_id,
             ingestion_id=self.ingestion_id,
             ingestion_type=self.ingestion_type,

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

@@ -17,11 +17,14 @@
 # under the License.
 from __future__ import annotations
 
-from functools import cached_property
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Any
 
-from airflow.models import BaseOperator
+from airflow.configuration import conf
+from airflow.exceptions import AirflowException
 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.mixins import aws_template_fields
 
 if TYPE_CHECKING:
     from mypy_boto3_redshift_data.type_defs import GetStatementResultResponseTypeDef
@@ -29,7 +32,7 @@ if TYPE_CHECKING:
     from airflow.utils.context import Context
 
 
-class RedshiftDataOperator(BaseOperator):
+class RedshiftDataOperator(AwsBaseOperator[RedshiftDataHook]):
     """
     Executes SQL Statements against an Amazon Redshift cluster using Redshift Data.
 
@@ -49,22 +52,29 @@ class RedshiftDataOperator(BaseOperator):
     :param poll_interval: how often in seconds to check the query status
     :param return_sql_result: if True will return the result of an SQL statement,
         if False (default) will return statement ID
-    :param aws_conn_id: aws connection to use
-    :param region: aws region to use
     :param workgroup_name: name of the Redshift Serverless workgroup. Mutually exclusive with
         `cluster_identifier`. Specify this parameter to query Redshift Serverless. More info
         https://docs.aws.amazon.com/redshift/latest/mgmt/working-with-serverless.html
+    :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 = (
+    aws_hook_class = RedshiftDataHook
+    template_fields = aws_template_fields(
         "cluster_identifier",
         "database",
         "sql",
         "db_user",
         "parameters",
         "statement_name",
-        "aws_conn_id",
-        "region",
         "workgroup_name",
     )
     template_ext = (".sql",)
@@ -84,9 +94,8 @@ class RedshiftDataOperator(BaseOperator):
         wait_for_completion: bool = True,
         poll_interval: int = 10,
         return_sql_result: bool = False,
-        aws_conn_id: str = "aws_default",
-        region: str | None = None,
         workgroup_name: str | None = None,
+        deferrable: bool = conf.getboolean("operators", "default_deferrable", fallback=False),
         **kwargs,
     ) -> None:
         super().__init__(**kwargs)
@@ -108,19 +117,18 @@ class RedshiftDataOperator(BaseOperator):
                 poll_interval,
             )
         self.return_sql_result = return_sql_result
-        self.aws_conn_id = aws_conn_id
-        self.region = region
         self.statement_id: str | None = None
-
-    @cached_property
-    def hook(self) -> RedshiftDataHook:
-        """Create and return an RedshiftDataHook."""
-        return RedshiftDataHook(aws_conn_id=self.aws_conn_id, region_name=self.region)
+        self.deferrable = deferrable
 
     def execute(self, context: Context) -> GetStatementResultResponseTypeDef | str:
         """Execute a statement against Amazon Redshift."""
         self.log.info("Executing statement: %s", self.sql)
 
+        # Set wait_for_completion to False so that it waits for the status in the deferred task.
+        wait_for_completion = self.wait_for_completion
+        if self.deferrable and self.wait_for_completion:
+            self.wait_for_completion = False
+
         self.statement_id = self.hook.execute_query(
             database=self.database,
             sql=self.sql,
@@ -131,10 +139,27 @@ class RedshiftDataOperator(BaseOperator):
             secret_arn=self.secret_arn,
             statement_name=self.statement_name,
             with_event=self.with_event,
-            wait_for_completion=self.wait_for_completion,
+            wait_for_completion=wait_for_completion,
             poll_interval=self.poll_interval,
         )
 
+        if self.deferrable:
+            is_finished = self.hook.check_query_is_finished(self.statement_id)
+            if not is_finished:
+                self.defer(
+                    timeout=self.execution_timeout,
+                    trigger=RedshiftDataTrigger(
+                        statement_id=self.statement_id,
+                        task_id=self.task_id,
+                        poll_interval=self.poll_interval,
+                        aws_conn_id=self.aws_conn_id,
+                        region_name=self.region_name,
+                        verify=self.verify,
+                        botocore_config=self.botocore_config,
+                    ),
+                    method_name="execute_complete",
+                )
+
         if self.return_sql_result:
             result = self.hook.conn.get_statement_result(Id=self.statement_id)
             self.log.debug("Statement result: %s", result)
@@ -142,6 +167,30 @@ class RedshiftDataOperator(BaseOperator):
         else:
             return self.statement_id
 
+    def execute_complete(
+        self, context: Context, event: dict[str, Any] | None = None
+    ) -> GetStatementResultResponseTypeDef | str:
+        if event is None:
+            err_msg = "Trigger error: event is None"
+            self.log.info(err_msg)
+            raise AirflowException(err_msg)
+
+        if event["status"] == "error":
+            msg = f"context: {context}, error message: {event['message']}"
+            raise AirflowException(msg)
+
+        statement_id = event["statement_id"]
+        if not statement_id:
+            raise AirflowException("statement_id should not be empty.")
+
+        self.log.info("%s completed successfully.", self.task_id)
+        if self.return_sql_result:
+            result = self.hook.conn.get_statement_result(Id=statement_id)
+            self.log.debug("Statement result: %s", result)
+            return result
+
+        return statement_id
+
     def on_kill(self) -> None:
         """Cancel the submitted redshift query."""
         if self.statement_id:

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

@@ -744,7 +744,6 @@ class S3ListOperator(BaseOperator):
     :param delimiter: the delimiter marks key hierarchy. (templated)
     :param aws_conn_id: The connection ID to use when connecting to S3 storage.
     :param verify: Whether or not to verify SSL certificates for S3 connection.
-    :param apply_wildcard: whether to treat '*' as a wildcard or a plain symbol in the prefix.
         By default SSL certificates are verified.
         You can provide the following values:
 
@@ -754,6 +753,7 @@ class S3ListOperator(BaseOperator):
         - ``path/to/cert/bundle.pem``: A filename of the CA cert bundle to uses.
                  You can specify this argument if you want to use a different
                  CA cert bundle than the one used by botocore.
+    :param apply_wildcard: whether to treat '*' as a wildcard or a plain symbol in the prefix.
 
 
     **Example**:

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

@@ -283,8 +283,20 @@ class SageMakerProcessingOperator(SageMakerBaseOperator):
             raise AirflowException(f"Sagemaker Processing Job creation failed: {response}")
 
         if self.deferrable and self.wait_for_completion:
+            response = self.hook.describe_processing_job(self.config["ProcessingJobName"])
+            status = response["ProcessingJobStatus"]
+            if status in self.hook.failed_states:
+                raise AirflowException(f"SageMaker job failed because {response['FailureReason']}")
+            elif status == "Completed":
+                self.log.info("%s completed successfully.", self.task_id)
+                return {"Processing": serialize(response)}
+
+            timeout = self.execution_timeout
+            if self.max_ingestion_time:
+                timeout = datetime.timedelta(seconds=self.max_ingestion_time)
+
             self.defer(
-                timeout=self.execution_timeout,
+                timeout=timeout,
                 trigger=SageMakerTrigger(
                     job_name=self.config["ProcessingJobName"],
                     job_type="Processing",
@@ -304,6 +316,7 @@ class SageMakerProcessingOperator(SageMakerBaseOperator):
         else:
             self.log.info(event["message"])
         self.serialized_job = serialize(self.hook.describe_processing_job(self.config["ProcessingJobName"]))
+        self.log.info("%s completed successfully.", self.task_id)
         return {"Processing": self.serialized_job}
 
     def get_openlineage_facets_on_complete(self, task_instance) -> OperatorLineage:
@@ -700,8 +713,24 @@ class SageMakerTransformOperator(SageMakerBaseOperator):
             raise AirflowException(f"Sagemaker transform Job creation failed: {response}")
 
         if self.deferrable and self.wait_for_completion:
+            response = self.hook.describe_transform_job(transform_config["TransformJobName"])
+            status = response["TransformJobStatus"]
+            if status in self.hook.failed_states:
+                raise AirflowException(f"SageMaker job failed because {response['FailureReason']}")
+
+            if status == "Completed":
+                self.log.info("%s completed successfully.", self.task_id)
+                return {
+                    "Model": serialize(self.hook.describe_model(transform_config["ModelName"])),
+                    "Transform": serialize(response),
+                }
+
+            timeout = self.execution_timeout
+            if self.max_ingestion_time:
+                timeout = datetime.timedelta(seconds=self.max_ingestion_time)
+
             self.defer(
-                timeout=self.execution_timeout,
+                timeout=timeout,
                 trigger=SageMakerTrigger(
                     job_name=transform_config["TransformJobName"],
                     job_type="Transform",
@@ -712,17 +741,18 @@ class SageMakerTransformOperator(SageMakerBaseOperator):
                 method_name="execute_complete",
             )
 
-        self.serialized_model = serialize(self.hook.describe_model(transform_config["ModelName"]))
-        self.serialized_transform = serialize(
-            self.hook.describe_transform_job(transform_config["TransformJobName"])
-        )
-        return {"Model": self.serialized_model, "Transform": self.serialized_transform}
+        return self.serialize_result()
 
-    def execute_complete(self, context, event=None):
-        if event["status"] != "success":
-            raise AirflowException(f"Error while running job: {event}")
-        else:
-            self.log.info(event["message"])
+    def execute_complete(self, context: Context, event: dict[str, Any] | None = None) -> dict[str, dict]:
+        if event is None:
+            err_msg = "Trigger error: event is None"
+            self.log.error(err_msg)
+            raise AirflowException(err_msg)
+
+        self.log.info(event["message"])
+        return self.serialize_result()
+
+    def serialize_result(self) -> dict[str, dict]:
         transform_config = self.config.get("Transform", self.config)
         self.serialized_model = serialize(self.hook.describe_model(transform_config["ModelName"]))
         self.serialized_transform = serialize(

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

@@ -18,18 +18,19 @@
 """This module contains sensors for AWS CloudFormation."""
 from __future__ import annotations
 
-from functools import cached_property
 from typing import TYPE_CHECKING, Sequence
 
+from airflow.providers.amazon.aws.sensors.base_aws import AwsBaseSensor
+from airflow.providers.amazon.aws.utils.mixins import aws_template_fields
+
 if TYPE_CHECKING:
     from airflow.utils.context import Context
 
 from airflow.exceptions import AirflowSkipException
 from airflow.providers.amazon.aws.hooks.cloud_formation import CloudFormationHook
-from airflow.sensors.base import BaseSensorOperator
 
 
-class CloudFormationCreateStackSensor(BaseSensorOperator):
+class CloudFormationCreateStackSensor(AwsBaseSensor[CloudFormationHook]):
     """
     Waits for a stack to be created successfully on AWS CloudFormation.
 
@@ -38,19 +39,25 @@ class CloudFormationCreateStackSensor(BaseSensorOperator):
         :ref:`howto/sensor:CloudFormationCreateStackSensor`
 
     :param stack_name: The name of the stack to wait for (templated)
-    :param aws_conn_id: ID of the Airflow connection where credentials and extra configuration are
-        stored
-    :param poke_interval: Time in seconds that the job should wait between each try
+    :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] = ("stack_name",)
+    aws_hook_class = CloudFormationHook
+    template_fields: Sequence[str] = aws_template_fields("stack_name")
     ui_color = "#C5CAE9"
 
-    def __init__(self, *, stack_name, aws_conn_id="aws_default", region_name=None, **kwargs):
+    def __init__(self, *, stack_name, **kwargs):
         super().__init__(**kwargs)
         self.stack_name = stack_name
-        self.aws_conn_id = aws_conn_id
-        self.region_name = region_name
 
     def poke(self, context: Context):
         stack_status = self.hook.get_stack_status(self.stack_name)
@@ -65,13 +72,8 @@ class CloudFormationCreateStackSensor(BaseSensorOperator):
             raise AirflowSkipException(message)
         raise ValueError(message)
 
-    @cached_property
-    def hook(self) -> CloudFormationHook:
-        """Create and return a CloudFormationHook."""
-        return CloudFormationHook(aws_conn_id=self.aws_conn_id, region_name=self.region_name)
 
-
-class CloudFormationDeleteStackSensor(BaseSensorOperator):
+class CloudFormationDeleteStackSensor(AwsBaseSensor[CloudFormationHook]):
     """
     Waits for a stack to be deleted successfully on AWS CloudFormation.
 
@@ -80,12 +82,20 @@ class CloudFormationDeleteStackSensor(BaseSensorOperator):
         :ref:`howto/sensor:CloudFormationDeleteStackSensor`
 
     :param stack_name: The name of the stack to wait for (templated)
-    :param aws_conn_id: ID of the Airflow connection where credentials and extra configuration are
-        stored
-    :param poke_interval: Time in seconds that the job should wait between each try
+    :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] = ("stack_name",)
+    aws_hook_class = CloudFormationHook
+    template_fields: Sequence[str] = aws_template_fields("stack_name")
     ui_color = "#C5CAE9"
 
     def __init__(
@@ -113,8 +123,3 @@ class CloudFormationDeleteStackSensor(BaseSensorOperator):
         if self.soft_fail:
             raise AirflowSkipException(message)
         raise ValueError(message)
-
-    @cached_property
-    def hook(self) -> CloudFormationHook:
-        """Create and return a CloudFormationHook."""
-        return CloudFormationHook(aws_conn_id=self.aws_conn_id, region_name=self.region_name)

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

@@ -17,47 +17,53 @@
 # under the License.
 from __future__ import annotations
 
-from functools import cached_property
 from typing import TYPE_CHECKING, Iterable, Sequence
 
 from deprecated import deprecated
 
 from airflow.exceptions import AirflowException, AirflowProviderDeprecationWarning, AirflowSkipException
 from airflow.providers.amazon.aws.hooks.dms import DmsHook
-from airflow.sensors.base import BaseSensorOperator
+from airflow.providers.amazon.aws.sensors.base_aws import AwsBaseSensor
+from airflow.providers.amazon.aws.utils.mixins import aws_template_fields
 
 if TYPE_CHECKING:
     from airflow.utils.context import Context
 
 
-class DmsTaskBaseSensor(BaseSensorOperator):
+class DmsTaskBaseSensor(AwsBaseSensor[DmsHook]):
     """
     Contains general sensor behavior for DMS task.
 
     Subclasses should set ``target_statuses`` and ``termination_statuses`` fields.
 
     :param replication_task_arn: AWS DMS replication task ARN
-    :param aws_conn_id: aws connection to uses
     :param target_statuses: the target statuses, sensor waits until
         the task reaches any of these states
     :param termination_statuses: the termination statuses, sensor fails when
         the task reaches any of these states
+    :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] = ("replication_task_arn",)
-    template_ext: Sequence[str] = ()
+    aws_hook_class = DmsHook
+    template_fields: Sequence[str] = aws_template_fields("replication_task_arn")
 
     def __init__(
         self,
         replication_task_arn: str,
-        aws_conn_id="aws_default",
         target_statuses: Iterable[str] | None = None,
         termination_statuses: Iterable[str] | None = None,
-        *args,
         **kwargs,
     ):
-        super().__init__(*args, **kwargs)
-        self.aws_conn_id = aws_conn_id
+        super().__init__(**kwargs)
         self.replication_task_arn = replication_task_arn
         self.target_statuses: Iterable[str] = target_statuses or []
         self.termination_statuses: Iterable[str] = termination_statuses or []
@@ -67,14 +73,8 @@ class DmsTaskBaseSensor(BaseSensorOperator):
         """Get DmsHook."""
         return self.hook
 
-    @cached_property
-    def hook(self) -> DmsHook:
-        return DmsHook(self.aws_conn_id)
-
     def poke(self, context: Context):
-        status: str | None = self.hook.get_task_status(self.replication_task_arn)
-
-        if not status:
+        if not (status := self.hook.get_task_status(self.replication_task_arn)):
             # TODO: remove this if check when min_airflow_version is set to higher than 2.7.1
             message = f"Failed to read task status, task with ARN {self.replication_task_arn} not found"
             if self.soft_fail:
@@ -105,15 +105,21 @@ class DmsTaskCompletedSensor(DmsTaskBaseSensor):
         :ref:`howto/sensor:DmsTaskCompletedSensor`
 
     :param replication_task_arn: AWS DMS replication task ARN
+    :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] = ("replication_task_arn",)
-    template_ext: Sequence[str] = ()
-
-    def __init__(self, *args, **kwargs):
-        super().__init__(*args, **kwargs)
-        self.target_statuses = ["stopped"]
-        self.termination_statuses = [
+    def __init__(self, **kwargs):
+        kwargs["target_statuses"] = ["stopped"]
+        kwargs["termination_statuses"] = [
             "creating",
             "deleting",
             "failed",
@@ -123,3 +129,4 @@ class DmsTaskCompletedSensor(DmsTaskBaseSensor):
             "ready",
             "testing",
         ]
+        super().__init__(**kwargs)

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

@@ -16,17 +16,17 @@
 # under the License.
 from __future__ import annotations
 
-from functools import cached_property
 from typing import TYPE_CHECKING, Any, Iterable, Sequence
 
 from airflow.providers.amazon.aws.hooks.dynamodb import DynamoDBHook
-from airflow.sensors.base import BaseSensorOperator
+from airflow.providers.amazon.aws.sensors.base_aws import AwsBaseSensor
+from airflow.providers.amazon.aws.utils.mixins import aws_template_fields
 
 if TYPE_CHECKING:
     from airflow.utils.context import Context
 
 
-class DynamoDBValueSensor(BaseSensorOperator):
+class DynamoDBValueSensor(AwsBaseSensor[DynamoDBHook]):
     """
     Waits for an attribute value to be present for an item in a DynamoDB table.
 
@@ -41,11 +41,20 @@ class DynamoDBValueSensor(BaseSensorOperator):
     :param attribute_value: DynamoDB attribute value
     :param sort_key_name: (optional) DynamoDB sort key name
     :param sort_key_value: (optional) DynamoDB sort key value
-    :param aws_conn_id: aws connection to use
-    :param region_name: aws region to use
+    :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] = (
+    aws_hook_class = DynamoDBHook
+    template_fields: Sequence[str] = aws_template_fields(
         "table_name",
         "partition_key_name",
         "partition_key_value",
@@ -64,8 +73,6 @@ class DynamoDBValueSensor(BaseSensorOperator):
         attribute_value: str | Iterable[str],
         sort_key_name: str | None = None,
         sort_key_value: str | None = None,
-        aws_conn_id: str | None = DynamoDBHook.default_conn_name,
-        region_name: str | None = None,
         **kwargs: Any,
     ):
         super().__init__(**kwargs)
@@ -76,8 +83,6 @@ class DynamoDBValueSensor(BaseSensorOperator):
         self.attribute_value = attribute_value
         self.sort_key_name = sort_key_name
         self.sort_key_value = sort_key_value
-        self.aws_conn_id = aws_conn_id
-        self.region_name = region_name
 
     def poke(self, context: Context) -> bool:
         """Test DynamoDB item for matching attribute value."""
@@ -108,8 +113,3 @@ class DynamoDBValueSensor(BaseSensorOperator):
             )
         except KeyError:
             return False
-
-    @cached_property
-    def hook(self) -> DynamoDBHook:
-        """Create and return a DynamoDBHook."""
-        return DynamoDBHook(self.aws_conn_id, region_name=self.region_name)