apache-airflow-providers-amazon 9.5.0rc2__py3-none-any.whl → 9.6.0__py3-none-any.whl

airflow/providers/amazon/__init__.py

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

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

@@ -27,15 +27,6 @@ from fastapi import FastAPI
 
 from airflow.api_fastapi.app import AUTH_MANAGER_FASTAPI_APP_PREFIX
 from airflow.api_fastapi.auth.managers.base_auth_manager import BaseAuthManager
-from airflow.api_fastapi.auth.managers.models.resource_details import (
-    AccessView,
-    BackfillDetails,
-    ConnectionDetails,
-    DagAccessEntity,
-    DagDetails,
-    PoolDetails,
-    VariableDetails,
-)
 from airflow.cli.cli_config import CLICommand, DefaultHelpParser, GroupCommand
 from airflow.configuration import conf
 from airflow.exceptions import AirflowOptionalProviderFeatureException
@@ -59,9 +50,16 @@ if TYPE_CHECKING:
         IsAuthorizedVariableRequest,
     )
     from airflow.api_fastapi.auth.managers.models.resource_details import (
+        AccessView,
         AssetAliasDetails,
         AssetDetails,
+        BackfillDetails,
         ConfigurationDetails,
+        ConnectionDetails,
+        DagAccessEntity,
+        DagDetails,
+        PoolDetails,
+        VariableDetails,
     )
     from airflow.api_fastapi.common.types import MenuItem
 
@@ -89,7 +87,7 @@ class AwsAuthManager(BaseAuthManager[AwsAuthManagerUser]):
 
     @cached_property
     def apiserver_endpoint(self) -> str:
-        return conf.get("api", "base_url")
+        return conf.get("api", "base_url", fallback="/")
 
     def deserialize_user(self, token: dict[str, Any]) -> AwsAuthManagerUser:
         return AwsAuthManagerUser(user_id=token.pop("sub"), **token)
@@ -256,7 +254,7 @@ class AwsAuthManager(BaseAuthManager[AwsAuthManagerUser]):
             {
                 "method": request["method"],
                 "entity_type": AvpEntities.CONNECTION,
-                "entity_id": cast(ConnectionDetails, request["details"]).conn_id
+                "entity_id": cast("ConnectionDetails", request["details"]).conn_id
                 if request.get("details")
                 else None,
             }
@@ -274,10 +272,10 @@ class AwsAuthManager(BaseAuthManager[AwsAuthManagerUser]):
             {
                 "method": request["method"],
                 "entity_type": AvpEntities.DAG,
-                "entity_id": cast(DagDetails, request["details"]).id if request.get("details") else None,
+                "entity_id": cast("DagDetails", request["details"]).id if request.get("details") else None,
                 "context": {
                     "dag_entity": {
-                        "string": cast(DagAccessEntity, request["access_entity"]).value,
+                        "string": cast("DagAccessEntity", request["access_entity"]).value,
                     },
                 }
                 if request.get("access_entity")
@@ -297,7 +295,7 @@ class AwsAuthManager(BaseAuthManager[AwsAuthManagerUser]):
             {
                 "method": request["method"],
                 "entity_type": AvpEntities.POOL,
-                "entity_id": cast(PoolDetails, request["details"]).name if request.get("details") else None,
+                "entity_id": cast("PoolDetails", request["details"]).name if request.get("details") else None,
             }
             for request in requests
         ]
@@ -313,7 +311,7 @@ class AwsAuthManager(BaseAuthManager[AwsAuthManagerUser]):
             {
                 "method": request["method"],
                 "entity_type": AvpEntities.VARIABLE,
-                "entity_id": cast(VariableDetails, request["details"]).key
+                "entity_id": cast("VariableDetails", request["details"]).key
                 if request.get("details")
                 else None,
             }

airflow/providers/amazon/aws/auth_manager/router/login.py

@@ -80,10 +80,12 @@ def login_callback(request: Request):
         username=saml_auth.get_nameid(),
         email=attributes["email"][0] if "email" in attributes else None,
     )
-    url = conf.get("api", "base_url")
+    url = conf.get("api", "base_url", fallback="/")
     token = get_auth_manager().generate_jwt(user)
     response = RedirectResponse(url=url, status_code=303)
-    response.set_cookie(COOKIE_NAME_JWT_TOKEN, token, secure=True)
+
+    secure = conf.has_option("api", "ssl_cert")
+    response.set_cookie(COOKIE_NAME_JWT_TOKEN, token, secure=secure)
     return response
 
 

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

@@ -49,12 +49,16 @@ from airflow.providers.amazon.aws.executors.utils.exponential_backoff_retry impo
     exponential_backoff_retry,
 )
 from airflow.providers.amazon.aws.hooks.ecs import EcsHook
+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
 from airflow.utils.state import State
 
 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.ecs.utils import (
         CommandType,
@@ -100,6 +104,11 @@ class AwsEcsExecutor(BaseExecutor):
     # AWS limits the maximum number of ARNs in the describe_tasks function.
     DESCRIBE_TASKS_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: EcsTaskCollection = EcsTaskCollection()
@@ -114,6 +123,31 @@ class AwsEcsExecutor(BaseExecutor):
 
         self.run_task_kwargs = self._load_run_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: list[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 start(self):
         """Call this when the Executor is run for the first time by the scheduler."""
         check_health = conf.getboolean(
@@ -278,7 +312,7 @@ class AwsEcsExecutor(BaseExecutor):
         if not has_exit_codes:
             return ""
         reasons = [
-            f'{container["container_arn"]} - {container["reason"]}'
+            f"{container['container_arn']} - {container['reason']}"
             for container in containers
             if "reason" in container
         ]
@@ -462,6 +496,24 @@ class AwsEcsExecutor(BaseExecutor):
         """Save the task to be executed in the next sync by inserting the commands into a queue."""
         if executor_config and ("name" in executor_config or "command" in executor_config):
             raise ValueError('Executor Config should never override "name" or "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"EcsExecutor doesn't know how to handle workload of type: {type(command[0])}"
+                )
+
         self.pending_tasks.append(
             EcsQueuedTask(key, command, queue, executor_config or {}, 1, timezone.utcnow())
         )

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

@@ -146,10 +146,10 @@ class AthenaSQLHook(AwsBaseHook, DbApiHook):
         creds = self.get_credentials(region_name=conn_params["region_name"])
 
         return URL.create(
-            f'awsathena+{conn_params["driver"]}',
+            f"awsathena+{conn_params['driver']}",
             username=creds.access_key,
             password=creds.secret_key,
-            host=f'athena.{conn_params["region_name"]}.{conn_params["aws_domain"]}',
+            host=f"athena.{conn_params['region_name']}.{conn_params['aws_domain']}",
             port=443,
             database=conn_params["schema_name"],
             query={"aws_session_token": creds.token, **self.conn.extra_dejson},

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

@@ -416,8 +416,7 @@ class BatchClientHook(AwsBaseHook):
                 )
         else:
             raise AirflowException(
-                f"AWS Batch job ({job_id}) description error: exceeded status_retries "
-                f"({self.status_retries})"
+                f"AWS Batch job ({job_id}) description error: exceeded status_retries ({self.status_retries})"
             )
 
     @staticmethod

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

@@ -30,7 +30,7 @@ import json
 import sys
 from copy import deepcopy
 from pathlib import Path
-from typing import TYPE_CHECKING, Callable
+from typing import TYPE_CHECKING, Any, Callable
 
 import botocore.client
 import botocore.exceptions
@@ -144,7 +144,12 @@ class BatchWaitersHook(BatchClientHook):
         return self._waiter_model
 
     def get_waiter(
-        self, waiter_name: str, _: dict[str, str] | None = None, deferrable: bool = False, client=None
+        self,
+        waiter_name: str,
+        parameters: dict[str, str] | None = None,
+        config_overrides: dict[str, Any] | None = None,
+        deferrable: bool = False,
+        client=None,
     ) -> botocore.waiter.Waiter:
         """
         Get an AWS Batch service waiter, using the configured ``.waiter_model``.
@@ -175,7 +180,10 @@ class BatchWaitersHook(BatchClientHook):
             the name (including the casing) of the key name in the waiter
             model file (typically this is CamelCasing); see ``.list_waiters``.
 
-        :param _: unused, just here to match the method signature in base_aws
+        :param parameters: unused, just here to match the method signature in base_aws
+        :param config_overrides: unused, just here to match the method signature in base_aws
+        :param deferrable: unused, just here to match the method signature in base_aws
+        :param client: unused, just here to match the method signature in base_aws
 
         :return: a waiter object for the named AWS Batch service
         """

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

@@ -292,7 +292,9 @@ class DmsHook(AwsBaseHook):
             return arn
 
         except ClientError as err:
-            err_str = f"Error: {err.get('Error','').get('Code','')}: {err.get('Error','').get('Message','')}"
+            err_str = (
+                f"Error: {err.get('Error', '').get('Code', '')}: {err.get('Error', '').get('Message', '')}"
+            )
             self.log.error("Error while creating replication config: %s", err_str)
             raise err
 

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

@@ -19,12 +19,13 @@ from __future__ import annotations
 
 import asyncio
 import time
+import warnings
 from functools import cached_property
 from typing import Any
 
 from botocore.exceptions import ClientError
 
-from airflow.exceptions import AirflowException
+from airflow.exceptions import AirflowException, AirflowProviderDeprecationWarning
 from airflow.providers.amazon.aws.hooks.base_aws import AwsBaseHook
 from airflow.providers.amazon.aws.hooks.logs import AwsLogsHook
 
@@ -145,7 +146,7 @@ class GlueJobHook(AwsBaseHook):
 
         return config
 
-    def list_jobs(self) -> list:
+    def describe_jobs(self) -> list:
         """
         Get list of Jobs.
 
@@ -154,6 +155,20 @@ class GlueJobHook(AwsBaseHook):
         """
         return self.conn.get_jobs()
 
+    def list_jobs(self) -> list:
+        """
+        Get list of Jobs.
+
+        .. deprecated::
+            - Use :meth:`describe_jobs` instead.
+        """
+        warnings.warn(
+            "The method `list_jobs` is deprecated. Use the method `describe_jobs` instead.",
+            AirflowProviderDeprecationWarning,
+            stacklevel=2,
+        )
+        return self.describe_jobs()
+
     def get_iam_execution_role(self) -> dict:
         try:
             iam_client = self.get_session(region_name=self.region_name).client(

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

@@ -26,7 +26,7 @@ from airflow.providers.amazon.aws.hooks.base_aws import AwsBaseHook
 
 class MwaaHook(AwsBaseHook):
     """
-    Interact with AWS Manager Workflows for Apache Airflow.
+    Interact with AWS Managed Workflows for Apache Airflow.
 
     Provide thin wrapper around :external+boto3:py:class:`boto3.client("mwaa") <MWAA.Client>`
 

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

@@ -67,7 +67,7 @@ class RedshiftHook(AwsBaseHook):
             for the cluster that is being created.
         :param params: Remaining AWS Create cluster API params.
         """
-        response = self.get_conn().create_cluster(
+        response = self.conn.create_cluster(
             ClusterIdentifier=cluster_identifier,
             NodeType=node_type,
             MasterUsername=master_username,
@@ -87,9 +87,9 @@ class RedshiftHook(AwsBaseHook):
         :param cluster_identifier: unique identifier of a cluster
         """
         try:
-            response = self.get_conn().describe_clusters(ClusterIdentifier=cluster_identifier)["Clusters"]
+            response = self.conn.describe_clusters(ClusterIdentifier=cluster_identifier)["Clusters"]
             return response[0]["ClusterStatus"] if response else None
-        except self.get_conn().exceptions.ClusterNotFoundFault:
+        except self.conn.exceptions.ClusterNotFoundFault:
             return "cluster_not_found"
 
     async def cluster_status_async(self, cluster_identifier: str) -> str:
@@ -115,7 +115,7 @@ class RedshiftHook(AwsBaseHook):
         """
         final_cluster_snapshot_identifier = final_cluster_snapshot_identifier or ""
 
-        response = self.get_conn().delete_cluster(
+        response = self.conn.delete_cluster(
             ClusterIdentifier=cluster_identifier,
             SkipFinalClusterSnapshot=skip_final_cluster_snapshot,
             FinalClusterSnapshotIdentifier=final_cluster_snapshot_identifier,
@@ -131,7 +131,7 @@ class RedshiftHook(AwsBaseHook):
 
         :param cluster_identifier: unique identifier of a cluster
         """
-        response = self.get_conn().describe_cluster_snapshots(ClusterIdentifier=cluster_identifier)
+        response = self.conn.describe_cluster_snapshots(ClusterIdentifier=cluster_identifier)
         if "Snapshots" not in response:
             return None
         snapshots = response["Snapshots"]
@@ -149,7 +149,7 @@ class RedshiftHook(AwsBaseHook):
         :param cluster_identifier: unique identifier of a cluster
         :param snapshot_identifier: unique identifier for a snapshot of a cluster
         """
-        response = self.get_conn().restore_from_cluster_snapshot(
+        response = self.conn.restore_from_cluster_snapshot(
             ClusterIdentifier=cluster_identifier, SnapshotIdentifier=snapshot_identifier
         )
         return response["Cluster"] if response["Cluster"] else None
@@ -175,7 +175,7 @@ class RedshiftHook(AwsBaseHook):
         """
         if tags is None:
             tags = []
-        response = self.get_conn().create_cluster_snapshot(
+        response = self.conn.create_cluster_snapshot(
             SnapshotIdentifier=snapshot_identifier,
             ClusterIdentifier=cluster_identifier,
             ManualSnapshotRetentionPeriod=retention_period,
@@ -192,11 +192,11 @@ class RedshiftHook(AwsBaseHook):
         :param snapshot_identifier: A unique identifier for the snapshot that you are requesting
         """
         try:
-            response = self.get_conn().describe_cluster_snapshots(
+            response = self.conn.describe_cluster_snapshots(
                 SnapshotIdentifier=snapshot_identifier,
             )
             snapshot = response.get("Snapshots")[0]
             snapshot_status: str = snapshot.get("Status")
             return snapshot_status
-        except self.get_conn().exceptions.ClusterSnapshotNotFoundFault:
+        except self.conn.exceptions.ClusterSnapshotNotFoundFault:
             return None

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

@@ -186,8 +186,7 @@ class RedshiftDataHook(AwsGenericHook["RedshiftDataAPIServiceClient"]):
                 RedshiftDataQueryFailedError if status == FAILED_STATE else RedshiftDataQueryAbortedError
             )
             raise exception_cls(
-                f"Statement {resp['Id']} terminated with status {status}. "
-                f"Response details: {pformat(resp)}"
+                f"Statement {resp['Id']} terminated with status {status}. Response details: {pformat(resp)}"
             )
 
         self.log.info("Query status: %s", status)

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

@@ -790,10 +790,6 @@ class S3Hook(AwsBaseHook):
                 "FAILURE: Inactivity Period passed, not enough objects found in %s",
                 path,
             )
-            return {
-                "status": "error",
-                "message": f"FAILURE: Inactivity Period passed, not enough objects found in {path}",
-            }
         return {
             "status": "pending",
             "previous_objects": previous_objects,

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

@@ -131,7 +131,7 @@ def secondary_training_status_message(
     status_strs = []
     for transition in transitions_to_print:
         message = transition["StatusMessage"]
-        time_utc = timezone.convert_to_utc(cast(datetime, job_description["LastModifiedTime"]))
+        time_utc = timezone.convert_to_utc(cast("datetime", job_description["LastModifiedTime"]))
         status_strs.append(f"{time_utc:%Y-%m-%d %H:%M:%S} {transition['Status']} - {message}")
 
     return "\n".join(status_strs)

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

@@ -25,6 +25,5 @@ class AthenaQueryResultsLink(BaseAwsLink):
     name = "Query Results"
     key = "_athena_query_results"
     format_str = (
-        BASE_AWS_CONSOLE_LINK + "/athena/home?region={region_name}#"
-        "/query-editor/history/{query_execution_id}"
+        BASE_AWS_CONSOLE_LINK + "/athena/home?region={region_name}#/query-editor/history/{query_execution_id}"
     )