Use a continuation token to get logs in ecs

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

@@ -171,17 +171,20 @@ def __init__(
         self.hook = AwsLogsHook(aws_conn_id=aws_conn_id, region_name=region_name)
 
     def run(self) -> None:
-        logs_to_skip = 0
+        continuation_token = AwsLogsHook.ContinuationToken()
         while not self.is_stopped():
             time.sleep(self.fetch_interval.total_seconds())
-            log_events = self._get_log_events(logs_to_skip)
+            log_events = self._get_log_events(continuation_token)
             for log_event in log_events:
                 self.logger.info(self._event_to_str(log_event))
-                logs_to_skip += 1
 
-    def _get_log_events(self, skip: int = 0) -> Generator:
+    def _get_log_events(self, skip_token: AwsLogsHook.ContinuationToken | None = None) -> Generator:
+        if skip_token is None:
+            skip_token = AwsLogsHook.ContinuationToken()
         try:
-            yield from self.hook.get_log_events(self.log_group, self.log_stream_name, skip=skip)
+            yield from self.hook.get_log_events(
+                self.log_group, self.log_stream_name, continuation_token=skip_token
+            )
         except ClientError as error:
             if error.response["Error"]["Code"] != "ResourceNotFoundException":
                 self.logger.warning("Error on retrieving Cloudwatch log events", error)

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

@@ -50,13 +50,20 @@ def __init__(self, *args, **kwargs) -> None:
         kwargs["client_type"] = "logs"
         super().__init__(*args, **kwargs)
 
+    class ContinuationToken:
+        """Just a wrapper around a str token to allow updating it from the caller."""
+
+        def __init__(self):
+            self.value: str | None = None
+
     def get_log_events(
         self,
         log_group: str,
         log_stream_name: str,
         start_time: int = 0,
         skip: int = 0,
         start_from_head: bool = True,
+        continuation_token: ContinuationToken | None = None,
     ) -> Generator:
         """
         A generator for log items in a single stream. This will yield all the
@@ -72,16 +79,20 @@ def get_log_events(
             This is for when there are multiple entries at the same timestamp.
         :param start_from_head: whether to start from the beginning (True) of the log or
             at the end of the log (False).
+        :param continuation_token: a token indicating where to read logs from.
+            Will be updated as this method reads new logs, to be reused in subsequent calls.
         :return: | A CloudWatch log event with the following key-value pairs:
                  |   'timestamp' (int): The time in milliseconds of the event.
                  |   'message' (str): The log event data.
                  |   'ingestionTime' (int): The time in milliseconds the event was ingested.
         """
+        if continuation_token is None:
+            continuation_token = AwsLogsHook.ContinuationToken()
+
         num_consecutive_empty_response = 0
-        next_token = None
         while True:
-            if next_token is not None:
-                token_arg: dict[str, str] = {"nextToken": next_token}
+            if continuation_token.value is not None:
+                token_arg: dict[str, str] = {"nextToken": continuation_token.value}
             else:
                 token_arg = {}
 
@@ -105,16 +116,16 @@ def get_log_events(
 
             yield from events
 
+            if continuation_token.value == response["nextForwardToken"]:
+                return
+
             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
+                num_consecutive_empty_response = 0
 
-            next_token = response["nextForwardToken"]
+            continuation_token.value = response["nextForwardToken"]

tests/providers/amazon/aws/hooks/test_ecs.py

@@ -18,12 +18,14 @@
 
 from datetime import timedelta
 from unittest import mock
+from unittest.mock import PropertyMock
 
 import pytest
 from botocore.exceptions import ClientError
 
 from airflow.providers.amazon.aws.exceptions import EcsOperatorError, EcsTaskFailToStart
 from airflow.providers.amazon.aws.hooks.ecs import EcsHook, EcsTaskLogFetcher, should_retry, should_retry_eni
+from airflow.providers.amazon.aws.hooks.logs import AwsLogsHook
 
 DEFAULT_CONN_ID: str = "aws_default"
 REGION: str = "us-east-1"
@@ -146,6 +148,20 @@ def test_get_log_events_with_unexpected_error(self, get_log_events_mock):
         with pytest.raises(Exception):
             next(self.log_fetcher._get_log_events())
 
+    @mock.patch.object(AwsLogsHook, "conn", new_callable=PropertyMock)
+    def test_get_log_events_updates_token(self, logs_conn_mock):
+        logs_conn_mock().get_log_events.return_value = {
+            "events": ["my_event"],
+            "nextForwardToken": "my_next_token",
+        }
+
+        token = AwsLogsHook.ContinuationToken()
+        list(self.log_fetcher._get_log_events(token))
+
+        assert token.value == "my_next_token"
+        # 2 calls expected, it's only on the second one that the stop condition old_token == next_token is met
+        assert logs_conn_mock().get_log_events.call_count == 2
+
     def test_event_to_str(self):
         events = [
             {"timestamp": 1617400267123, "message": "First"},