Repository: apache/airflow

Issue:
Enhance Airflow Logs API to fetch logs from Amazon Cloudwatch with time range
### Apache Airflow version

Other Airflow 2 version (please specify below)

### What happened

MWAA Version: 2.4.3
Airflow Version: 2.4.3

Airflow Logs currently do not fetch logs from Cloudwatch without time range, so when the cloudwatch logs are large and CloudWatch log streams are OLD, the airflow UI cannot display logs with error message:

```
*** Reading remote log from Cloudwatch log_group: airflow-cdp-airflow243-XXXX-Task log_stream: dag_id=<DAG_NAME>/run_id=scheduled__2023-07-27T07_25_00+00_00/task_id=<TASK_ID>/attempt=1.log.
Could not read remote logs from log_group: airflow-cdp-airflow243-XXXXXX-Task log_group: airflow-cdp-airflow243-XXXX-Task log_stream: dag_id=<DAG_NAME>/run_id=scheduled__2023-07-27T07_25_00+00_00/task_id=<TASK_ID>/attempt=1.log
```

The Airflow API need to pass start and end timestamps to GetLogEvents API from Amazon CloudWatch to resolve this error and it also improves performance of fetching logs.

This is critical issue for customers when they would like to fetch logs to investigate failed pipelines form few days to weeks old

### What you think should happen instead

The Airflow API need to pass start and end timestamps to GetLogEvents API from Amazon CloudWatch to resolve this error.
This should also improve performance of fetching logs.


### How to reproduce

This issue is intermittent and happens mostly on FAILD tasks.

1. Log onto Amazon MWAA Service
2. Open Airflow UI
3. Select DAG
4. Select the Failed Tasks 
5. Select Logs
You should see error message like below in the logs:
```
*** Reading remote log from Cloudwatch log_group: airflow-cdp-airflow243-XXXX-Task log_stream: dag_id=<DAG_NAME>/run_id=scheduled__2023-07-27T07_25_00+00_00/task_id=<TASK_ID>/attempt=1.log.
Could not read remote logs from log_group: airflow-cdp-airflow243-XXXXXX-Task log_group: airflow-cdp-airflow243-XXXX-Task log_stream: dag_id=<DAG_NAME>/run_id=scheduled__2023-07-27T07_25_00+00_00/task_id=<TASK_ID>/attempt=1.log
```

### Operating System

Running with Amazon MWAA

### Versions of Apache Airflow Providers

apache-airflow-providers-amazon==8.3.1
apache-airflow==2.4.3

### Deployment

Amazon (AWS) MWAA

### Deployment details

_No response_

### Anything else

_No response_

### Are you willing to submit PR?

- [ ] Yes I am willing to submit a PR!

### Code of Conduct

- [X] I agree to follow this project's [Code of Conduct](https://github.com/apache/airflow/blob/main/CODE_OF_CONDUCT.md)



Discussion:
Thanks for opening your first issue here! Be sure to follow the issue template! If you are willing to raise PR to address this issue please do so, no need to wait for approval.

@rushabh-lokhande would you like to work on fixing this issue?
Hey @rushabh-lokhande, in your opinion, where would the values for start time and end time come from? I'm guessing from the start and end time of the DAG itself?
> Hey @rushabh-lokhande, in your opinion, where would the values for start time and end time come from? I'm guessing from the start and end time of the DAG itself?

Yes @ivica-k we will pick up start time and end time with UTC conversion if needed, as ClopudWatch translates the dates to UTC. Also, even if we send start_date and end_date instead of start_timestamps and end_timestamps, that should narrow down Amazon Cloudwatch logs window and should improve performance.

Resolving patch:
diff --git a/airflow/providers/amazon/aws/hooks/logs.py b/airflow/providers/amazon/aws/hooks/logs.py
--- a/airflow/providers/amazon/aws/hooks/logs.py
+++ b/airflow/providers/amazon/aws/hooks/logs.py
@@ -22,6 +22,7 @@
 
 from airflow.exceptions import AirflowProviderDeprecationWarning
 from airflow.providers.amazon.aws.hooks.base_aws import AwsBaseHook
+from airflow.utils.helpers import prune_dict
 
 # 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.
@@ -60,6 +61,7 @@ def get_log_events(
         log_group: str,
         log_stream_name: str,
         start_time: int = 0,
+        end_time: int | None = None,
         skip: int = 0,
         start_from_head: bool | None = None,
         continuation_token: ContinuationToken | None = None,
@@ -72,7 +74,9 @@ def get_log_events(
 
         :param log_group: The name of the log group.
         :param log_stream_name: The name of the specific stream.
-        :param start_time: The time stamp value to start reading the logs from (default: 0).
+        :param start_time: The timestamp value in ms to start reading the logs from (default: 0).
+        :param end_time: The timestamp value in ms to stop reading the logs from (default: None).
+            If None is provided, reads it until the end of the log stream
         :param skip: The number of log entries to skip at the start (default: 0).
             This is for when there are multiple entries at the same timestamp.
         :param start_from_head: Deprecated. Do not use with False, logs would be retrieved out of order.
@@ -110,11 +114,16 @@ def get_log_events(
                 token_arg = {}
 
             response = self.conn.get_log_events(
-                logGroupName=log_group,
-                logStreamName=log_stream_name,
-                startTime=start_time,
-                startFromHead=start_from_head,
-                **token_arg,
+                **prune_dict(
+                    {
+                        "logGroupName": log_group,
+                        "logStreamName": log_stream_name,
+                        "startTime": start_time,
+                        "endTime": end_time,
+                        "startFromHead": start_from_head,
+                        **token_arg,
+                    }
+                )
             )
 
             events = response["events"]
diff --git a/airflow/providers/amazon/aws/log/cloudwatch_task_handler.py b/airflow/providers/amazon/aws/log/cloudwatch_task_handler.py
--- a/airflow/providers/amazon/aws/log/cloudwatch_task_handler.py
+++ b/airflow/providers/amazon/aws/log/cloudwatch_task_handler.py
@@ -17,13 +17,15 @@
 # under the License.
 from __future__ import annotations
 
-from datetime import datetime
+from datetime import datetime, timedelta
 from functools import cached_property
 
 import watchtower
 
 from airflow.configuration import conf
+from airflow.models import TaskInstance
 from airflow.providers.amazon.aws.hooks.logs import AwsLogsHook
+from airflow.providers.amazon.aws.utils import datetime_to_epoch_utc_ms
 from airflow.utils.log.file_task_handler import FileTaskHandler
 from airflow.utils.log.logging_mixin import LoggingMixin
 
@@ -90,7 +92,8 @@ def _read(self, task_instance, try_number, metadata=None):
         try:
             return (
                 f"*** Reading remote log from Cloudwatch log_group: {self.log_group} "
-                f"log_stream: {stream_name}.\n{self.get_cloudwatch_logs(stream_name=stream_name)}\n",
+                f"log_stream: {stream_name}.\n"
+                f"{self.get_cloudwatch_logs(stream_name=stream_name, task_instance=task_instance)}\n",
                 {"end_of_log": True},
             )
         except Exception as e:
@@ -103,17 +106,29 @@ def _read(self, task_instance, try_number, metadata=None):
             log += local_log
             return log, metadata
 
-    def get_cloudwatch_logs(self, stream_name: str) -> str:
+    def get_cloudwatch_logs(self, stream_name: str, task_instance: TaskInstance) -> str:
         """
         Return all logs from the given log stream.
 
         :param stream_name: name of the Cloudwatch log stream to get all logs from
+        :param task_instance: the task instance to get logs about
         :return: string of all logs from the given log stream
         """
+        start_time = (
+            0 if task_instance.start_date is None else datetime_to_epoch_utc_ms(task_instance.start_date)
+        )
+        # If there is an end_date to the task instance, fetch logs until that date + 30 seconds
+        # 30 seconds is an arbitrary buffer so that we don't miss any logs that were emitted
+        end_time = (
+            None
+            if task_instance.end_date is None
+            else datetime_to_epoch_utc_ms(task_instance.end_date + timedelta(seconds=30))
+        )
         events = self.hook.get_log_events(
             log_group=self.log_group,
             log_stream_name=stream_name,
-            start_from_head=True,
+            start_time=start_time,
+            end_time=end_time,
         )
         return "\n".join(self._event_to_str(event) for event in events)
 
diff --git a/airflow/providers/amazon/aws/utils/__init__.py b/airflow/providers/amazon/aws/utils/__init__.py
--- a/airflow/providers/amazon/aws/utils/__init__.py
+++ b/airflow/providers/amazon/aws/utils/__init__.py
@@ -18,7 +18,7 @@
 
 import logging
 import re
-from datetime import datetime
+from datetime import datetime, timezone
 from enum import Enum
 
 from airflow.utils.helpers import prune_dict
@@ -55,6 +55,11 @@ def datetime_to_epoch_ms(date_time: datetime) -> int:
     return int(date_time.timestamp() * 1_000)
 
 
+def datetime_to_epoch_utc_ms(date_time: datetime) -> int:
+    """Convert a datetime object to an epoch integer (milliseconds) in UTC timezone."""
+    return int(date_time.replace(tzinfo=timezone.utc).timestamp() * 1_000)
+
+
 def datetime_to_epoch_us(date_time: datetime) -> int:
     """Convert a datetime object to an epoch integer (microseconds)."""
     return int(date_time.timestamp() * 1_000_000)


Test patch:
diff --git a/tests/providers/amazon/aws/hooks/test_logs.py b/tests/providers/amazon/aws/hooks/test_logs.py
--- a/tests/providers/amazon/aws/hooks/test_logs.py
+++ b/tests/providers/amazon/aws/hooks/test_logs.py
@@ -18,7 +18,7 @@
 from __future__ import annotations
 
 from unittest import mock
-from unittest.mock import patch
+from unittest.mock import ANY, patch
 
 import pytest
 from moto import mock_logs
@@ -29,7 +29,7 @@
 @mock_logs
 class TestAwsLogsHook:
     @pytest.mark.parametrize(
-        "get_log_events_response, num_skip_events, expected_num_events",
+        "get_log_events_response, num_skip_events, expected_num_events, end_time",
         [
             # 3 empty responses with different tokens
             (
@@ -40,6 +40,7 @@ class TestAwsLogsHook:
                 ],
                 0,
                 0,
+                None,
             ),
             # 2 events on the second response with same token
             (
@@ -49,6 +50,7 @@ class TestAwsLogsHook:
                 ],
                 0,
                 2,
+                None,
             ),
             # Different tokens, 2 events on the second response then 3 empty responses
             (
@@ -63,6 +65,7 @@ class TestAwsLogsHook:
                 ],
                 0,
                 2,
+                10,
             ),
             # 2 events on the second response, then 2 empty responses, then 2 consecutive responses with
             # 2 events with the same token
@@ -79,20 +82,36 @@ class TestAwsLogsHook:
                 ],
                 0,
                 6,
+                20,
             ),
         ],
     )
     @patch("airflow.providers.amazon.aws.hooks.logs.AwsLogsHook.conn", new_callable=mock.PropertyMock)
-    def test_get_log_events(self, mock_conn, get_log_events_response, num_skip_events, expected_num_events):
+    def test_get_log_events(
+        self, mock_conn, get_log_events_response, num_skip_events, expected_num_events, end_time
+    ):
         mock_conn().get_log_events.side_effect = get_log_events_response
 
+        log_group_name = "example-group"
+        log_stream_name = "example-log-stream"
         hook = AwsLogsHook(aws_conn_id="aws_default", region_name="us-east-1")
         events = hook.get_log_events(
-            log_group="example-group",
-            log_stream_name="example-log-stream",
+            log_group=log_group_name,
+            log_stream_name=log_stream_name,
             skip=num_skip_events,
+            end_time=end_time,
         )
 
         events = list(events)
 
         assert len(events) == expected_num_events
+        kwargs = {
+            "logGroupName": log_group_name,
+            "logStreamName": log_stream_name,
+            "startFromHead": True,
+            "startTime": 0,
+            "nextToken": ANY,
+        }
+        if end_time:
+            kwargs["endTime"] = end_time
+        mock_conn().get_log_events.assert_called_with(**kwargs)
diff --git a/tests/providers/amazon/aws/log/test_cloudwatch_task_handler.py b/tests/providers/amazon/aws/log/test_cloudwatch_task_handler.py
--- a/tests/providers/amazon/aws/log/test_cloudwatch_task_handler.py
+++ b/tests/providers/amazon/aws/log/test_cloudwatch_task_handler.py
@@ -18,7 +18,7 @@
 from __future__ import annotations
 
 import time
-from datetime import datetime as dt
+from datetime import datetime as dt, timedelta
 from unittest import mock
 from unittest.mock import call
 
@@ -31,6 +31,7 @@
 from airflow.operators.empty import EmptyOperator
 from airflow.providers.amazon.aws.hooks.logs import AwsLogsHook
 from airflow.providers.amazon.aws.log.cloudwatch_task_handler import CloudwatchTaskHandler
+from airflow.providers.amazon.aws.utils import datetime_to_epoch_utc_ms
 from airflow.utils.session import create_session
 from airflow.utils.state import State
 from airflow.utils.timezone import datetime
@@ -60,7 +61,6 @@ def setup_tests(self, create_log_template, tmp_path_factory):
             self.local_log_location,
             f"arn:aws:logs:{self.region_name}:11111111:log-group:{self.remote_log_group}",
         )
-        self.cloudwatch_task_handler.hook
 
         date = datetime(2020, 1, 1)
         dag_id = "dag_for_testing_cloudwatch_task_handler"
@@ -154,6 +154,39 @@ def test_read(self):
             [{"end_of_log": True}],
         )
 
+    @pytest.mark.parametrize(
+        "start_date, end_date, expected_start_time, expected_end_time",
+        [
+            (None, None, 0, None),
+            (datetime(2020, 1, 1), None, datetime_to_epoch_utc_ms(datetime(2020, 1, 1)), None),
+            (
+                None,
+                datetime(2020, 1, 2),
+                0,
+                datetime_to_epoch_utc_ms(datetime(2020, 1, 2) + timedelta(seconds=30)),
+            ),
+            (
+                datetime(2020, 1, 1),
+                datetime(2020, 1, 2),
+                datetime_to_epoch_utc_ms(datetime(2020, 1, 1)),
+                datetime_to_epoch_utc_ms(datetime(2020, 1, 2) + timedelta(seconds=30)),
+            ),
+        ],
+    )
+    @mock.patch.object(AwsLogsHook, "get_log_events")
+    def test_get_cloudwatch_logs(
+        self, mock_get_log_events, start_date, end_date, expected_start_time, expected_end_time
+    ):
+        self.ti.start_date = start_date
+        self.ti.end_date = end_date
+        self.cloudwatch_task_handler.get_cloudwatch_logs(self.remote_log_stream, self.ti)
+        mock_get_log_events.assert_called_once_with(
+            log_group=self.remote_log_group,
+            log_stream_name=self.remote_log_stream,
+            start_time=expected_start_time,
+            end_time=expected_end_time,
+        )
+
     def test_close_prevents_duplicate_calls(self):
         with mock.patch("watchtower.CloudWatchLogHandler.close") as mock_log_handler_close:
             with mock.patch("airflow.utils.log.file_task_handler.FileTaskHandler.set_context"):