apache-airflow-providers-amazon 9.9.0rc1__py3-none-any.whl → 9.10.0rc1__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.10.0"
 
 if packaging.version.parse(packaging.version.parse(airflow_version).base_version) < packaging.version.parse(
     "2.10.0"

airflow/providers/amazon/aws/auth_manager/avp/facade.py

@@ -37,6 +37,13 @@ from airflow.utils.log.logging_mixin import LoggingMixin
 
 if TYPE_CHECKING:
     from airflow.api_fastapi.auth.managers.base_auth_manager import ResourceMethod
+
+    try:
+        from airflow.api_fastapi.auth.managers.base_auth_manager import ExtendedResourceMethod
+    except ImportError:
+        from airflow.api_fastapi.auth.managers.base_auth_manager import (
+            ResourceMethod as ExtendedResourceMethod,
+        )
     from airflow.providers.amazon.aws.auth_manager.user import AwsAuthManagerUser
 
 
@@ -48,7 +55,7 @@ NB_REQUESTS_PER_BATCH = 30
 class IsAuthorizedRequest(TypedDict, total=False):
     """Represent the parameters of ``is_authorized`` method in AVP facade."""
 
-    method: ResourceMethod
+    method: ExtendedResourceMethod
     entity_type: AvpEntities
     entity_id: str | None
     context: dict | None

airflow/providers/amazon/aws/auth_manager/aws_auth_manager.py

@@ -44,10 +44,7 @@ from airflow.providers.amazon.version_compat import AIRFLOW_V_3_0_PLUS
 if TYPE_CHECKING:
     from airflow.api_fastapi.auth.managers.base_auth_manager import ResourceMethod
     from airflow.api_fastapi.auth.managers.models.batch_apis import (
-        IsAuthorizedConnectionRequest,
         IsAuthorizedDagRequest,
-        IsAuthorizedPoolRequest,
-        IsAuthorizedVariableRequest,
     )
     from airflow.api_fastapi.auth.managers.models.resource_details import (
         AccessView,
@@ -247,24 +244,6 @@ class AwsAuthManager(BaseAuthManager[AwsAuthManagerUser]):
 
         return [menu_item for menu_item in menu_items if _has_access_to_menu_item(requests[menu_item.value])]
 
-    def batch_is_authorized_connection(
-        self,
-        requests: Sequence[IsAuthorizedConnectionRequest],
-        *,
-        user: AwsAuthManagerUser,
-    ) -> bool:
-        facade_requests: Sequence[IsAuthorizedRequest] = [
-            {
-                "method": request["method"],
-                "entity_type": AvpEntities.CONNECTION,
-                "entity_id": cast("ConnectionDetails", request["details"]).conn_id
-                if request.get("details")
-                else None,
-            }
-            for request in requests
-        ]
-        return self.avp_facade.batch_is_authorized(requests=facade_requests, user=user)
-
     def batch_is_authorized_dag(
         self,
         requests: Sequence[IsAuthorizedDagRequest],
@@ -288,40 +267,6 @@ class AwsAuthManager(BaseAuthManager[AwsAuthManagerUser]):
         ]
         return self.avp_facade.batch_is_authorized(requests=facade_requests, user=user)
 
-    def batch_is_authorized_pool(
-        self,
-        requests: Sequence[IsAuthorizedPoolRequest],
-        *,
-        user: AwsAuthManagerUser,
-    ) -> bool:
-        facade_requests: Sequence[IsAuthorizedRequest] = [
-            {
-                "method": request["method"],
-                "entity_type": AvpEntities.POOL,
-                "entity_id": cast("PoolDetails", request["details"]).name if request.get("details") else None,
-            }
-            for request in requests
-        ]
-        return self.avp_facade.batch_is_authorized(requests=facade_requests, user=user)
-
-    def batch_is_authorized_variable(
-        self,
-        requests: Sequence[IsAuthorizedVariableRequest],
-        *,
-        user: AwsAuthManagerUser,
-    ) -> bool:
-        facade_requests: Sequence[IsAuthorizedRequest] = [
-            {
-                "method": request["method"],
-                "entity_type": AvpEntities.VARIABLE,
-                "entity_id": cast("VariableDetails", request["details"]).key
-                if request.get("details")
-                else None,
-            }
-            for request in requests
-        ]
-        return self.avp_facade.batch_is_authorized(requests=facade_requests, user=user)
-
     def filter_authorized_dag_ids(
         self,
         *,

airflow/providers/amazon/aws/bundles/__init__.py

@@ -0,0 +1,16 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.

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

@@ -0,0 +1,152 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+import os
+from pathlib import Path
+
+import structlog
+
+from airflow.dag_processing.bundles.base import BaseDagBundle
+from airflow.exceptions import AirflowException
+from airflow.providers.amazon.aws.hooks.base_aws import AwsBaseHook
+from airflow.providers.amazon.aws.hooks.s3 import S3Hook
+
+
+class S3DagBundle(BaseDagBundle):
+    """
+    S3 DAG bundle - exposes a directory in S3 as a DAG bundle.
+
+    This allows Airflow to load DAGs directly from an S3 bucket.
+
+    :param aws_conn_id: Airflow connection ID for AWS.  Defaults to AwsBaseHook.default_conn_name.
+    :param bucket_name: The name of the S3 bucket containing the DAG files.
+    :param prefix:  Optional subdirectory within the S3 bucket where the DAGs are stored.
+                    If None, DAGs are assumed to be at the root of the bucket (Optional).
+    """
+
+    supports_versioning = False
+
+    def __init__(
+        self,
+        *,
+        aws_conn_id: str = AwsBaseHook.default_conn_name,
+        bucket_name: str,
+        prefix: str = "",
+        **kwargs,
+    ) -> None:
+        super().__init__(**kwargs)
+        self.aws_conn_id = aws_conn_id
+        self.bucket_name = bucket_name
+        self.prefix = prefix
+        # Local path where S3 DAGs are downloaded
+        self.s3_dags_dir: Path = self.base_dir
+
+        log = structlog.get_logger(__name__)
+        self._log = log.bind(
+            bundle_name=self.name,
+            version=self.version,
+            bucket_name=self.bucket_name,
+            prefix=self.prefix,
+            aws_conn_id=self.aws_conn_id,
+        )
+        self._s3_hook: S3Hook | None = None
+
+    def _initialize(self):
+        with self.lock():
+            if not self.s3_dags_dir.exists():
+                self._log.info("Creating local DAGs directory: %s", self.s3_dags_dir)
+                os.makedirs(self.s3_dags_dir)
+
+            if not self.s3_dags_dir.is_dir():
+                raise AirflowException(f"Local DAGs path: {self.s3_dags_dir} is not a directory.")
+
+            if not self.s3_hook.check_for_bucket(bucket_name=self.bucket_name):
+                raise AirflowException(f"S3 bucket '{self.bucket_name}' does not exist.")
+
+            if self.prefix:
+                # don't check when prefix is ""
+                if not self.s3_hook.check_for_prefix(
+                    bucket_name=self.bucket_name, prefix=self.prefix, delimiter="/"
+                ):
+                    raise AirflowException(
+                        f"S3 prefix 's3://{self.bucket_name}/{self.prefix}' does not exist."
+                    )
+            self.refresh()
+
+    def initialize(self) -> None:
+        self._initialize()
+        super().initialize()
+
+    @property
+    def s3_hook(self):
+        if self._s3_hook is None:
+            try:
+                self._s3_hook: S3Hook = S3Hook(aws_conn_id=self.aws_conn_id)  # Initialize S3 hook.
+            except AirflowException as e:
+                self._log.warning("Could not create S3Hook for connection %s: %s", self.aws_conn_id, e)
+        return self._s3_hook
+
+    def __repr__(self):
+        return (
+            f"<S3DagBundle("
+            f"name={self.name!r}, "
+            f"bucket_name={self.bucket_name!r}, "
+            f"prefix={self.prefix!r}, "
+            f"version={self.version!r}"
+            f")>"
+        )
+
+    def get_current_version(self) -> str | None:
+        """Return the current version of the DAG bundle. Currently not supported."""
+        return None
+
+    @property
+    def path(self) -> Path:
+        """Return the local path to the DAG files."""
+        return self.s3_dags_dir  # Path where DAGs are downloaded.
+
+    def refresh(self) -> None:
+        """Refresh the DAG bundle by re-downloading the DAGs from S3."""
+        if self.version:
+            raise AirflowException("Refreshing a specific version is not supported")
+
+        with self.lock():
+            self._log.debug(
+                "Downloading DAGs from s3://%s/%s to %s", self.bucket_name, self.prefix, self.s3_dags_dir
+            )
+            self.s3_hook.sync_to_local_dir(
+                bucket_name=self.bucket_name,
+                s3_prefix=self.prefix,
+                local_dir=self.s3_dags_dir,
+                delete_stale=True,
+            )
+
+    def view_url(self, version: str | None = None) -> str | None:
+        """Return a URL for viewing the DAGs in S3. Currently, versioning is not supported."""
+        if self.version:
+            raise AirflowException("S3 url with version is not supported")
+
+        # https://<bucket-name>.s3.<region>.amazonaws.com/<object-key>
+        url = f"https://{self.bucket_name}.s3"
+        if self.s3_hook.region_name:
+            url += f".{self.s3_hook.region_name}"
+        url += ".amazonaws.com"
+        if self.prefix:
+            url += f"/{self.prefix}"
+
+        return url

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,19 +57,28 @@ 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.amazon.version_compat import BaseHook
 from airflow.providers.common.compat.version_compat import AIRFLOW_V_3_0_PLUS
 from airflow.providers_manager import ProvidersManager
 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
+    from airflow.sdk.exceptions import AirflowRuntimeError, ErrorType
 
 if TYPE_CHECKING:
     from aiobotocore.session import AioSession
@@ -607,19 +619,16 @@ class AwsGenericHook(BaseHook, Generic[BaseAwsConnection]):
         """Get the Airflow Connection object and wrap it in helper (cached)."""
         connection = None
         if self.aws_conn_id:
-            possible_exceptions: tuple[type[Exception], ...]
-
-            if AIRFLOW_V_3_0_PLUS:
-                possible_exceptions = (AirflowNotFoundException, AirflowRuntimeError)
-            else:
-                possible_exceptions = (AirflowNotFoundException,)
-
             try:
                 connection = self.get_connection(self.aws_conn_id)
-            except possible_exceptions as e:
-                if isinstance(
-                    e, AirflowNotFoundException
-                ) or f"Connection with ID {self.aws_conn_id} not found" in str(e):
+            except Exception as e:
+                not_found_exc_via_core = isinstance(e, AirflowNotFoundException)
+                not_found_exc_via_task_sdk = (
+                    AIRFLOW_V_3_0_PLUS
+                    and isinstance(e, AirflowRuntimeError)
+                    and e.error.error == ErrorType.CONNECTION_NOT_FOUND
+                )
+                if not_found_exc_via_core or not_found_exc_via_task_sdk:
                     self.log.warning(
                         "Unable to find AWS Connection ID '%s', switching to empty.", self.aws_conn_id
                     )
@@ -627,7 +636,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 +1050,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