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

airflow/providers/amazon/__init__.py

@@ -29,7 +29,7 @@ from airflow import __version__ as airflow_version
 
 __all__ = ["__version__"]
 
-__version__ = "9.9.0"
+__version__ = "9.9.1"
 
 if packaging.version.parse(packaging.version.parse(airflow_version).base_version) < packaging.version.parse(
     "2.10.0"

airflow/providers/amazon/aws/executors/batch/batch_executor.py

@@ -36,11 +36,15 @@ from airflow.providers.amazon.aws.executors.utils.exponential_backoff_retry impo
     exponential_backoff_retry,
 )
 from airflow.providers.amazon.aws.hooks.batch_client import BatchClientHook
+from airflow.providers.amazon.version_compat import AIRFLOW_V_3_0_PLUS
 from airflow.stats import Stats
 from airflow.utils import timezone
 from airflow.utils.helpers import merge_dicts
 
 if TYPE_CHECKING:
+    from sqlalchemy.orm import Session
+
+    from airflow.executors import workloads
     from airflow.models.taskinstance import TaskInstance, TaskInstanceKey
 from airflow.providers.amazon.aws.executors.batch.boto_schema import (
     BatchDescribeJobsResponseSchema,
@@ -97,6 +101,11 @@ class AwsBatchExecutor(BaseExecutor):
     # AWS only allows a maximum number of JOBs in the describe_jobs function
     DESCRIBE_JOBS_BATCH_SIZE = 99
 
+    if TYPE_CHECKING and AIRFLOW_V_3_0_PLUS:
+        # In the v3 path, we store workloads, not commands as strings.
+        # TODO: TaskSDK: move this type change into BaseExecutor
+        queued_tasks: dict[TaskInstanceKey, workloads.All]  # type: ignore[assignment]
+
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
         self.active_workers = BatchJobCollection()
@@ -106,6 +115,30 @@ class AwsBatchExecutor(BaseExecutor):
         self.IS_BOTO_CONNECTION_HEALTHY = False
         self.submit_job_kwargs = self._load_submit_kwargs()
 
+    def queue_workload(self, workload: workloads.All, session: Session | None) -> None:
+        from airflow.executors import workloads
+
+        if not isinstance(workload, workloads.ExecuteTask):
+            raise RuntimeError(f"{type(self)} cannot handle workloads of type {type(workload)}")
+        ti = workload.ti
+        self.queued_tasks[ti.key] = workload
+
+    def _process_workloads(self, workloads: Sequence[workloads.All]) -> None:
+        from airflow.executors.workloads import ExecuteTask
+
+        # Airflow V3 version
+        for w in workloads:
+            if not isinstance(w, ExecuteTask):
+                raise RuntimeError(f"{type(self)} cannot handle workloads of type {type(w)}")
+            command = [w]
+            key = w.ti.key
+            queue = w.ti.queue
+            executor_config = w.ti.executor_config or {}
+
+            del self.queued_tasks[key]
+            self.execute_async(key=key, command=command, queue=queue, executor_config=executor_config)  # type: ignore[arg-type]
+            self.running.add(key)
+
     def check_health(self):
         """Make a test API call to check the health of the Batch Executor."""
         success_status = "succeeded."
@@ -343,6 +376,24 @@ class AwsBatchExecutor(BaseExecutor):
         if executor_config and "command" in executor_config:
             raise ValueError('Executor Config should never override "command"')
 
+        if len(command) == 1:
+            from airflow.executors.workloads import ExecuteTask
+
+            if isinstance(command[0], ExecuteTask):
+                workload = command[0]
+                ser_input = workload.model_dump_json()
+                command = [
+                    "python",
+                    "-m",
+                    "airflow.sdk.execution_time.execute_workload",
+                    "--json-string",
+                    ser_input,
+                ]
+            else:
+                raise ValueError(
+                    f"BatchExecutor doesn't know how to handle workload of type: {type(command[0])}"
+                )
+
         self.pending_jobs.append(
             BatchQueuedJob(
                 key=key,

airflow/providers/amazon/aws/executors/ecs/utils.py

@@ -25,9 +25,9 @@ from __future__ import annotations
 
 import datetime
 from collections import defaultdict
-from collections.abc import Sequence
+from collections.abc import Callable, Sequence
 from dataclasses import dataclass
-from typing import TYPE_CHECKING, Any, Callable
+from typing import TYPE_CHECKING, Any
 
 from inflection import camelize
 

airflow/providers/amazon/aws/executors/utils/exponential_backoff_retry.py

@@ -17,8 +17,8 @@
 from __future__ import annotations
 
 import logging
+from collections.abc import Callable
 from datetime import datetime, timedelta
-from typing import Callable
 
 from airflow.utils import timezone
 

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

@@ -18,8 +18,9 @@ from __future__ import annotations
 
 import asyncio
 import logging
+from collections.abc import Callable
 from functools import partial
-from typing import TYPE_CHECKING, Any, Callable
+from typing import TYPE_CHECKING, Any
 
 import requests
 from botocore import UNSIGNED

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

@@ -111,7 +111,14 @@ class AthenaSQLHook(AwsBaseHook, DbApiHook):
                 connection.login = athena_conn.login
                 connection.password = athena_conn.password
                 connection.schema = athena_conn.schema
-                connection.set_extra(json.dumps({**athena_conn.extra_dejson, **connection.extra_dejson}))
+                merged_extra = {**athena_conn.extra_dejson, **connection.extra_dejson}
+                try:
+                    extra_json = json.dumps(merged_extra)
+                    connection.extra = extra_json
+                except (TypeError, ValueError):
+                    raise ValueError(
+                        f"Encountered non-JSON in `extra` field for connection {self.aws_conn_id!r}."
+                    )
             except AirflowNotFoundException:
                 connection = athena_conn
                 connection.conn_type = "aws"
@@ -120,7 +127,10 @@ class AthenaSQLHook(AwsBaseHook, DbApiHook):
                 )
 
         return AwsConnectionWrapper(
-            conn=connection, region_name=self._region_name, botocore_config=self._config, verify=self._verify
+            conn=connection,
+            region_name=self._region_name,
+            botocore_config=self._config,
+            verify=self._verify,
         )
 
     @property

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

@@ -31,10 +31,11 @@ import json
 import logging
 import os
 import warnings
+from collections.abc import Callable
 from copy import deepcopy
 from functools import cached_property, wraps
 from pathlib import Path
-from typing import TYPE_CHECKING, Any, Callable, Generic, TypeVar, Union
+from typing import TYPE_CHECKING, Any, Generic, TypeVar, Union
 
 import boto3
 import botocore
@@ -43,6 +44,8 @@ import jinja2
 import requests
 import tenacity
 from asgiref.sync import sync_to_async
+from boto3.resources.base import ServiceResource
+from botocore.client import BaseClient
 from botocore.config import Config
 from botocore.waiter import Waiter, WaiterModel
 from dateutil.tz import tzlocal
@@ -54,16 +57,29 @@ from airflow.exceptions import (
     AirflowNotFoundException,
     AirflowProviderDeprecationWarning,
 )
-from airflow.hooks.base import BaseHook
 from airflow.providers.amazon.aws.utils.connection_wrapper import AwsConnectionWrapper
 from airflow.providers.amazon.aws.utils.identifiers import generate_uuid
 from airflow.providers.amazon.aws.utils.suppress import return_on_error
 from airflow.providers.common.compat.version_compat import AIRFLOW_V_3_0_PLUS
 from airflow.providers_manager import ProvidersManager
+
+try:
+    from airflow.sdk import BaseHook
+except ImportError:
+    from airflow.hooks.base import BaseHook  # type: ignore[attr-defined,no-redef]
 from airflow.utils.helpers import exactly_one
 from airflow.utils.log.logging_mixin import LoggingMixin
 
-BaseAwsConnection = TypeVar("BaseAwsConnection", bound=Union[boto3.client, boto3.resource])
+# We need to set typeignore, sadly without it Sphinx build and mypy don't agree.
+# ideally the code should be:
+# BaseAwsConnection = TypeVar("BaseAwsConnection", bound=BaseClient | ServiceResource)
+# but if we do that Sphinx complains about:
+# TypeError: unsupported operand type(s) for |: 'BaseClient' and 'ServiceResource'
+# If we change to Union syntax then mypy is not happy with UP007 Use `X | Y` for type annotations
+# The only way to workaround it for now is to keep the union syntax with ignore for mypy
+# We should try to resolve this later.
+BaseAwsConnection = TypeVar("BaseAwsConnection", bound=Union[BaseClient, ServiceResource])  # type: ignore[operator] # noqa: UP007
+
 
 if AIRFLOW_V_3_0_PLUS:
     from airflow.sdk.exceptions import AirflowRuntimeError
@@ -627,7 +643,10 @@ class AwsGenericHook(BaseHook, Generic[BaseAwsConnection]):
                     raise
 
         return AwsConnectionWrapper(
-            conn=connection, region_name=self._region_name, botocore_config=self._config, verify=self._verify
+            conn=connection,  # type: ignore[arg-type]
+            region_name=self._region_name,
+            botocore_config=self._config,
+            verify=self._verify,
         )
 
     def _resolve_service_name(self, is_resource_type: bool = False) -> str:
@@ -1038,7 +1057,7 @@ class AwsGenericHook(BaseHook, Generic[BaseAwsConnection]):
             return WaiterModel(model_config).waiter_names
 
 
-class AwsBaseHook(AwsGenericHook[Union[boto3.client, boto3.resource]]):
+class AwsBaseHook(AwsGenericHook[Union[boto3.client, boto3.resource]]):  # type: ignore[operator] # noqa: UP007
     """
     Base class for interact with AWS.
 

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

@@ -30,7 +30,8 @@ from __future__ import annotations
 import itertools
 import random
 import time
-from typing import TYPE_CHECKING, Callable, Protocol, runtime_checkable
+from collections.abc import Callable
+from typing import TYPE_CHECKING, Protocol, runtime_checkable
 
 import botocore.client
 import botocore.exceptions

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

@@ -28,9 +28,10 @@ from __future__ import annotations
 
 import json
 import sys
+from collections.abc import Callable
 from copy import deepcopy
 from pathlib import Path
-from typing import TYPE_CHECKING, Any, Callable
+from typing import TYPE_CHECKING, Any
 
 import botocore.client
 import botocore.exceptions

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

@@ -66,9 +66,13 @@ class ChimeWebhookHook(HttpHook):
         :return: Endpoint(str) for chime webhook.
         """
         conn = self.get_connection(conn_id)
-        token = conn.get_password()
+        token = conn.password
         if token is None:
             raise AirflowException("Webhook token field is missing and is required.")
+        if not conn.schema:
+            raise AirflowException("Webook schema field is missing and is required")
+        if not conn.host:
+            raise AirflowException("Webhook host field is missing and is required.")
         url = conn.schema + "://" + conn.host
         endpoint = url + token
         # Check to make sure the endpoint matches what Chime expects

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

@@ -19,7 +19,8 @@ from __future__ import annotations
 
 import functools
 import time
-from typing import Callable, TypeVar
+from collections.abc import Callable
+from typing import TypeVar
 
 from airflow.exceptions import AirflowException
 from airflow.providers.amazon.aws.hooks.base_aws import AwsBaseHook

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

@@ -23,11 +23,10 @@ import json
 import os
 import sys
 import tempfile
-from collections.abc import Generator
+from collections.abc import Callable, Generator
 from contextlib import contextmanager
 from enum import Enum
 from functools import partial
-from typing import Callable
 
 from botocore.exceptions import ClientError
 from botocore.signers import RequestSigner

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
 

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,9 +25,13 @@ 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
 
+try:
+    from airflow.sdk import BaseHook
+except ImportError:
+    from airflow.hooks.base import BaseHook  # type: ignore[attr-defined,no-redef]
+
 
 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(