apache-airflow-providers-amazon 7.4.1rc1__py3-none-any.whl → 8.0.0rc2__py3-none-any.whl

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

@@ -24,7 +24,6 @@ This module contains AWS Athena hook.
 """
 from __future__ import annotations
 
-import warnings
 from time import sleep
 from typing import Any
 
@@ -224,7 +223,6 @@ class AthenaHook(AwsBaseHook):
     def poll_query_status(
         self,
         query_execution_id: str,
-        max_tries: int | None = None,
         max_polling_attempts: int | None = None,
     ) -> str | None:
         """
@@ -232,21 +230,8 @@ class AthenaHook(AwsBaseHook):
         Returns one of the final states
 
         :param query_execution_id: Id of submitted athena query
-        :param max_tries: Deprecated - Use max_polling_attempts instead
         :param max_polling_attempts: Number of times to poll for query state before function exits
         """
-        if max_tries:
-            warnings.warn(
-                f"Passing 'max_tries' to {self.__class__.__name__}.poll_query_status is deprecated "
-                f"and will be removed in a future release. Please use 'max_polling_attempts' instead.",
-                DeprecationWarning,
-                stacklevel=2,
-            )
-            if max_polling_attempts and max_polling_attempts != max_tries:
-                raise Exception("max_polling_attempts must be the same value as max_tries")
-            else:
-                max_polling_attempts = max_tries
-
         try_number = 1
         final_query_state = None  # Query state when query reaches final state or max_polling_attempts reached
         while True:

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

@@ -30,7 +30,6 @@ import json
 import logging
 import os
 import uuid
-import warnings
 from copy import deepcopy
 from functools import wraps
 from os import PathLike
@@ -58,7 +57,6 @@ from airflow.exceptions import (
 )
 from airflow.hooks.base import BaseHook
 from airflow.providers.amazon.aws.utils.connection_wrapper import AwsConnectionWrapper
-from airflow.providers.amazon.aws.waiters.base_waiter import BaseBotoWaiter
 from airflow.providers_manager import ProvidersManager
 from airflow.utils.helpers import exactly_one
 from airflow.utils.log.logging_mixin import LoggingMixin
@@ -72,12 +70,15 @@ if TYPE_CHECKING:
 
 class BaseSessionFactory(LoggingMixin):
     """
-    Base AWS Session Factory class to handle boto3 session creation.
+    Base AWS Session Factory class to handle synchronous and async boto session creation.
     It can handle most of the AWS supported authentication methods.
 
     User can also derive from this class to have full control of boto3 session
     creation or to support custom federation.
 
+    Note: Not all features implemented for synchronous sessions are available for async
+    sessions.
+
     .. seealso::
         - :ref:`howto/connection:aws:session-factory`
     """
@@ -127,17 +128,50 @@ class BaseSessionFactory(LoggingMixin):
         """Assume Role ARN from AWS Connection"""
         return self.conn.role_arn
 
-    def create_session(self) -> boto3.session.Session:
-        """Create boto3 Session from connection config."""
+    def _apply_session_kwargs(self, session):
+        if self.conn.session_kwargs.get("profile_name", None) is not None:
+            session.set_config_variable("profile", self.conn.session_kwargs["profile_name"])
+
+        if (
+            self.conn.session_kwargs.get("aws_access_key_id", None)
+            or self.conn.session_kwargs.get("aws_secret_access_key", None)
+            or self.conn.session_kwargs.get("aws_session_token", None)
+        ):
+            session.set_credentials(
+                self.conn.session_kwargs["aws_access_key_id"],
+                self.conn.session_kwargs["aws_secret_access_key"],
+                self.conn.session_kwargs["aws_session_token"],
+            )
+
+        if self.conn.session_kwargs.get("region_name", None) is not None:
+            session.set_config_variable("region", self.conn.session_kwargs["region_name"])
+
+    def get_async_session(self):
+        from aiobotocore.session import get_session as async_get_session
+
+        return async_get_session()
+
+    def create_session(self, deferrable: bool = False) -> boto3.session.Session:
+        """Create boto3 or aiobotocore Session from connection config."""
         if not self.conn:
             self.log.info(
                 "No connection ID provided. Fallback on boto3 credential strategy (region_name=%r). "
                 "See: https://boto3.amazonaws.com/v1/documentation/api/latest/guide/configuration.html",
                 self.region_name,
             )
-            return boto3.session.Session(region_name=self.region_name)
+            if deferrable:
+                session = self.get_async_session()
+                self._apply_session_kwargs(session)
+                return session
+            else:
+                return boto3.session.Session(region_name=self.region_name)
         elif not self.role_arn:
-            return self.basic_session
+            if deferrable:
+                session = self.get_async_session()
+                self._apply_session_kwargs(session)
+                return session
+            else:
+                return self.basic_session
 
         # Values stored in ``AwsConnectionWrapper.session_kwargs`` are intended to be used only
         # to create the initial boto3 session.
@@ -150,12 +184,18 @@ class BaseSessionFactory(LoggingMixin):
         assume_session_kwargs = {}
         if self.conn.region_name:
             assume_session_kwargs["region_name"] = self.conn.region_name
-        return self._create_session_with_assume_role(session_kwargs=assume_session_kwargs)
+        return self._create_session_with_assume_role(
+            session_kwargs=assume_session_kwargs, deferrable=deferrable
+        )
 
     def _create_basic_session(self, session_kwargs: dict[str, Any]) -> boto3.session.Session:
         return boto3.session.Session(**session_kwargs)
 
-    def _create_session_with_assume_role(self, session_kwargs: dict[str, Any]) -> boto3.session.Session:
+    def _create_session_with_assume_role(
+        self, session_kwargs: dict[str, Any], deferrable: bool = False
+    ) -> boto3.session.Session:
+        from aiobotocore.session import get_session as async_get_session
+
         if self.conn.assume_role_method == "assume_role_with_web_identity":
             # Deferred credentials have no initial credentials
             credential_fetcher = self._get_web_identity_credential_fetcher()
@@ -172,10 +212,10 @@ class BaseSessionFactory(LoggingMixin):
                 method="sts-assume-role",
             )
 
-        session = botocore.session.get_session()
+        session = async_get_session() if deferrable else botocore.session.get_session()
+
         session._credentials = credentials
-        region_name = self.basic_session.region_name
-        session.set_config_variable("region", region_name)
+        session.set_config_variable("region", self.basic_session.region_name)
 
         return boto3.session.Session(botocore_session=session, **session_kwargs)
 
@@ -362,34 +402,6 @@ class BaseSessionFactory(LoggingMixin):
     def _strip_invalid_session_name_characters(self, role_session_name: str) -> str:
         return slugify(role_session_name, regex_pattern=r"[^\w+=,.@-]+")
 
-    def _get_region_name(self) -> str | None:
-        warnings.warn(
-            "`BaseSessionFactory._get_region_name` method deprecated and will be removed "
-            "in a future releases. Please use `BaseSessionFactory.region_name` property instead.",
-            DeprecationWarning,
-            stacklevel=2,
-        )
-        return self.region_name
-
-    def _read_role_arn_from_extra_config(self) -> str | None:
-        warnings.warn(
-            "`BaseSessionFactory._read_role_arn_from_extra_config` method deprecated and will be removed "
-            "in a future releases. Please use `BaseSessionFactory.role_arn` property instead.",
-            DeprecationWarning,
-            stacklevel=2,
-        )
-        return self.role_arn
-
-    def _read_credentials_from_connection(self) -> tuple[str | None, str | None]:
-        warnings.warn(
-            "`BaseSessionFactory._read_credentials_from_connection` method deprecated and will be removed "
-            "in a future releases. Please use `BaseSessionFactory.conn.aws_access_key_id` and "
-            "`BaseSessionFactory.aws_secret_access_key` properties instead.",
-            DeprecationWarning,
-            stacklevel=2,
-        )
-        return self.conn.aws_access_key_id, self.conn.aws_secret_access_key
-
 
 class AwsGenericHook(BaseHook, Generic[BaseAwsConnection]):
     """
@@ -531,13 +543,8 @@ class AwsGenericHook(BaseHook, Generic[BaseAwsConnection]):
             try:
                 connection = self.get_connection(self.aws_conn_id)
             except AirflowNotFoundException:
-                warnings.warn(
-                    f"Unable to find AWS Connection ID '{self.aws_conn_id}', switching to empty. "
-                    "This behaviour is deprecated and will be removed in a future releases. "
-                    "Please provide existed AWS connection ID or if required boto3 credential strategy "
-                    "explicit set AWS Connection ID to None.",
-                    DeprecationWarning,
-                    stacklevel=2,
+                self.log.warning(
+                    "Unable to find AWS Connection ID '%s', switching to empty.", self.aws_conn_id
                 )
 
         return AwsConnectionWrapper(
@@ -564,11 +571,11 @@ class AwsGenericHook(BaseHook, Generic[BaseAwsConnection]):
         """Verify or not SSL certificates boto3 client/resource read-only property."""
         return self.conn_config.verify
 
-    def get_session(self, region_name: str | None = None) -> boto3.session.Session:
+    def get_session(self, region_name: str | None = None, deferrable: bool = False) -> boto3.session.Session:
         """Get the underlying boto3.session.Session(region_name=region_name)."""
         return SessionFactory(
             conn=self.conn_config, region_name=region_name, config=self.config
-        ).create_session()
+        ).create_session(deferrable=deferrable)
 
     def _get_config(self, config: Config | None = None) -> Config:
         """
@@ -591,10 +598,19 @@ class AwsGenericHook(BaseHook, Generic[BaseAwsConnection]):
         self,
         region_name: str | None = None,
         config: Config | None = None,
+        deferrable: bool = False,
     ) -> boto3.client:
         """Get the underlying boto3 client using boto3 session"""
         client_type = self.client_type
-        session = self.get_session(region_name=region_name)
+        session = self.get_session(region_name=region_name, deferrable=deferrable)
+        if not isinstance(session, boto3.session.Session):
+            return session.create_client(
+                client_type,
+                endpoint_url=self.conn_config.endpoint_url,
+                config=self._get_config(config),
+                verify=self.verify,
+            )
+
         return session.client(
             client_type,
             endpoint_url=self.conn_config.endpoint_url,
@@ -634,6 +650,14 @@ class AwsGenericHook(BaseHook, Generic[BaseAwsConnection]):
         else:
             return self.get_resource_type(region_name=self.region_name)
 
+    @property
+    def async_conn(self):
+        """Get an aiobotocore client to use for async operations."""
+        if not self.client_type:
+            raise ValueError("client_type must be specified.")
+
+        return self.get_client_type(region_name=self.region_name, deferrable=True)
+
     @cached_property
     def conn_client_meta(self) -> ClientMeta:
         """Get botocore client metadata from Hook connection (cached)."""
@@ -730,17 +754,6 @@ class AwsGenericHook(BaseHook, Generic[BaseAwsConnection]):
 
         return retry_decorator
 
-    def _get_credentials(self, region_name: str | None) -> tuple[boto3.session.Session, str | None]:
-        warnings.warn(
-            "`AwsGenericHook._get_credentials` method deprecated and will be removed in a future releases. "
-            "Please use `AwsGenericHook.get_session` method and "
-            "`AwsGenericHook.conn_config.endpoint_url` property instead.",
-            DeprecationWarning,
-            stacklevel=2,
-        )
-
-        return self.get_session(region_name=region_name), self.conn_config.endpoint_url
-
     @staticmethod
     def get_ui_field_behaviour() -> dict[str, Any]:
         """Returns custom UI field behaviour for AWS Connection."""
@@ -794,21 +807,39 @@ class AwsGenericHook(BaseHook, Generic[BaseAwsConnection]):
 
     @cached_property
     def waiter_path(self) -> PathLike[str] | None:
-        path = Path(__file__).parents[1].joinpath(f"waiters/{self.client_type}.json").resolve()
+        filename = self.client_type if self.client_type else self.resource_type
+        path = Path(__file__).parents[1].joinpath(f"waiters/{filename}.json").resolve()
         return path if path.exists() else None
 
-    def get_waiter(self, waiter_name: str, parameters: dict[str, str] | None = None) -> Waiter:
+    def get_waiter(
+        self,
+        waiter_name: str,
+        parameters: dict[str, str] | None = None,
+        deferrable: bool = False,
+        client=None,
+    ) -> Waiter:
         """
         First checks if there is a custom waiter with the provided waiter_name and
         uses that if it exists, otherwise it will check the service client for a
         waiter that matches the name and pass that through.
 
+        If `deferrable` is True, the waiter will be an AIOWaiter, generated from the
+        client that is passed as a parameter. If `deferrable` is True, `client` must be
+        provided.
+
         :param waiter_name: The name of the waiter.  The name should exactly match the
             name of the key in the waiter model file (typically this is CamelCase).
         :param parameters: will scan the waiter config for the keys of that dict, and replace them with the
             corresponding value. If a custom waiter has such keys to be expanded, they need to be provided
             here.
+        :param deferrable: If True, the waiter is going to be an async custom waiter.
+
         """
+        from airflow.providers.amazon.aws.waiters.base_waiter import BaseBotoWaiter
+
+        if deferrable and not client:
+            raise ValueError("client must be provided for a deferrable waiter.")
+        client = client or self.conn
         if self.waiter_path and (waiter_name in self._list_custom_waiters()):
             # Technically if waiter_name is in custom_waiters then self.waiter_path must
             # exist but MyPy doesn't like the fact that self.waiter_path could be None.
@@ -816,7 +847,9 @@ class AwsGenericHook(BaseHook, Generic[BaseAwsConnection]):
                 config = json.loads(config_file.read())
 
             config = self._apply_parameters_value(config, waiter_name, parameters)
-            return BaseBotoWaiter(client=self.conn, model_config=config).waiter(waiter_name)
+            return BaseBotoWaiter(client=client, model_config=config, deferrable=deferrable).waiter(
+                waiter_name
+            )
         # If there is no custom waiter found for the provided name,
         # then try checking the service's official waiters.
         return self.conn.get_waiter(waiter_name)
@@ -985,7 +1018,7 @@ class BaseAsyncSessionFactory(BaseSessionFactory):
             aio_session.set_config_variable("region", region_name)
         return aio_session
 
-    def create_session(self) -> AioSession:
+    def create_session(self, deferrable: bool = False) -> AioSession:
         """Create aiobotocore Session from connection and config."""
         if not self._conn:
             self.log.info("No connection ID provided. Fallback on boto3 credential strategy")

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

@@ -414,43 +414,76 @@ class BatchClientHook(AwsBaseHook):
         return matching_jobs[0]
 
     def get_job_awslogs_info(self, job_id: str) -> dict[str, str] | None:
+        all_info = self.get_job_all_awslogs_info(job_id)
+        if not all_info:
+            return None
+        if len(all_info) > 1:
+            self.log.warning(
+                f"AWS Batch job ({job_id}) has more than one log stream, " f"only returning the first one."
+            )
+        return all_info[0]
+
+    def get_job_all_awslogs_info(self, job_id: str) -> list[dict[str, str]]:
         """
         Parse job description to extract AWS CloudWatch information.
 
         :param job_id: AWS Batch Job ID
         """
-        job_container_desc = self.get_job_description(job_id=job_id).get("container", {})
-        log_configuration = job_container_desc.get("logConfiguration", {})
-
-        # In case if user select other "logDriver" rather than "awslogs"
-        # than CloudWatch logging should be disabled.
-        # If user not specify anything than expected that "awslogs" will use
-        # with default settings:
-        #   awslogs-group = /aws/batch/job
-        #   awslogs-region = `same as AWS Batch Job region`
-        log_driver = log_configuration.get("logDriver", "awslogs")
-        if log_driver != "awslogs":
+        job_desc = self.get_job_description(job_id=job_id)
+
+        job_node_properties = job_desc.get("nodeProperties", {})
+        job_container_desc = job_desc.get("container", {})
+
+        if job_node_properties:
+            # one log config per node
+            log_configs = [
+                p.get("container", {}).get("logConfiguration", {})
+                for p in job_node_properties.get("nodeRangeProperties", {})
+            ]
+            # one stream name per attempt
+            stream_names = [a.get("container", {}).get("logStreamName") for a in job_desc.get("attempts", [])]
+        elif job_container_desc:
+            log_configs = [job_container_desc.get("logConfiguration", {})]
+            stream_name = job_container_desc.get("logStreamName")
+            stream_names = [stream_name] if stream_name is not None else []
+        else:
+            raise AirflowException(
+                f"AWS Batch job ({job_id}) is not a supported job type. "
+                "Supported job types: container, array, multinode."
+            )
+
+        # If the user selected another logDriver than "awslogs", then CloudWatch logging is disabled.
+        if any([c.get("logDriver", "awslogs") != "awslogs" for c in log_configs]):
             self.log.warning(
-                "AWS Batch job (%s) uses logDriver (%s). AWS CloudWatch logging disabled.", job_id, log_driver
+                f"AWS Batch job ({job_id}) uses non-aws log drivers. AWS CloudWatch logging disabled."
             )
-            return None
+            return []
 
-        awslogs_stream_name = job_container_desc.get("logStreamName")
-        if not awslogs_stream_name:
-            # In case of call this method on very early stage of running AWS Batch
-            # there is possibility than AWS CloudWatch Stream Name not exists yet.
-            # AWS CloudWatch Stream Name also not created in case of misconfiguration.
-            self.log.warning("AWS Batch job (%s) doesn't create AWS CloudWatch Stream.", job_id)
-            return None
+        if not stream_names:
+            # If this method is called very early after starting the AWS Batch job,
+            # there is a possibility that the AWS CloudWatch Stream Name would not exist yet.
+            # This can also happen in case of misconfiguration.
+            self.log.warning(f"AWS Batch job ({job_id}) doesn't have any AWS CloudWatch Stream.")
+            return []
 
         # Try to get user-defined log configuration options
-        log_options = log_configuration.get("options", {})
-
-        return {
-            "awslogs_stream_name": awslogs_stream_name,
-            "awslogs_group": log_options.get("awslogs-group", "/aws/batch/job"),
-            "awslogs_region": log_options.get("awslogs-region", self.conn_region_name),
-        }
+        log_options = [c.get("options", {}) for c in log_configs]
+
+        # cross stream names with options (i.e. attempts X nodes) to generate all log infos
+        result = []
+        for stream in stream_names:
+            for option in log_options:
+                result.append(
+                    {
+                        "awslogs_stream_name": stream,
+                        # If the user did not specify anything, the default settings are:
+                        #   awslogs-group = /aws/batch/job
+                        #   awslogs-region = `same as AWS Batch Job region`
+                        "awslogs_group": option.get("awslogs-group", "/aws/batch/job"),
+                        "awslogs_region": option.get("awslogs-region", self.conn_region_name),
+                    }
+                )
+        return result
 
     @staticmethod
     def add_jitter(delay: int | float, width: int | float = 1, minima: int | float = 0) -> float:

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

@@ -138,7 +138,9 @@ class BatchWaitersHook(BatchClientHook):
         """
         return self._waiter_model
 
-    def get_waiter(self, waiter_name: str, _: dict[str, str] | None = None) -> botocore.waiter.Waiter:
+    def get_waiter(
+        self, waiter_name: str, _: dict[str, str] | None = None, deferrable: bool = False, client=None
+    ) -> botocore.waiter.Waiter:
         """
         Get an AWS Batch service waiter, using the configured ``.waiter_model``.
 

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

@@ -20,14 +20,12 @@ from __future__ import annotations
 import json
 import warnings
 from time import sleep
-from typing import Any, Callable
+from typing import Any
 
 from botocore.exceptions import ClientError
 
-from airflow.compat.functools import cached_property
 from airflow.exceptions import AirflowException, AirflowNotFoundException
 from airflow.providers.amazon.aws.hooks.base_aws import AwsBaseHook
-from airflow.providers.amazon.aws.utils.waiter import get_state, waiter
 from airflow.utils.helpers import prune_dict
 
 
@@ -254,66 +252,41 @@ class EmrServerlessHook(AwsBaseHook):
         kwargs["client_type"] = "emr-serverless"
         super().__init__(*args, **kwargs)
 
-    @cached_property
-    def conn(self):
-        """Get the underlying boto3 EmrServerlessAPIService client (cached)"""
-        return super().conn
-
-    # This method should be replaced with boto waiters which would implement timeouts and backoff nicely.
-    def waiter(
-        self,
-        get_state_callable: Callable,
-        get_state_args: dict,
-        parse_response: list,
-        desired_state: set,
-        failure_states: set,
-        object_type: str,
-        action: str,
-        countdown: int = 25 * 60,
-        check_interval_seconds: int = 60,
-    ) -> None:
+    def cancel_running_jobs(self, application_id: str, waiter_config: dict = {}):
         """
-        Will run the sensor until it turns True.
-
-        :param get_state_callable: A callable to run until it returns True
-        :param get_state_args: Arguments to pass to get_state_callable
-        :param parse_response: Dictionary keys to extract state from response of get_state_callable
-        :param desired_state: Wait until the getter returns this value
-        :param failure_states: A set of states which indicate failure and should throw an
-            exception if any are reached before the desired_state
-        :param object_type: Used for the reporting string. What are you waiting for? (application, job, etc)
-        :param action: Used for the reporting string. What action are you waiting for? (created, deleted, etc)
-        :param countdown: Total amount of time the waiter should wait for the desired state
-            before timing out (in seconds). Defaults to 25 * 60 seconds.
-        :param check_interval_seconds: Number of seconds waiter should wait before attempting
-            to retry get_state_callable. Defaults to 60 seconds.
+        List all jobs in an intermediate state and cancel them.
+        Then wait for those jobs to reach a terminal state.
+        Note: if new jobs are triggered while this operation is ongoing,
+        it's going to time out and return an error.
         """
-        warnings.warn(
-            """This method is deprecated.
-            Please use `airflow.providers.amazon.aws.utils.waiter.waiter`.""",
-            DeprecationWarning,
-            stacklevel=2,
-        )
-        waiter(
-            get_state_callable=get_state_callable,
-            get_state_args=get_state_args,
-            parse_response=parse_response,
-            desired_state=desired_state,
-            failure_states=failure_states,
-            object_type=object_type,
-            action=action,
-            countdown=countdown,
-            check_interval_seconds=check_interval_seconds,
-        )
-
-    def get_state(self, response, keys) -> str:
-        warnings.warn(
-            """This method is deprecated.
-            Please use `airflow.providers.amazon.aws.utils.waiter.get_state`.""",
-            DeprecationWarning,
-            stacklevel=2,
+        paginator = self.conn.get_paginator("list_job_runs")
+        results_per_response = 50
+        iterator = paginator.paginate(
+            applicationId=application_id,
+            states=list(self.JOB_INTERMEDIATE_STATES),
+            PaginationConfig={
+                "PageSize": results_per_response,
+            },
         )
-        return get_state(response=response, keys=keys)
+        count = 0
+        for r in iterator:
+            job_ids = [jr["id"] for jr in r["jobRuns"]]
+            count += len(job_ids)
+            if len(job_ids) > 0:
+                self.log.info(
+                    "Cancelling %s pending job(s) for the application %s so that it can be stopped",
+                    len(job_ids),
+                    application_id,
+                )
+                for job_id in job_ids:
+                    self.conn.cancel_job_run(applicationId=application_id, jobRunId=job_id)
+        if count > 0:
+            self.log.info("now waiting for the %s cancelled job(s) to terminate", count)
+            self.get_waiter("no_job_running").wait(
+                applicationId=application_id,
+                states=list(self.JOB_INTERMEDIATE_STATES.union({"CANCELLING"})),
+                WaiterConfig=waiter_config,
+            )
 
 
 class EmrContainerHook(AwsBaseHook):
@@ -483,7 +456,6 @@ class EmrContainerHook(AwsBaseHook):
     def poll_query_status(
         self,
         job_id: str,
-        max_tries: int | None = None,
         poll_interval: int = 30,
         max_polling_attempts: int | None = None,
     ) -> str | None:
@@ -492,22 +464,9 @@ class EmrContainerHook(AwsBaseHook):
         Returns one of the final states.
 
         :param job_id: The ID of the job run request.
-        :param max_tries: Deprecated - Use max_polling_attempts instead
         :param poll_interval: Time (in seconds) to wait between calls to check query status on EMR
         :param max_polling_attempts: Number of times to poll for query state before function exits
         """
-        if max_tries:
-            warnings.warn(
-                f"Method `{self.__class__.__name__}.max_tries` is deprecated and will be removed "
-                "in a future release.  Please use method `max_polling_attempts` instead.",
-                DeprecationWarning,
-                stacklevel=2,
-            )
-            if max_polling_attempts and max_polling_attempts != max_tries:
-                raise Exception("max_polling_attempts must be the same value as max_tries")
-            else:
-                max_polling_attempts = max_tries
-
         try_number = 1
         final_query_state = None  # Query state when query reaches final state or max_polling_attempts reached
 

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

@@ -25,6 +25,14 @@ from typing import Generator
 
 from airflow.providers.amazon.aws.hooks.base_aws import AwsBaseHook
 
+# Guidance received from the AWS team regarding the correct way to check for the end of a stream is that the
+# value of the nextForwardToken is the same in subsequent calls.
+# The issue with this approach is, it can take a huge amount of time (e.g. 20 seconds) to retrieve logs using
+# this approach. As an intermediate solution, we decided to stop fetching logs if 3 consecutive responses
+# are empty.
+# See PR https://github.com/apache/airflow/pull/20814
+NUM_CONSECUTIVE_EMPTY_RESPONSE_EXIT_THRESHOLD = 3
+
 
 class AwsLogsHook(AwsBaseHook):
     """
@@ -69,14 +77,15 @@ class AwsLogsHook(AwsBaseHook):
                  |   'message' (str): The log event data.
                  |   'ingestionTime' (int): The time in milliseconds the event was ingested.
         """
+        num_consecutive_empty_response = 0
         next_token = None
         while True:
             if next_token is not None:
-                token_arg: dict[str, str] | None = {"nextToken": next_token}
+                token_arg: dict[str, str] = {"nextToken": next_token}
             else:
                 token_arg = {}
 
-            response = self.get_conn().get_log_events(
+            response = self.conn.get_log_events(
                 logGroupName=log_group,
                 logStreamName=log_stream_name,
                 startTime=start_time,
@@ -96,7 +105,16 @@ class AwsLogsHook(AwsBaseHook):
 
             yield from events
 
-            if next_token != response["nextForwardToken"]:
-                next_token = response["nextForwardToken"]
+            if not event_count:
+                num_consecutive_empty_response += 1
+                if num_consecutive_empty_response >= NUM_CONSECUTIVE_EMPTY_RESPONSE_EXIT_THRESHOLD:
+                    # Exit if there are more than NUM_CONSECUTIVE_EMPTY_RESPONSE_EXIT_THRESHOLD consecutive
+                    # empty responses
+                    return
+            elif next_token != response["nextForwardToken"]:
+                num_consecutive_empty_response = 0
             else:
+                # Exit if the value of nextForwardToken is same in subsequent calls
                 return
+
+            next_token = response["nextForwardToken"]