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

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

@@ -24,6 +24,14 @@ from functools import cached_property
 from typing import Any
 
 from botocore.exceptions import ClientError
+from tenacity import (
+    AsyncRetrying,
+    Retrying,
+    before_sleep_log,
+    retry_if_exception,
+    stop_after_attempt,
+    wait_exponential,
+)
 
 from airflow.exceptions import AirflowException, AirflowProviderDeprecationWarning
 from airflow.providers.amazon.aws.hooks.base_aws import AwsBaseHook
@@ -46,11 +54,11 @@ class GlueJobHook(AwsBaseHook):
     :param script_location: path to etl script on s3
     :param retry_limit: Maximum number of times to retry this job if it fails
     :param num_of_dpus: Number of AWS Glue DPUs to allocate to this Job
-    :param region_name: aws region name (example: us-east-1)
     :param iam_role_name: AWS IAM Role for Glue Job Execution. If set `iam_role_arn` must equal None.
     :param iam_role_arn: AWS IAM Role ARN for Glue Job Execution, If set `iam_role_name` must equal None.
     :param create_job_kwargs: Extra arguments for Glue Job Creation
     :param update_config: Update job configuration on Glue (default: False)
+    :param api_retry_args: An optional dictionary with arguments passed to ``tenacity.Retrying`` & ``tenacity.AsyncRetrying`` classes.
 
     Additional arguments (such as ``aws_conn_id``) may be specified and
     are passed down to the underlying AwsBaseHook.
@@ -80,6 +88,7 @@ class GlueJobHook(AwsBaseHook):
         create_job_kwargs: dict | None = None,
         update_config: bool = False,
         job_poll_interval: int | float = 6,
+        api_retry_args: dict[Any, Any] | None = None,
         *args,
         **kwargs,
     ):
@@ -96,6 +105,17 @@ class GlueJobHook(AwsBaseHook):
         self.update_config = update_config
         self.job_poll_interval = job_poll_interval
 
+        self.retry_config: dict[str, Any] = {
+            "retry": retry_if_exception(self._should_retry_on_error),
+            "wait": wait_exponential(multiplier=1, min=1, max=60),
+            "stop": stop_after_attempt(5),
+            "before_sleep": before_sleep_log(self.log, log_level=20),
+            "reraise": True,
+        }
+
+        if api_retry_args:
+            self.retry_config.update(api_retry_args)
+
         worker_type_exists = "WorkerType" in self.create_job_kwargs
         num_workers_exists = "NumberOfWorkers" in self.create_job_kwargs
 
@@ -116,6 +136,29 @@ class GlueJobHook(AwsBaseHook):
         kwargs["client_type"] = "glue"
         super().__init__(*args, **kwargs)
 
+    def _should_retry_on_error(self, exception: BaseException) -> bool:
+        """
+        Determine if an exception should trigger a retry.
+
+        :param exception: The exception that occurred
+        :return: True if the exception should trigger a retry, False otherwise
+        """
+        if isinstance(exception, ClientError):
+            error_code = exception.response.get("Error", {}).get("Code", "")
+            retryable_errors = {
+                "ThrottlingException",
+                "RequestLimitExceeded",
+                "ServiceUnavailable",
+                "InternalFailure",
+                "InternalServerError",
+                "TooManyRequestsException",
+                "RequestTimeout",
+                "RequestTimeoutException",
+                "HttpTimeoutException",
+            }
+            return error_code in retryable_errors
+        return False
+
     def create_glue_job_config(self) -> dict:
         default_command = {
             "Name": "glueetl",
@@ -217,8 +260,21 @@ class GlueJobHook(AwsBaseHook):
         :param run_id: The job-run ID of the predecessor job run
         :return: State of the Glue job
         """
-        job_run = self.conn.get_job_run(JobName=job_name, RunId=run_id, PredecessorsIncluded=True)
-        return job_run["JobRun"]["JobRunState"]
+        for attempt in Retrying(**self.retry_config):
+            with attempt:
+                try:
+                    job_run = self.conn.get_job_run(JobName=job_name, RunId=run_id, PredecessorsIncluded=True)
+                    return job_run["JobRun"]["JobRunState"]
+                except ClientError as e:
+                    self.log.error("Failed to get job state for job %s run %s: %s", job_name, run_id, e)
+                    raise
+                except Exception as e:
+                    self.log.error(
+                        "Unexpected error getting job state for job %s run %s: %s", job_name, run_id, e
+                    )
+                    raise
+        # This should never be reached due to reraise=True, but mypy needs it
+        raise RuntimeError("Unexpected end of retry loop")
 
     async def async_get_job_state(self, job_name: str, run_id: str) -> str:
         """
@@ -226,9 +282,22 @@ class GlueJobHook(AwsBaseHook):
 
         The async version of get_job_state.
         """
-        async with await self.get_async_conn() as client:
-            job_run = await client.get_job_run(JobName=job_name, RunId=run_id)
-        return job_run["JobRun"]["JobRunState"]
+        async for attempt in AsyncRetrying(**self.retry_config):
+            with attempt:
+                try:
+                    async with await self.get_async_conn() as client:
+                        job_run = await client.get_job_run(JobName=job_name, RunId=run_id)
+                    return job_run["JobRun"]["JobRunState"]
+                except ClientError as e:
+                    self.log.error("Failed to get job state for job %s run %s: %s", job_name, run_id, e)
+                    raise
+                except Exception as e:
+                    self.log.error(
+                        "Unexpected error getting job state for job %s run %s: %s", job_name, run_id, e
+                    )
+                    raise
+        # This should never be reached due to reraise=True, but mypy needs it
+        raise RuntimeError("Unexpected end of retry loop")
 
     @cached_property
     def logs_hook(self):
@@ -372,7 +441,7 @@ class GlueJobHook(AwsBaseHook):
         )
         return None
 
-    def has_job(self, job_name) -> bool:
+    def has_job(self, job_name: str) -> bool:
         """
         Check if the job already exists.
 
@@ -422,6 +491,9 @@ class GlueJobHook(AwsBaseHook):
 
         :return:Name of the Job
         """
+        if self.job_name is None:
+            raise ValueError("job_name must be set to get or create a Glue job")
+
         if self.has_job(self.job_name):
             return self.job_name
 
@@ -441,6 +513,9 @@ class GlueJobHook(AwsBaseHook):
 
         :return:Name of the Job
         """
+        if self.job_name is None:
+            raise ValueError("job_name must be set to create or update a Glue job")
+
         config = self.create_glue_job_config()
 
         if self.has_job(self.job_name):

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

@@ -20,7 +20,8 @@
 from __future__ import annotations
 
 import time
-from typing import TYPE_CHECKING, Callable
+from collections.abc import Callable
+from typing import TYPE_CHECKING
 
 from airflow.exceptions import AirflowException, AirflowNotFoundException
 from airflow.providers.amazon.aws.hooks.base_aws import AwsGenericHook

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

@@ -28,7 +28,7 @@ import os
 import re
 import shutil
 import time
-from collections.abc import AsyncIterator
+from collections.abc import AsyncIterator, Callable
 from contextlib import suppress
 from copy import deepcopy
 from datetime import datetime
@@ -37,7 +37,7 @@ from inspect import signature
 from io import BytesIO
 from pathlib import Path
 from tempfile import NamedTemporaryFile, gettempdir
-from typing import TYPE_CHECKING, Any, Callable
+from typing import TYPE_CHECKING, Any
 from urllib.parse import urlsplit
 from uuid import uuid4
 
@@ -635,6 +635,10 @@ class S3Hook(AwsBaseHook):
         delimiter: str | None = "/",
     ) -> list[Any]:
         """Get a list of files in the bucket."""
+        # Validate that bucket_keys is in fact a list, otherwise, the characters will be split
+        if isinstance(bucket_keys, str):
+            bucket_keys = [bucket_keys]
+
         keys: list[Any] = []
         for key in bucket_keys:
             prefix = key
@@ -652,7 +656,9 @@ class S3Hook(AwsBaseHook):
             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)))
+                    keys.extend(
+                        k.get("Key") for k in page["Contents"] if isinstance(k.get("Size"), (int, float))
+                    )
         return keys
 
     async def _list_keys_async(
@@ -1683,3 +1689,80 @@ class S3Hook(AwsBaseHook):
         """
         s3_client = self.get_conn()
         s3_client.delete_bucket_tagging(Bucket=bucket_name)
+
+    def _sync_to_local_dir_delete_stale_local_files(self, current_s3_objects: list[Path], local_dir: Path):
+        current_s3_keys = {key for key in current_s3_objects}
+
+        for item in local_dir.iterdir():
+            item: Path  # type: ignore[no-redef]
+            absolute_item_path = item.resolve()
+
+            if absolute_item_path not in current_s3_keys:
+                try:
+                    if item.is_file():
+                        item.unlink(missing_ok=True)
+                        self.log.debug("Deleted stale local file: %s", item)
+                    elif item.is_dir():
+                        # delete only when the folder is empty
+                        if not os.listdir(item):
+                            item.rmdir()
+                            self.log.debug("Deleted stale empty directory: %s", item)
+                    else:
+                        self.log.debug("Skipping stale item of unknown type: %s", item)
+                except OSError as e:
+                    self.log.error("Error deleting stale item %s: %s", item, e)
+                    raise e
+
+    def _sync_to_local_dir_if_changed(self, s3_bucket, s3_object, local_target_path: Path):
+        should_download = False
+        download_msg = ""
+        if not local_target_path.exists():
+            should_download = True
+            download_msg = f"Local file {local_target_path} does not exist."
+        else:
+            local_stats = local_target_path.stat()
+
+            if s3_object.size != local_stats.st_size:
+                should_download = True
+                download_msg = (
+                    f"S3 object size ({s3_object.size}) and local file size ({local_stats.st_size}) differ."
+                )
+
+            s3_last_modified = s3_object.last_modified
+            if local_stats.st_mtime < s3_last_modified.microsecond:
+                should_download = True
+                download_msg = f"S3 object last modified ({s3_last_modified.microsecond}) and local file last modified ({local_stats.st_mtime}) differ."
+
+        if should_download:
+            s3_bucket.download_file(s3_object.key, local_target_path)
+            self.log.debug(
+                "%s Downloaded %s to %s", download_msg, s3_object.key, local_target_path.as_posix()
+            )
+        else:
+            self.log.debug(
+                "Local file %s is up-to-date with S3 object %s. Skipping download.",
+                local_target_path.as_posix(),
+                s3_object.key,
+            )
+
+    def sync_to_local_dir(self, bucket_name: str, local_dir: Path, s3_prefix="", delete_stale: bool = True):
+        """Download S3 files from the S3 bucket to the local directory."""
+        self.log.debug("Downloading data from s3://%s/%s to %s", bucket_name, s3_prefix, local_dir)
+
+        local_s3_objects = []
+        s3_bucket = self.get_bucket(bucket_name)
+        for obj in s3_bucket.objects.filter(Prefix=s3_prefix):
+            obj_path = Path(obj.key)
+            local_target_path = local_dir.joinpath(obj_path.relative_to(s3_prefix))
+            if not local_target_path.parent.exists():
+                local_target_path.parent.mkdir(parents=True, exist_ok=True)
+                self.log.debug("Created local directory: %s", local_target_path.parent)
+            self._sync_to_local_dir_if_changed(
+                s3_bucket=s3_bucket, s3_object=obj, local_target_path=local_target_path
+            )
+            local_s3_objects.append(local_target_path)
+
+        if delete_stale:
+            self._sync_to_local_dir_delete_stale_local_files(
+                current_s3_objects=local_s3_objects, local_dir=local_dir
+            )

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

@@ -23,10 +23,10 @@ import tarfile
 import tempfile
 import time
 from collections import Counter, namedtuple
-from collections.abc import AsyncGenerator, Generator
+from collections.abc import AsyncGenerator, Callable, Generator
 from datetime import datetime
 from functools import partial
-from typing import Any, Callable, cast
+from typing import Any, cast
 
 from asgiref.sync import sync_to_async
 from botocore.exceptions import ClientError

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

@@ -25,8 +25,8 @@ from sagemaker_studio import ClientConfig
 from sagemaker_studio.sagemaker_studio_api import SageMakerStudioAPI
 
 from airflow.exceptions import AirflowException
-from airflow.hooks.base import BaseHook
 from airflow.providers.amazon.aws.utils.sagemaker_unified_studio import is_local_runner
+from airflow.providers.amazon.version_compat import BaseHook
 
 
 class SageMakerNotebookHook(BaseHook):

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

@@ -20,20 +20,13 @@ from __future__ import annotations
 from typing import TYPE_CHECKING, ClassVar
 
 from airflow.providers.amazon.aws.utils.suppress import return_on_error
-from airflow.providers.amazon.version_compat import AIRFLOW_V_3_0_PLUS
+from airflow.providers.amazon.version_compat import BaseOperatorLink, XCom
 
 if TYPE_CHECKING:
     from airflow.models import BaseOperator
     from airflow.models.taskinstancekey import TaskInstanceKey
     from airflow.utils.context import Context
 
-if AIRFLOW_V_3_0_PLUS:
-    from airflow.sdk import BaseOperatorLink
-    from airflow.sdk.execution_time.xcom import XCom
-else:
-    from airflow.models import XCom  # type: ignore[no-redef]
-    from airflow.models.baseoperatorlink import BaseOperatorLink  # type: ignore[no-redef]
-
 
 BASE_AWS_CONSOLE_LINK = "https://console.{aws_domain}"
 
@@ -94,8 +87,7 @@ class BaseAwsLink(BaseOperatorLink):
         if not operator.do_xcom_push:
             return
 
-        operator.xcom_push(
-            context,
+        context["ti"].xcom_push(
             key=cls.key,
             value={
                 "region_name": region_name,

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

@@ -19,13 +19,13 @@ from __future__ import annotations
 
 from collections.abc import Sequence
 
-from airflow.models import BaseOperator
 from airflow.providers.amazon.aws.utils.mixins import (
     AwsBaseHookMixin,
     AwsHookParams,
     AwsHookType,
     aws_template_fields,
 )
+from airflow.providers.amazon.version_compat import BaseOperator
 from airflow.utils.types import NOTSET, ArgNotSet
 
 

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

@@ -140,28 +140,12 @@ class BatchOperator(AwsBaseOperator[BatchClientHook]):
         "retry_strategy": "json",
     }
 
-    @property
-    def operator_extra_links(self):
-        op_extra_links = [BatchJobDetailsLink()]
-
-        if self.is_mapped:
-            wait_for_completion = self.partial_kwargs.get(
-                "wait_for_completion"
-            ) or self.expand_input.value.get("wait_for_completion")
-            array_properties = self.partial_kwargs.get("array_properties") or self.expand_input.value.get(
-                "array_properties"
-            )
-        else:
-            wait_for_completion = self.wait_for_completion
-            array_properties = self.array_properties
-
-        if wait_for_completion:
-            op_extra_links.extend([BatchJobDefinitionLink(), BatchJobQueueLink()])
-        if not array_properties:
-            # There is no CloudWatch Link to the parent Batch Job available.
-            op_extra_links.append(CloudWatchEventsLink())
-
-        return tuple(op_extra_links)
+    operator_extra_links = (
+        BatchJobDetailsLink(),
+        BatchJobDefinitionLink(),
+        BatchJobQueueLink(),
+        CloudWatchEventsLink(),
+    )
 
     def __init__(
         self,

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

@@ -526,7 +526,7 @@ class EcsRunTaskOperator(EcsBaseOperator):
             self._start_task()
 
         if self.do_xcom_push:
-            self.xcom_push(context, key="ecs_task_arn", value=self.arn)
+            context["ti"].xcom_push(key="ecs_task_arn", value=self.arn)
 
         if self.deferrable:
             self.defer(

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

@@ -60,7 +60,6 @@ class GlueJobOperator(AwsBaseOperator[GlueJobHook]):
     :param script_args: etl script arguments and AWS Glue arguments (templated)
     :param retry_limit: The maximum number of times to retry this job if it fails
     :param num_of_dpus: Number of AWS Glue DPUs to allocate to this Job.
-    :param region_name: aws region name (example: us-east-1)
     :param s3_bucket: S3 bucket where logs and local etl script will be uploaded
     :param iam_role_name: AWS IAM Role for Glue Job Execution. If set `iam_role_arn` must equal None.
     :param iam_role_arn: AWS IAM ARN for Glue Job Execution. If set `iam_role_name` must equal None.
@@ -78,7 +77,20 @@ class GlueJobOperator(AwsBaseOperator[GlueJobHook]):
         of limiting concurrency, Glue needs 5-10 seconds to clean up resources.
         Thus if status is returned immediately it might end up in case of more than 1 concurrent run.
         It is recommended to set this parameter to 10 when you are using concurrency=1.
-        For more information see: https://repost.aws/questions/QUaKgpLBMPSGWO0iq2Fob_bw/glue-run-concurrent-jobs#ANFpCL2fRnQRqgDFuIU_rpvA
+        For more information see:
+        https://repost.aws/questions/QUaKgpLBMPSGWO0iq2Fob_bw/glue-run-concurrent-jobs#ANFpCL2fRnQRqgDFuIU_rpvA
+    :param waiter_delay: Time in seconds to wait between status checks. (default: 60)
+    :param waiter_max_attempts: Maximum number of attempts to check for job completion. (default: 20)
+    :param aws_conn_id: The Airflow connection used for AWS credentials.
+        If this is ``None`` or empty then the default boto3 behaviour is used. If
+        running Airflow in a distributed manner and aws_conn_id is None or
+        empty, then default boto3 configuration would be used (and must be
+        maintained on each worker node).
+    :param region_name: AWS region_name. If not specified then the default boto3 behaviour is used.
+    :param verify: Whether or not to verify SSL certificates. See:
+        https://boto3.amazonaws.com/v1/documentation/api/latest/reference/core/session.html
+    :param botocore_config: Configuration dictionary (key-values) for botocore client. See:
+        https://botocore.amazonaws.com/v1/documentation/api/latest/reference/config.html
     """
 
     aws_hook_class = GlueJobHook
@@ -122,9 +134,11 @@ class GlueJobOperator(AwsBaseOperator[GlueJobHook]):
         verbose: bool = False,
         replace_script_file: bool = False,
         update_config: bool = False,
-        job_poll_interval: int | float = 6,
         stop_job_run_on_kill: bool = False,
         sleep_before_return: int = 0,
+        job_poll_interval: int | float = 6,
+        waiter_delay: int = 60,
+        waiter_max_attempts: int = 75,
         **kwargs,
     ):
         super().__init__(**kwargs)
@@ -152,6 +166,8 @@ class GlueJobOperator(AwsBaseOperator[GlueJobHook]):
         self._job_run_id: str | None = None
         self.sleep_before_return: int = sleep_before_return
         self.s3_script_location: str | None = None
+        self.waiter_delay = waiter_delay
+        self.waiter_max_attempts = waiter_max_attempts
 
     @property
     def _hook_parameters(self):
@@ -231,7 +247,9 @@ class GlueJobOperator(AwsBaseOperator[GlueJobHook]):
                     run_id=self._job_run_id,
                     verbose=self.verbose,
                     aws_conn_id=self.aws_conn_id,
-                    job_poll_interval=self.job_poll_interval,
+                    waiter_delay=self.waiter_delay,
+                    waiter_max_attempts=self.waiter_max_attempts,
+                    region_name=self.region_name,
                 ),
                 method_name="execute_complete",
             )
@@ -254,7 +272,7 @@ class GlueJobOperator(AwsBaseOperator[GlueJobHook]):
 
         if validated_event["status"] != "success":
             raise AirflowException(f"Error in glue job: {validated_event}")
-        return validated_event["value"]
+        return validated_event["run_id"]
 
     def on_kill(self):
         """Cancel the running AWS Glue Job."""
@@ -282,7 +300,6 @@ class GlueDataQualityOperator(AwsBaseOperator[GlueDataQualityHook]):
     :param description: A description of the data quality ruleset.
     :param update_rule_set: To update existing ruleset, Set this flag to True. (default: False)
     :param data_quality_ruleset_kwargs: Extra arguments for RuleSet.
-
     :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
@@ -378,7 +395,6 @@ class GlueDataQualityRuleSetEvaluationRunOperator(AwsBaseOperator[GlueDataQualit
     :param deferrable: If True, the operator will wait asynchronously for the job to stop.
         This implies waiting for completion. This mode requires aiobotocore module to be installed.
         (default: False)
-
     :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
@@ -543,7 +559,6 @@ class GlueDataQualityRuleRecommendationRunOperator(AwsBaseOperator[GlueDataQuali
     :param deferrable: If True, the operator will wait asynchronously for the job to stop.
         This implies waiting for completion. This mode requires aiobotocore module to be installed.
         (default: False)
-
     :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

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

@@ -159,7 +159,7 @@ class RedshiftDataOperator(AwsBaseOperator[RedshiftDataHook]):
         self.statement_id: str = query_execution_output.statement_id
 
         if query_execution_output.session_id:
-            self.xcom_push(context, key="session_id", value=query_execution_output.session_id)
+            context["ti"].xcom_push(key="session_id", value=query_execution_output.session_id)
 
         if self.deferrable and self.wait_for_completion:
             is_finished: bool = self.hook.check_query_is_finished(self.statement_id)

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

@@ -20,8 +20,8 @@ import datetime
 import json
 import time
 import urllib
-from collections.abc import Sequence
-from typing import TYPE_CHECKING, Any, Callable, ClassVar
+from collections.abc import Callable, Sequence
+from typing import TYPE_CHECKING, Any, ClassVar
 
 from botocore.exceptions import ClientError
 

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

@@ -24,7 +24,6 @@ from typing import TYPE_CHECKING
 
 from airflow.configuration import conf
 from airflow.exceptions import AirflowException
-from airflow.models import BaseOperator
 from airflow.providers.amazon.aws.hooks.sagemaker_unified_studio import (
     SageMakerNotebookHook,
 )
@@ -34,6 +33,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.amazon.version_compat import BaseOperator
 
 if TYPE_CHECKING:
     from airflow.utils.context import Context

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

@@ -25,7 +25,7 @@ from airflow.providers.amazon.aws.utils.mixins import (
     AwsHookType,
     aws_template_fields,
 )
-from airflow.sensors.base import BaseSensorOperator
+from airflow.providers.amazon.version_compat import BaseSensorOperator
 from airflow.utils.types import NOTSET, ArgNotSet