Extend task context logging support for remote logging using AWS S3

airflow/providers/amazon/aws/log/s3_task_handler.py

@@ -17,10 +17,12 @@
 # under the License.
 from __future__ import annotations
 
+import logging
 import os
 import pathlib
 import shutil
 from functools import cached_property
+from typing import TYPE_CHECKING
 
 from packaging.version import Version
 
@@ -29,6 +31,9 @@
 from airflow.utils.log.file_task_handler import FileTaskHandler
 from airflow.utils.log.logging_mixin import LoggingMixin
 
+if TYPE_CHECKING:
+    from airflow.models.taskinstance import TaskInstance
+
 
 def get_default_delete_local_copy():
     """Load delete_local_logs conf if Airflow version > 2.6 and return False if not.
@@ -55,6 +60,7 @@ def __init__(
         self, base_log_folder: str, s3_log_folder: str, filename_template: str | None = None, **kwargs
     ):
         super().__init__(base_log_folder, filename_template)
+        self.handler: logging.FileHandler | None = None
         self.remote_base = s3_log_folder
         self.log_relative_path = ""
         self._hook = None
@@ -71,14 +77,20 @@ def hook(self):
             aws_conn_id=conf.get("logging", "REMOTE_LOG_CONN_ID"), transfer_config_args={"use_threads": False}
         )
 
-    def set_context(self, ti):
-        super().set_context(ti)
+    def set_context(self, ti: TaskInstance, *, identifier: str | None = None) -> None:
+        if getattr(self, "supports_task_context_logging", False):
+            super().set_context(ti, identifier=identifier)
+        else:
+            super().set_context(ti)
         # Local location and remote location is needed to open and
         # upload local log file to S3 remote storage.
+        if TYPE_CHECKING:
+            assert self.handler is not None
+
         full_path = self.handler.baseFilename
         self.log_relative_path = pathlib.Path(full_path).relative_to(self.local_base).as_posix()
         is_trigger_log_context = getattr(ti, "is_trigger_log_context", False)
-        self.upload_on_close = is_trigger_log_context or not ti.raw
+        self.upload_on_close = is_trigger_log_context or not getattr(ti, "raw", None)
         # Clear the file first so that duplicate data is not uploaded
         # when re-using the same path (e.g. with rescheduled sensors)
         if self.upload_on_close: