apache-airflow-providers-amazon 9.8.0rc1__py3-none-any.whl → 9.9.0rc1__py3-none-any.whl

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

@@ -193,6 +193,7 @@ class S3Hook(AwsBaseHook):
     ) -> None:
         kwargs["client_type"] = "s3"
         kwargs["aws_conn_id"] = aws_conn_id
+        self._requester_pays = kwargs.pop("requester_pays", False)
 
         if transfer_config_args and not isinstance(transfer_config_args, dict):
             raise TypeError(f"transfer_config_args expected dict, got {type(transfer_config_args).__name__}.")
@@ -409,12 +410,15 @@ class S3Hook(AwsBaseHook):
         }
 
         paginator = self.get_conn().get_paginator("list_objects_v2")
-        response = paginator.paginate(
-            Bucket=bucket_name,
-            Prefix=prefix,
-            Delimiter=delimiter,
-            PaginationConfig=config,
-        )
+        params = {
+            "Bucket": bucket_name,
+            "Prefix": prefix,
+            "Delimiter": delimiter,
+            "PaginationConfig": config,
+        }
+        if self._requester_pays:
+            params["RequestPayer"] = "requester"
+        response = paginator.paginate(**params)
 
         prefixes: list[str] = []
         for page in response:
@@ -437,7 +441,13 @@ class S3Hook(AwsBaseHook):
         """
         head_object_val: dict[str, Any] | None = None
         try:
-            head_object_val = await client.head_object(Bucket=bucket_name, Key=key)
+            params = {
+                "Bucket": bucket_name,
+                "Key": key,
+            }
+            if self._requester_pays:
+                params["RequestPayer"] = "requester"
+            head_object_val = await client.head_object(**params)
             return head_object_val
         except ClientError as e:
             if e.response["ResponseMetadata"]["HTTPStatusCode"] == 404:
@@ -472,12 +482,15 @@ class S3Hook(AwsBaseHook):
         }
 
         paginator = client.get_paginator("list_objects_v2")
-        response = paginator.paginate(
-            Bucket=bucket_name,
-            Prefix=prefix,
-            Delimiter=delimiter,
-            PaginationConfig=config,
-        )
+        params = {
+            "Bucket": bucket_name,
+            "Prefix": prefix,
+            "Delimiter": delimiter,
+            "PaginationConfig": config,
+        }
+        if self._requester_pays:
+            params["RequestPayer"] = "requester"
+        response = paginator.paginate(**params)
 
         prefixes = []
         async for page in response:
@@ -501,7 +514,14 @@ class S3Hook(AwsBaseHook):
         prefix = re.split(r"[\[\*\?]", key, 1)[0] if key else ""
         delimiter = ""
         paginator = client.get_paginator("list_objects_v2")
-        response = paginator.paginate(Bucket=bucket_name, Prefix=prefix, Delimiter=delimiter)
+        params = {
+            "Bucket": bucket_name,
+            "Prefix": prefix,
+            "Delimiter": delimiter,
+        }
+        if self._requester_pays:
+            params["RequestPayer"] = "requester"
+        response = paginator.paginate(**params)
         async for page in response:
             if "Contents" in page:
                 for row in page["Contents"]:
@@ -622,14 +642,21 @@ class S3Hook(AwsBaseHook):
                 prefix = re.split(r"[\[*?]", key, 1)[0]
 
             paginator = client.get_paginator("list_objects_v2")
-            response = paginator.paginate(Bucket=bucket, Prefix=prefix, Delimiter=delimiter)
+            params = {
+                "Bucket": bucket,
+                "Prefix": prefix,
+                "Delimiter": delimiter,
+            }
+            if self._requester_pays:
+                params["RequestPayer"] = "requester"
+            response = paginator.paginate(**params)
             async for page in response:
                 if "Contents" in page:
                     keys.extend(k for k in page["Contents"] if isinstance(k.get("Size"), (int, float)))
         return keys
 
-    @staticmethod
     async def _list_keys_async(
+        self,
         client: AioBaseClient,
         bucket_name: str | None = None,
         prefix: str | None = None,
@@ -655,12 +682,15 @@ class S3Hook(AwsBaseHook):
         }
 
         paginator = client.get_paginator("list_objects_v2")
-        response = paginator.paginate(
-            Bucket=bucket_name,
-            Prefix=prefix,
-            Delimiter=delimiter,
-            PaginationConfig=config,
-        )
+        params = {
+            "Bucket": bucket_name,
+            "Prefix": prefix,
+            "Delimiter": delimiter,
+            "PaginationConfig": config,
+        }
+        if self._requester_pays:
+            params["RequestPayer"] = "requester"
+        response = paginator.paginate(**params)
 
         keys = []
         async for page in response:
@@ -863,13 +893,16 @@ class S3Hook(AwsBaseHook):
         }
 
         paginator = self.get_conn().get_paginator("list_objects_v2")
-        response = paginator.paginate(
-            Bucket=bucket_name,
-            Prefix=_prefix,
-            Delimiter=delimiter,
-            PaginationConfig=config,
-            StartAfter=start_after_key,
-        )
+        params = {
+            "Bucket": bucket_name,
+            "Prefix": _prefix,
+            "Delimiter": delimiter,
+            "PaginationConfig": config,
+            "StartAfter": start_after_key,
+        }
+        if self._requester_pays:
+            params["RequestPayer"] = "requester"
+        response = paginator.paginate(**params)
 
         keys: list[str] = []
         for page in response:
@@ -909,7 +942,14 @@ class S3Hook(AwsBaseHook):
         }
 
         paginator = self.get_conn().get_paginator("list_objects_v2")
-        response = paginator.paginate(Bucket=bucket_name, Prefix=prefix, PaginationConfig=config)
+        params = {
+            "Bucket": bucket_name,
+            "Prefix": prefix,
+            "PaginationConfig": config,
+        }
+        if self._requester_pays:
+            params["RequestPayer"] = "requester"
+        response = paginator.paginate(**params)
 
         files = []
         for page in response:
@@ -931,7 +971,13 @@ class S3Hook(AwsBaseHook):
         :return: metadata of an object
         """
         try:
-            return self.get_conn().head_object(Bucket=bucket_name, Key=key)
+            params = {
+                "Bucket": bucket_name,
+                "Key": key,
+            }
+            if self._requester_pays:
+                params["RequestPayer"] = "requester"
+            return self.get_conn().head_object(**params)
         except ClientError as e:
             if e.response["ResponseMetadata"]["HTTPStatusCode"] == 404:
                 return None
@@ -975,8 +1021,11 @@ class S3Hook(AwsBaseHook):
                 if arg_name in S3Transfer.ALLOWED_DOWNLOAD_ARGS
             }
 
+        params = sanitize_extra_args()
+        if self._requester_pays:
+            params["RequestPayer"] = "requester"
         obj = self.resource.Object(bucket_name, key)
-        obj.load(**sanitize_extra_args())
+        obj.load(**params)
         return obj
 
     @unify_bucket_name_and_key
@@ -1022,11 +1071,14 @@ class S3Hook(AwsBaseHook):
         """
         expression = expression or "SELECT * FROM S3Object"
         expression_type = expression_type or "SQL"
+        extra_args = {}
 
         if input_serialization is None:
             input_serialization = {"CSV": {}}
         if output_serialization is None:
             output_serialization = {"CSV": {}}
+        if self._requester_pays:
+            extra_args["RequestPayer"] = "requester"
 
         response = self.get_conn().select_object_content(
             Bucket=bucket_name,
@@ -1035,6 +1087,7 @@ class S3Hook(AwsBaseHook):
             ExpressionType=expression_type,
             InputSerialization=input_serialization,
             OutputSerialization=output_serialization,
+            ExtraArgs=extra_args,
         )
 
         return b"".join(
@@ -1124,6 +1177,8 @@ class S3Hook(AwsBaseHook):
                     filename = filename_gz
         if acl_policy:
             extra_args["ACL"] = acl_policy
+        if self._requester_pays:
+            extra_args["RequestPayer"] = "requester"
 
         client = self.get_conn()
         client.upload_file(
@@ -1270,6 +1325,8 @@ class S3Hook(AwsBaseHook):
             extra_args["ServerSideEncryption"] = "AES256"
         if acl_policy:
             extra_args["ACL"] = acl_policy
+        if self._requester_pays:
+            extra_args["RequestPayer"] = "requester"
 
         client = self.get_conn()
         client.upload_fileobj(
@@ -1330,6 +1387,8 @@ class S3Hook(AwsBaseHook):
             kwargs["ACL"] = acl_policy
         if meta_data_directive:
             kwargs["MetadataDirective"] = meta_data_directive
+        if self._requester_pays:
+            kwargs["RequestPayer"] = "requester"
 
         dest_bucket_name, dest_bucket_key = self.get_s3_bucket_key(
             dest_bucket_name, dest_bucket_key, "dest_bucket_name", "dest_bucket_key"
@@ -1412,12 +1471,17 @@ class S3Hook(AwsBaseHook):
             keys = [keys]
 
         s3 = self.get_conn()
+        extra_kwargs = {}
+        if self._requester_pays:
+            extra_kwargs["RequestPayer"] = "requester"
 
         # We can only send a maximum of 1000 keys per request.
         # For details see:
         # https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/s3.html#S3.Client.delete_objects
         for chunk in chunks(keys, chunk_size=1000):
-            response = s3.delete_objects(Bucket=bucket, Delete={"Objects": [{"Key": k} for k in chunk]})
+            response = s3.delete_objects(
+                Bucket=bucket, Delete={"Objects": [{"Key": k} for k in chunk]}, **extra_kwargs
+            )
             deleted_keys = [x["Key"] for x in response.get("Deleted", [])]
             self.log.info("Deleted: %s", deleted_keys)
             if "Errors" in response:
@@ -1496,9 +1560,12 @@ class S3Hook(AwsBaseHook):
             file = NamedTemporaryFile(dir=local_path, prefix="airflow_tmp_", delete=False)  # type: ignore
 
         with file:
+            extra_args = {**self.extra_args}
+            if self._requester_pays:
+                extra_args["RequestPayer"] = "requester"
             s3_obj.download_fileobj(
                 file,
-                ExtraArgs=self.extra_args,
+                ExtraArgs=extra_args,
                 Config=self.transfer_config,
             )
         get_hook_lineage_collector().add_input_asset(

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

@@ -60,6 +60,8 @@ class SnsHook(AwsBaseHook):
         message: str,
         subject: str | None = None,
         message_attributes: dict | None = None,
+        message_deduplication_id: str | None = None,
+        message_group_id: str | None = None,
     ):
         """
         Publish a message to a SNS topic or an endpoint.
@@ -77,7 +79,10 @@ class SnsHook(AwsBaseHook):
             - str = String
             - int, float = Number
             - iterable = String.Array
-
+        :param message_deduplication_id: Every message must have a unique message_deduplication_id.
+            This parameter applies only to FIFO (first-in-first-out) topics.
+        :param message_group_id: Tag that specifies that a message belongs to a specific message group.
+            This parameter applies only to FIFO (first-in-first-out) topics.
         """
         publish_kwargs: dict[str, str | dict] = {
             "TargetArn": target_arn,
@@ -88,6 +93,10 @@ class SnsHook(AwsBaseHook):
         # Construct args this way because boto3 distinguishes from missing args and those set to None
         if subject:
             publish_kwargs["Subject"] = subject
+        if message_deduplication_id:
+            publish_kwargs["MessageDeduplicationId"] = message_deduplication_id
+        if message_group_id:
+            publish_kwargs["MessageGroupId"] = message_group_id
         if message_attributes:
             publish_kwargs["MessageAttributes"] = {
                 key: _get_message_attribute(val) for key, val in message_attributes.items()

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

@@ -125,10 +125,10 @@ class CloudWatchRemoteLogIO(LoggingMixin):  # noqa: D101
         def proc(logger: structlog.typing.WrappedLogger, method_name: str, event: structlog.typing.EventDict):
             if not logger or not (stream_name := relative_path_from_logger(logger)):
                 return event
-            # Only init the handler stream_name once. We cannot do it above when we init the handler because
-            # we don't yet know the log path at that point.
-            if not _handler.log_stream_name:
-                _handler.log_stream_name = stream_name.as_posix().replace(":", "_")
+            # We can't set the log stream name in the above init handler because
+            # the log path isn't known at that stage.
+            # Instead, we should always rely on the path (log stream name) provided by the logger.
+            _handler.log_stream_name = stream_name.as_posix().replace(":", "_")
             name = event.get("logger_name") or event.get("logger", "")
             level = structlog.stdlib.NAME_TO_LEVEL.get(method_name.lower(), logging.INFO)
             msg = copy.copy(event)
@@ -149,7 +149,14 @@ class CloudWatchRemoteLogIO(LoggingMixin):  # noqa: D101
         return (proc,)
 
     def close(self):
-        self.handler.close()
+        # Use the flush method to ensure all logs are sent to CloudWatch.
+        # Closing the handler sets `shutting_down` to True, which prevents any further logs from being sent.
+        # When `shutting_down` is True, means the logging system is in the process of shutting down,
+        # during which it attempts to flush the logs which are queued.
+        if self.handler is None or self.handler.shutting_down:
+            return
+
+        self.handler.flush()
 
     def upload(self, path: os.PathLike | str, ti: RuntimeTI):
         # No-op, as we upload via the processor as we go

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

@@ -32,7 +32,6 @@ from typing import TYPE_CHECKING, Any
 
 from airflow.configuration import conf
 from airflow.exceptions import AirflowException
-from airflow.models.mappedoperator import MappedOperator
 from airflow.providers.amazon.aws.hooks.batch_client import BatchClientHook
 from airflow.providers.amazon.aws.links.batch import (
     BatchJobDefinitionLink,
@@ -145,7 +144,7 @@ class BatchOperator(AwsBaseOperator[BatchClientHook]):
     def operator_extra_links(self):
         op_extra_links = [BatchJobDetailsLink()]
 
-        if isinstance(self, MappedOperator):
+        if self.is_mapped:
             wait_for_completion = self.partial_kwargs.get(
                 "wait_for_completion"
             ) or self.expand_input.value.get("wait_for_completion")

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

@@ -98,13 +98,11 @@ class CloudFormationDeleteStackOperator(AwsBaseOperator[CloudFormationHook]):
         *,
         stack_name: str,
         cloudformation_parameters: dict | None = None,
-        aws_conn_id: str | None = "aws_default",
         **kwargs,
     ):
         super().__init__(**kwargs)
         self.cloudformation_parameters = cloudformation_parameters or {}
         self.stack_name = stack_name
-        self.aws_conn_id = aws_conn_id
 
     def execute(self, context: Context):
         self.log.info("CloudFormation Parameters: %s", self.cloudformation_parameters)

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

@@ -289,7 +289,6 @@ class ComprehendCreateDocumentClassifierOperator(AwsBaseOperator[ComprehendHook]
         waiter_delay: int = 60,
         waiter_max_attempts: int = 20,
         deferrable: bool = conf.getboolean("operators", "default_deferrable", fallback=False),
-        aws_conn_id: str | None = "aws_default",
         **kwargs,
     ):
         super().__init__(**kwargs)
@@ -305,7 +304,6 @@ class ComprehendCreateDocumentClassifierOperator(AwsBaseOperator[ComprehendHook]
         self.waiter_delay = waiter_delay
         self.waiter_max_attempts = waiter_max_attempts
         self.deferrable = deferrable
-        self.aws_conn_id = aws_conn_id
 
     def execute(self, context: Context) -> str:
         if self.output_data_config:

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

@@ -91,7 +91,6 @@ class DmsCreateTaskOperator(AwsBaseOperator[DmsHook]):
         table_mappings: dict,
         migration_type: str = "full-load",
         create_task_kwargs: dict | None = None,
-        aws_conn_id: str | None = "aws_default",
         **kwargs,
     ):
         super().__init__(**kwargs)
@@ -102,7 +101,6 @@ class DmsCreateTaskOperator(AwsBaseOperator[DmsHook]):
         self.migration_type = migration_type
         self.table_mappings = table_mappings
         self.create_task_kwargs = create_task_kwargs or {}
-        self.aws_conn_id = aws_conn_id
 
     def execute(self, context: Context):
         """

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

@@ -517,7 +517,7 @@ class EcsRunTaskOperator(EcsBaseOperator):
         if self.reattach:
             # Generate deterministic UUID which refers to unique TaskInstanceKey
             ti: TaskInstance = context["ti"]
-            self._started_by = generate_uuid(*map(str, ti.key.primary))
+            self._started_by = generate_uuid(*map(str, [ti.dag_id, ti.task_id, ti.run_id, ti.map_index]))
             self.log.info("Try to find run with startedBy=%r", self._started_by)
             self._try_reattach_task(started_by=self._started_by)
 

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

@@ -1056,6 +1056,7 @@ class EksPodOperator(KubernetesPodOperator):
             in_cluster=self.in_cluster,
             namespace=self.namespace,
             name=self.pod_name,
+            trigger_kwargs={"eks_cluster_name": cluster_name},
             **kwargs,
         )
         # There is no need to manage the kube_config file, as it will be generated automatically.
@@ -1072,3 +1073,15 @@ class EksPodOperator(KubernetesPodOperator):
             eks_cluster_name=self.cluster_name, pod_namespace=self.namespace
         ) as self.config_file:
             return super().execute(context)
+
+    def trigger_reentry(self, context: Context, event: dict[str, Any]) -> Any:
+        eks_hook = EksHook(
+            aws_conn_id=self.aws_conn_id,
+            region_name=self.region,
+        )
+        eks_cluster_name = event["eks_cluster_name"]
+        pod_namespace = event["namespace"]
+        with eks_hook.generate_config_file(
+            eks_cluster_name=eks_cluster_name, pod_namespace=pod_namespace
+        ) as self.config_file:
+            return super().trigger_reentry(context, event)

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

@@ -984,7 +984,7 @@ class EmrServerlessCreateApplicationOperator(AwsBaseOperator[EmrServerlessHook])
     :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
-    :waiter_max_attempts: Number of times the waiter should poll the application to check the state.
+    :param waiter_max_attempts: Number of times the waiter should poll the application to check the state.
         If not set, the waiter will use its default value.
     :param waiter_delay: Number of seconds between polling the state of the application.
     :param deferrable: If True, the operator will wait asynchronously for application to be created.
@@ -1128,7 +1128,7 @@ class EmrServerlessStartJobOperator(AwsBaseOperator[EmrServerlessHook]):
     :param verify: Whether or not to verify SSL certificates. See:
         https://boto3.amazonaws.com/v1/documentation/api/latest/reference/core/session.html
     :param name: Name for the EMR Serverless job. If not provided, a default name will be assigned.
-    :waiter_max_attempts: Number of times the waiter should poll the application to check the state.
+    :param waiter_max_attempts: Number of times the waiter should poll the application to check the state.
         If not set, the waiter will use its default value.
     :param waiter_delay: Number of seconds between polling the state of the job run.
     :param deferrable: If True, the operator will wait asynchronously for the crawl to complete.
@@ -1438,7 +1438,7 @@ class EmrServerlessStopApplicationOperator(AwsBaseOperator[EmrServerlessHook]):
         Otherwise, trying to stop an app with running jobs will return an error.
         If you want to wait for the jobs to finish gracefully, use
         :class:`airflow.providers.amazon.aws.sensors.emr.EmrServerlessJobSensor`
-    :waiter_max_attempts: Number of times the waiter should poll the application to check the state.
+    :param waiter_max_attempts: Number of times the waiter should poll the application to check the state.
         Default is 25.
     :param waiter_delay: Number of seconds between polling the state of the application.
         Default is 60 seconds.
@@ -1573,7 +1573,7 @@ class EmrServerlessDeleteApplicationOperator(EmrServerlessStopApplicationOperato
     :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
-    :waiter_max_attempts: Number of times the waiter should poll the application to check the state.
+    :param waiter_max_attempts: Number of times the waiter should poll the application to check the state.
         Defaults to 25.
     :param waiter_delay: Number of seconds between polling the state of the application.
         Defaults to 60 seconds.

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

@@ -313,7 +313,6 @@ class GlueDataQualityOperator(AwsBaseOperator[GlueDataQualityHook]):
         description: str = "AWS Glue Data Quality Rule Set With Airflow",
         update_rule_set: bool = False,
         data_quality_ruleset_kwargs: dict | None = None,
-        aws_conn_id: str | None = "aws_default",
         **kwargs,
     ):
         super().__init__(**kwargs)
@@ -322,7 +321,6 @@ class GlueDataQualityOperator(AwsBaseOperator[GlueDataQualityHook]):
         self.description = description
         self.update_rule_set = update_rule_set
         self.data_quality_ruleset_kwargs = data_quality_ruleset_kwargs or {}
-        self.aws_conn_id = aws_conn_id
 
     def validate_inputs(self) -> None:
         if not self.ruleset.startswith("Rules") or not self.ruleset.endswith("]"):
@@ -421,7 +419,6 @@ class GlueDataQualityRuleSetEvaluationRunOperator(AwsBaseOperator[GlueDataQualit
         waiter_delay: int = 60,
         waiter_max_attempts: int = 20,
         deferrable: bool = conf.getboolean("operators", "default_deferrable", fallback=False),
-        aws_conn_id: str | None = "aws_default",
         **kwargs,
     ):
         super().__init__(**kwargs)
@@ -437,7 +434,6 @@ class GlueDataQualityRuleSetEvaluationRunOperator(AwsBaseOperator[GlueDataQualit
         self.waiter_delay = waiter_delay
         self.waiter_max_attempts = waiter_max_attempts
         self.deferrable = deferrable
-        self.aws_conn_id = aws_conn_id
 
     def validate_inputs(self) -> None:
         glue_table = self.datasource.get("GlueTable", {})
@@ -584,7 +580,6 @@ class GlueDataQualityRuleRecommendationRunOperator(AwsBaseOperator[GlueDataQuali
         waiter_delay: int = 60,
         waiter_max_attempts: int = 20,
         deferrable: bool = conf.getboolean("operators", "default_deferrable", fallback=False),
-        aws_conn_id: str | None = "aws_default",
         **kwargs,
     ):
         super().__init__(**kwargs)
@@ -598,7 +593,6 @@ class GlueDataQualityRuleRecommendationRunOperator(AwsBaseOperator[GlueDataQuali
         self.waiter_delay = waiter_delay
         self.waiter_max_attempts = waiter_max_attempts
         self.deferrable = deferrable
-        self.aws_conn_id = aws_conn_id
 
     def execute(self, context: Context) -> str:
         glue_table = self.datasource.get("GlueTable", {})

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

@@ -54,12 +54,8 @@ class RdsBaseOperator(AwsBaseOperator[RdsHook]):
     def __init__(
         self,
         *args,
-        aws_conn_id: str | None = "aws_conn_id",
-        region_name: str | None = None,
         **kwargs,
     ):
-        self.aws_conn_id = aws_conn_id
-        self.region_name = region_name
         super().__init__(*args, **kwargs)
 
         self._await_interval = 60  # seconds