apache-airflow-providers-amazon 8.17.0rc2__py3-none-any.whl → 8.18.0rc2__py3-none-any.whl

airflow/providers/amazon/__init__.py

@@ -27,7 +27,7 @@ import packaging.version
 
 __all__ = ["__version__"]
 
-__version__ = "8.17.0"
+__version__ = "8.18.0"
 
 try:
     from airflow import __version__ as airflow_version

airflow/providers/amazon/aws/auth_manager/cli/avp_commands.py

@@ -55,7 +55,7 @@ def init_avp(args):
     if not is_new_policy_store:
         print(
             f"Since an existing policy store with description '{args.policy_store_description}' has been found in Amazon Verified Permissions, "
-            "the CLI nade no changes to this policy store for security reasons. "
+            "the CLI made no changes to this policy store for security reasons. "
             "Any modification to this policy store must be done manually.",
         )
     else:
@@ -80,7 +80,7 @@ def update_schema(args):
 
 
 def _get_client():
-    """Returns Amazon Verified Permissions client."""
+    """Return Amazon Verified Permissions client."""
     region_name = conf.get(CONF_SECTION_NAME, CONF_REGION_NAME_KEY)
     return boto3.client("verifiedpermissions", region_name=region_name)
 
@@ -115,7 +115,7 @@ def _create_policy_store(client: BaseClient, args) -> tuple[str | None, bool]:
         print(f"No policy store with description '{args.policy_store_description}' found, creating one.")
         if args.dry_run:
             print(
-                "Dry run, not creating the policy store with description '{args.policy_store_description}'."
+                f"Dry run, not creating the policy store with description '{args.policy_store_description}'."
             )
             return None, True
 

airflow/providers/amazon/aws/auth_manager/cli/definition.py

@@ -35,6 +35,14 @@ ARG_DRY_RUN = Arg(
     action="store_true",
 )
 
+# AWS IAM Identity Center
+ARG_INSTANCE_NAME = Arg(("--instance-name",), help="Instance name in Identity Center", default="Airflow")
+
+ARG_APPLICATION_NAME = Arg(
+    ("--application-name",), help="Application name in Identity Center", default="Airflow"
+)
+
+
 # Amazon Verified Permissions
 ARG_POLICY_STORE_DESCRIPTION = Arg(
     ("--policy-store-description",), help="Policy store description", default="Airflow"
@@ -47,6 +55,12 @@ ARG_POLICY_STORE_ID = Arg(("--policy-store-id",), help="Policy store ID")
 ################
 
 AWS_AUTH_MANAGER_COMMANDS = (
+    ActionCommand(
+        name="init-identity-center",
+        help="Initialize AWS IAM identity Center resources to be used by AWS manager",
+        func=lazy_load_command("airflow.providers.amazon.aws.auth_manager.cli.idc_commands.init_idc"),
+        args=(ARG_INSTANCE_NAME, ARG_APPLICATION_NAME, ARG_DRY_RUN, ARG_VERBOSE),
+    ),
     ActionCommand(
         name="init-avp",
         help="Initialize Amazon Verified resources to be used by AWS manager",

airflow/providers/amazon/aws/auth_manager/cli/idc_commands.py

@@ -0,0 +1,148 @@
+# 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.
+"""User sub-commands."""
+from __future__ import annotations
+
+import logging
+from typing import TYPE_CHECKING
+
+import boto3
+from botocore.exceptions import ClientError
+
+from airflow.configuration import conf
+from airflow.exceptions import AirflowOptionalProviderFeatureException
+from airflow.providers.amazon.aws.auth_manager.constants import CONF_REGION_NAME_KEY, CONF_SECTION_NAME
+from airflow.utils import cli as cli_utils
+
+try:
+    from airflow.utils.providers_configuration_loader import providers_configuration_loaded
+except ImportError:
+    raise AirflowOptionalProviderFeatureException(
+        "Failed to import avp_commands. This feature is only available in Airflow "
+        "version >= 2.8.0 where Auth Managers are introduced."
+    )
+
+if TYPE_CHECKING:
+    from botocore.client import BaseClient
+
+log = logging.getLogger(__name__)
+
+
+@cli_utils.action_cli
+@providers_configuration_loaded
+def init_idc(args):
+    """Initialize AWS IAM Identity Center resources."""
+    client = _get_client()
+
+    # Create the instance if needed
+    instance_arn = _create_instance(client, args)
+
+    # Create the application if needed
+    _create_application(client, instance_arn, args)
+
+    if not args.dry_run:
+        print("AWS IAM Identity Center resources created successfully.")
+
+
+def _get_client():
+    """Return AWS IAM Identity Center client."""
+    region_name = conf.get(CONF_SECTION_NAME, CONF_REGION_NAME_KEY)
+    return boto3.client("sso-admin", region_name=region_name)
+
+
+def _create_instance(client: BaseClient, args) -> str | None:
+    """Create if needed AWS IAM Identity Center instance."""
+    instances = client.list_instances()
+
+    if args.verbose:
+        log.debug("Instances found: %s", instances)
+
+    if len(instances["Instances"]) > 0:
+        print(
+            f"There is already an instance configured in AWS IAM Identity Center: '{instances['Instances'][0]['InstanceArn']}'. "
+            "No need to create a new one."
+        )
+        return instances["Instances"][0]["InstanceArn"]
+    else:
+        print("No instance configured in AWS IAM Identity Center, creating one.")
+        if args.dry_run:
+            print("Dry run, not creating the instance.")
+            return None
+
+        response = client.create_instance(Name=args.instance_name)
+        if args.verbose:
+            log.debug("Response from create_instance: %s", response)
+
+        print(f"Instance created: '{response['InstanceArn']}'")
+
+        return response["InstanceArn"]
+
+
+def _create_application(client: BaseClient, instance_arn: str | None, args) -> str | None:
+    """Create if needed AWS IAM identity Center application."""
+    paginator = client.get_paginator("list_applications")
+    pages = paginator.paginate(InstanceArn=instance_arn or "")
+    applications = [application for page in pages for application in page["Applications"]]
+    existing_applications = [
+        application for application in applications if application["Name"] == args.application_name
+    ]
+
+    if args.verbose:
+        log.debug("Applications found: %s", applications)
+        log.debug("Existing applications found: %s", existing_applications)
+
+    if len(existing_applications) > 0:
+        print(
+            f"There is already an application named '{args.application_name}' in AWS IAM Identity Center: '{existing_applications[0]['ApplicationArn']}'. "
+            "Using this application."
+        )
+        return existing_applications[0]["ApplicationArn"]
+    else:
+        print(f"No application named {args.application_name} found, creating one.")
+        if args.dry_run:
+            print("Dry run, not creating the application.")
+            return None
+
+        try:
+            response = client.create_application(
+                ApplicationProviderArn="arn:aws:sso::aws:applicationProvider/custom-saml",
+                Description="Application automatically created through the Airflow CLI. This application is used to access Airflow environment.",
+                InstanceArn=instance_arn,
+                Name=args.application_name,
+                PortalOptions={
+                    "SignInOptions": {
+                        "Origin": "IDENTITY_CENTER",
+                    },
+                    "Visibility": "ENABLED",
+                },
+                Status="ENABLED",
+            )
+            if args.verbose:
+                log.debug("Response from create_application: %s", response)
+        except ClientError as e:
+            # This is needed because as of today, the create_application in AWS Identity Center does not support SAML application
+            # Remove this part when it is supported
+            if "is not supported for this action" in e.response["Error"]["Message"]:
+                print(
+                    "Creation of SAML applications is only supported in AWS console today. "
+                    "Please create the application through the console."
+                )
+            raise
+
+        print(f"Application created: '{response['ApplicationArn']}'")
+
+        return response["ApplicationArn"]

airflow/providers/amazon/aws/auth_manager/views/auth.py

@@ -71,7 +71,7 @@ class AwsAuthManagerAuthenticationViews(AirflowBaseView):
     @expose("/login_callback", methods=("GET", "POST"))
     def login_callback(self):
         """
-        Callback where the user is redirected to after successful login.
+        Redirect the user to this callback after successful login.
 
         CSRF protection needs to be disabled otherwise the callback won't work.
         """

airflow/providers/amazon/aws/executors/ecs/Dockerfile

@@ -37,7 +37,7 @@ COPY <<"EOF" /install_dags_entrypoint.sh
 #!/bin/bash
 
 echo "Downloading DAGs from S3 bucket"
-aws s3 sync "$S3_URL" "$CONTAINER_DAG_PATH"
+aws s3 sync "$S3_URI" "$CONTAINER_DAG_PATH"
 
 /home/airflow/entrypoint.sh "$@"
 EOF
@@ -98,8 +98,8 @@ ENV CONTAINER_DAG_PATH=$container_dag_path
 
 
 # Use these arguments to load DAGs onto the container from S3
-ARG s3_url
-ENV S3_URL=$s3_url
+ARG s3_uri
+ENV S3_URI=$s3_uri
 # If using S3 bucket as source of DAGs, uncommenting the next ENTRYPOINT command will overwrite this one.
 ENTRYPOINT ["/usr/bin/dumb-init", "--", "/home/airflow/entrypoint.sh"]
 

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

@@ -61,7 +61,7 @@ class BotoTaskSchema(Schema):
 
     @post_load
     def make_task(self, data, **kwargs):
-        """Overwrites marshmallow load() to return an instance of EcsExecutorTask instead of a dictionary."""
+        """Overwrite marshmallow load() to return an EcsExecutorTask instance instead of a dictionary."""
         # Imported here to avoid circular import.
         from airflow.providers.amazon.aws.executors.ecs.utils import EcsExecutorTask
 

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

@@ -42,9 +42,13 @@ from airflow.providers.amazon.aws.executors.ecs.utils import (
     EcsQueuedTask,
     EcsTaskCollection,
 )
-from airflow.providers.amazon.aws.executors.utils.exponential_backoff_retry import exponential_backoff_retry
+from airflow.providers.amazon.aws.executors.utils.exponential_backoff_retry import (
+    calculate_next_attempt_delay,
+    exponential_backoff_retry,
+)
 from airflow.providers.amazon.aws.hooks.ecs import EcsHook
 from airflow.utils import timezone
+from airflow.utils.helpers import merge_dicts
 from airflow.utils.state import State
 
 if TYPE_CHECKING:
@@ -108,7 +112,7 @@ class AwsEcsExecutor(BaseExecutor):
         self.run_task_kwargs = self._load_run_kwargs()
 
     def start(self):
-        """This is called by the scheduler when the Executor is being run for the first time."""
+        """Call this when the Executor is run for the first time by the scheduler."""
         check_health = conf.getboolean(
             CONFIG_GROUP_NAME, AllEcsConfigKeys.CHECK_HEALTH_ON_STARTUP, fallback=False
         )
@@ -213,7 +217,7 @@ class AwsEcsExecutor(BaseExecutor):
             self.log.exception("Failed to sync %s", self.__class__.__name__)
 
     def sync_running_tasks(self):
-        """Checks and update state on all running tasks."""
+        """Check and update state on all running tasks."""
         all_task_arns = self.active_workers.get_all_arns()
         if not all_task_arns:
             self.log.debug("No active Airflow tasks, skipping sync.")
@@ -300,7 +304,14 @@ class AwsEcsExecutor(BaseExecutor):
             )
             self.active_workers.increment_failure_count(task_key)
             self.pending_tasks.appendleft(
-                EcsQueuedTask(task_key, task_cmd, queue, exec_info, failure_count + 1)
+                EcsQueuedTask(
+                    task_key,
+                    task_cmd,
+                    queue,
+                    exec_info,
+                    failure_count + 1,
+                    timezone.utcnow() + calculate_next_attempt_delay(failure_count),
+                )
             )
         else:
             self.log.error(
@@ -313,7 +324,7 @@ class AwsEcsExecutor(BaseExecutor):
 
     def attempt_task_runs(self):
         """
-        Takes tasks from the pending_tasks queue, and attempts to find an instance to run it on.
+        Take tasks from the pending_tasks queue, and attempts to find an instance to run it on.
 
         If the launch type is EC2, this will attempt to place tasks on empty EC2 instances.  If
             there are no EC2 instances available, no task is placed and this function will be
@@ -331,6 +342,8 @@ class AwsEcsExecutor(BaseExecutor):
             exec_config = ecs_task.executor_config
             attempt_number = ecs_task.attempt_number
             _failure_reasons = []
+            if timezone.utcnow() < ecs_task.next_attempt_time:
+                continue
             try:
                 run_task_response = self._run_task(task_key, cmd, queue, exec_config)
             except NoCredentialsError:
@@ -361,6 +374,9 @@ class AwsEcsExecutor(BaseExecutor):
                 # Make sure the number of attempts does not exceed MAX_RUN_TASK_ATTEMPTS
                 if int(attempt_number) <= int(self.__class__.MAX_RUN_TASK_ATTEMPTS):
                     ecs_task.attempt_number += 1
+                    ecs_task.next_attempt_time = timezone.utcnow() + calculate_next_attempt_delay(
+                        attempt_number
+                    )
                     self.pending_tasks.appendleft(ecs_task)
                 else:
                     self.log.error(
@@ -393,8 +409,8 @@ class AwsEcsExecutor(BaseExecutor):
         The command and executor config will be placed in the container-override
         section of the JSON request before calling Boto3's "run_task" function.
         """
-        run_task_api = self._run_task_kwargs(task_id, cmd, queue, exec_config)
-        boto_run_task = self.ecs.run_task(**run_task_api)
+        run_task_kwargs = self._run_task_kwargs(task_id, cmd, queue, exec_config)
+        boto_run_task = self.ecs.run_task(**run_task_kwargs)
         run_task_response = BotoRunTaskSchema().load(boto_run_task)
         return run_task_response
 
@@ -402,30 +418,32 @@ class AwsEcsExecutor(BaseExecutor):
         self, task_id: TaskInstanceKey, cmd: CommandType, queue: str, exec_config: ExecutorConfigType
     ) -> dict:
         """
-        Overrides the Airflow command to update the container overrides so kwargs are specific to this task.
+        Update the Airflow command by modifying container overrides for task-specific kwargs.
 
         One last chance to modify Boto3's "run_task" kwarg params before it gets passed into the Boto3 client.
         """
-        run_task_api = deepcopy(self.run_task_kwargs)
-        container_override = self.get_container(run_task_api["overrides"]["containerOverrides"])
+        run_task_kwargs = deepcopy(self.run_task_kwargs)
+        run_task_kwargs = merge_dicts(run_task_kwargs, exec_config)
+        container_override = self.get_container(run_task_kwargs["overrides"]["containerOverrides"])
         container_override["command"] = cmd
-        container_override.update(exec_config)
 
         # Inject the env variable to configure logging for containerized execution environment
         if "environment" not in container_override:
             container_override["environment"] = []
         container_override["environment"].append({"name": "AIRFLOW_IS_EXECUTOR_CONTAINER", "value": "true"})
 
-        return run_task_api
+        return run_task_kwargs
 
     def execute_async(self, key: TaskInstanceKey, command: CommandType, queue=None, executor_config=None):
         """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"')
-        self.pending_tasks.append(EcsQueuedTask(key, command, queue, executor_config or {}, 1))
+        self.pending_tasks.append(
+            EcsQueuedTask(key, command, queue, executor_config or {}, 1, timezone.utcnow())
+        )
 
     def end(self, heartbeat_interval=10):
-        """Waits for all currently running tasks to end, and doesn't launch any tasks."""
+        """Wait for all currently running tasks to end, and don't launch any tasks."""
         try:
             while True:
                 self.sync()
@@ -465,8 +483,13 @@ class AwsEcsExecutor(BaseExecutor):
         return ecs_executor_run_task_kwargs
 
     def get_container(self, container_list):
-        """Searches task list for core Airflow container."""
+        """Search task list for core Airflow container."""
         for container in container_list:
-            if container["name"] == self.container_name:
-                return container
+            try:
+                if container["name"] == self.container_name:
+                    return container
+            except KeyError:
+                raise EcsExecutorException(
+                    'container "name" must be provided in "containerOverrides" configuration'
+                )
         raise KeyError(f"No such container found by container name: {self.container_name}")

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

@@ -23,6 +23,7 @@ Data classes and utility functions used by the ECS executor.
 
 from __future__ import annotations
 
+import datetime
 from collections import defaultdict
 from dataclasses import dataclass
 from typing import TYPE_CHECKING, Any, Callable, Dict, List
@@ -58,6 +59,7 @@ class EcsQueuedTask:
     queue: str
     executor_config: ExecutorConfigType
     attempt_number: int
+    next_attempt_time: datetime.datetime
 
 
 @dataclass
@@ -125,9 +127,9 @@ class EcsExecutorTask:
 
     def get_task_state(self) -> str:
         """
-        This is the primary logic that handles state in an ECS task.
+        Determine the state of an ECS task based on its status and other relevant attributes.
 
-        It will determine if a status is:
+        It can return one of the following statuses:
             QUEUED - Task is being provisioned.
             RUNNING - Task is launched on ECS.
             REMOVED - Task provisioning has failed for some reason. See `stopped_reason`.
@@ -171,7 +173,7 @@ class EcsTaskCollection:
         exec_config: ExecutorConfigType,
         attempt_number: int,
     ):
-        """Adds a task to the collection."""
+        """Add a task to the collection."""
         arn = task.task_arn
         self.tasks[arn] = task
         self.key_to_arn[airflow_task_key] = arn
@@ -180,7 +182,7 @@ class EcsTaskCollection:
         self.key_to_failure_counts[airflow_task_key] = attempt_number
 
     def update_task(self, task: EcsExecutorTask):
-        """Updates the state of the given task based on task ARN."""
+        """Update the state of the given task based on task ARN."""
         self.tasks[task.task_arn] = task
 
     def task_by_key(self, task_key: TaskInstanceKey) -> EcsExecutorTask:
@@ -193,7 +195,7 @@ class EcsTaskCollection:
         return self.tasks[arn]
 
     def pop_by_key(self, task_key: TaskInstanceKey) -> EcsExecutorTask:
-        """Deletes task from collection based off of Airflow Task Instance Key."""
+        """Delete task from collection based off of Airflow Task Instance Key."""
         arn = self.key_to_arn[task_key]
         task = self.tasks[arn]
         del self.key_to_arn[task_key]
@@ -225,11 +227,11 @@ class EcsTaskCollection:
         return self.key_to_task_info[task_key]
 
     def __getitem__(self, value):
-        """Gets a task by AWS ARN."""
+        """Get a task by AWS ARN."""
         return self.task_by_arn(value)
 
     def __len__(self):
-        """Determines the number of tasks in collection."""
+        """Determine the number of tasks in collection."""
         return len(self.tasks)
 
 

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

@@ -25,6 +25,21 @@ from airflow.utils import timezone
 log = logging.getLogger(__name__)
 
 
+def calculate_next_attempt_delay(
+    attempt_number: int,
+    max_delay: int = 60 * 2,
+    exponent_base: int = 4,
+) -> timedelta:
+    """
+    Calculate the exponential backoff (in seconds) until the next attempt.
+
+    :param attempt_number: Number of attempts since last success.
+    :param max_delay: Maximum delay in seconds between retries. Default 120.
+    :param exponent_base: Exponent base to calculate delay. Default 4.
+    """
+    return timedelta(seconds=min((exponent_base**attempt_number), max_delay))
+
+
 def exponential_backoff_retry(
     last_attempt_time: datetime,
     attempts_since_last_successful: int,
@@ -34,7 +49,7 @@ def exponential_backoff_retry(
     exponent_base: int = 4,
 ) -> None:
     """
-    Retries a callable function with exponential backoff between attempts if it raises an exception.
+    Retry a callable function with exponential backoff between attempts if it raises an exception.
 
     :param last_attempt_time: Timestamp of last attempt call.
     :param attempts_since_last_successful: Number of attempts since last success.
@@ -47,8 +62,10 @@ def exponential_backoff_retry(
         log.error("Max attempts reached. Exiting.")
         return
 
-    delay = min((exponent_base**attempts_since_last_successful), max_delay)
-    next_retry_time = last_attempt_time + timedelta(seconds=delay)
+    next_retry_time = last_attempt_time + calculate_next_attempt_delay(
+        attempt_number=attempts_since_last_successful, max_delay=max_delay, exponent_base=exponent_base
+    )
+
     current_time = timezone.utcnow()
 
     if current_time >= next_retry_time:
@@ -56,5 +73,7 @@ def exponential_backoff_retry(
             callable_function()
         except Exception:
             log.exception("Error calling %r", callable_function.__name__)
-            next_delay = min((exponent_base ** (attempts_since_last_successful + 1)), max_delay)
+            next_delay = calculate_next_attempt_delay(
+                attempts_since_last_successful + 1, max_delay, exponent_base
+            )
             log.info("Waiting for %s seconds before retrying.", next_delay)

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

@@ -25,7 +25,7 @@ This module contains AWS Athena hook.
 from __future__ import annotations
 
 import warnings
-from typing import TYPE_CHECKING, Any
+from typing import TYPE_CHECKING, Any, Collection
 
 from airflow.exceptions import AirflowException, AirflowProviderDeprecationWarning
 from airflow.providers.amazon.aws.hooks.base_aws import AwsBaseHook
@@ -34,6 +34,19 @@ from airflow.providers.amazon.aws.utils.waiter_with_logging import wait
 if TYPE_CHECKING:
     from botocore.paginate import PageIterator
 
+MULTI_LINE_QUERY_LOG_PREFIX = "\n\t\t"
+
+
+def query_params_to_string(params: dict[str, str | Collection[str]]) -> str:
+    result = ""
+    for key, value in params.items():
+        if key == "QueryString":
+            value = (
+                MULTI_LINE_QUERY_LOG_PREFIX + str(value).replace("\n", MULTI_LINE_QUERY_LOG_PREFIX).rstrip()
+            )
+        result += f"\t{key}: {value}\n"
+    return result.rstrip()
+
 
 class AthenaHook(AwsBaseHook):
     """Interact with Amazon Athena.
@@ -115,7 +128,7 @@ class AthenaHook(AwsBaseHook):
         if client_request_token:
             params["ClientRequestToken"] = client_request_token
         if self.log_query:
-            self.log.info("Running Query with params: %s", params)
+            self.log.info("Running Query with params:\n%s", query_params_to_string(params))
         response = self.get_conn().start_query_execution(**params)
         query_execution_id = response["QueryExecutionId"]
         self.log.info("Query execution id: %s", query_execution_id)

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

@@ -29,7 +29,6 @@ import inspect
 import json
 import logging
 import os
-import warnings
 from copy import deepcopy
 from functools import cached_property, wraps
 from pathlib import Path
@@ -44,6 +43,7 @@ import tenacity
 from botocore.config import Config
 from botocore.waiter import Waiter, WaiterModel
 from dateutil.tz import tzlocal
+from deprecated import deprecated
 from slugify import slugify
 
 from airflow.configuration import conf
@@ -314,7 +314,7 @@ class BaseSessionFactory(LoggingMixin):
             idp_request_retry_kwargs = saml_config["idp_request_retry_kwargs"]
             self.log.info("idp_request_retry_kwargs= %s", idp_request_retry_kwargs)
             from requests.adapters import HTTPAdapter
-            from requests.packages.urllib3.util.retry import Retry
+            from urllib3.util.retry import Retry
 
             retry_strategy = Retry(**idp_request_retry_kwargs)
             adapter = HTTPAdapter(max_retries=retry_strategy)
@@ -1020,6 +1020,13 @@ except ImportError:
     pass
 
 
+@deprecated(
+    reason=(
+        "`airflow.providers.amazon.aws.hook.base_aws.BaseAsyncSessionFactory` "
+        "has been deprecated and will be removed in future"
+    ),
+    category=AirflowProviderDeprecationWarning,
+)
 class BaseAsyncSessionFactory(BaseSessionFactory):
     """
     Base AWS Session Factory class to handle aiobotocore session creation.
@@ -1029,12 +1036,6 @@ class BaseAsyncSessionFactory(BaseSessionFactory):
     """
 
     def __init__(self, *args, **kwargs):
-        warnings.warn(
-            "airflow.providers.amazon.aws.hook.base_aws.BaseAsyncSessionFactory has been deprecated and "
-            "will be removed in future",
-            AirflowProviderDeprecationWarning,
-            stacklevel=2,
-        )
         super().__init__(*args, **kwargs)
 
     async def get_role_credentials(self) -> dict:
@@ -1113,6 +1114,13 @@ class BaseAsyncSessionFactory(BaseSessionFactory):
         return self._get_session_with_assume_role()
 
 
+@deprecated(
+    reason=(
+        "`airflow.providers.amazon.aws.hook.base_aws.AwsBaseAsyncHook` "
+        "has been deprecated and will be removed in future"
+    ),
+    category=AirflowProviderDeprecationWarning,
+)
 class AwsBaseAsyncHook(AwsBaseHook):
     """Interacts with AWS using aiobotocore asynchronously.
 
@@ -1129,12 +1137,6 @@ class AwsBaseAsyncHook(AwsBaseHook):
     """
 
     def __init__(self, *args, **kwargs):
-        warnings.warn(
-            "airflow.providers.amazon.aws.hook.base_aws.AwsBaseAsyncHook has been deprecated and "
-            "will be removed in future",
-            AirflowProviderDeprecationWarning,
-            stacklevel=2,
-        )
         super().__init__(*args, **kwargs)
 
     def get_async_session(self) -> AioSession:

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

@@ -383,6 +383,7 @@ class EmrContainerHook(AwsBaseHook):
         configuration_overrides: dict | None = None,
         client_request_token: str | None = None,
         tags: dict | None = None,
+        retry_max_attempts: int | None = None,
     ) -> str:
         """
         Submit a job to the EMR Containers API and return the job ID.
@@ -402,6 +403,7 @@ class EmrContainerHook(AwsBaseHook):
         :param client_request_token: The client idempotency token of the job run request.
             Use this if you want to specify a unique ID to prevent two jobs from getting started.
         :param tags: The tags assigned to job runs.
+        :param retry_max_attempts: The maximum number of attempts on the job's driver.
         :return: The ID of the job run request.
         """
         params = {
@@ -415,6 +417,10 @@ class EmrContainerHook(AwsBaseHook):
         }
         if client_request_token:
             params["clientToken"] = client_request_token
+        if retry_max_attempts:
+            params["retryPolicyConfiguration"] = {
+                "maxAttempts": retry_max_attempts,
+            }
 
         response = self.conn.start_job_run(**params)