Extend task context logging support for remote logging using WASB (Az…

…ure Blob Storage) (apache#32972)

With the addition of task context logging feature in PR apache#32646,
this PR extends the feature to Azure Blob Storage (WASB)  when is
it set as remote logging store. Here, backward compatibility is
ensured for older versions of Airflow that do not have the feature
included in Airflow Core.

airflow/providers/microsoft/azure/log/wasb_task_handler.py

@@ -95,8 +95,11 @@ def hook(self):
             )
             return None
 
-    def set_context(self, ti: TaskInstance, **kwargs) -> None:
-        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 Wasb remote storage.
         if TYPE_CHECKING:
@@ -105,7 +108,7 @@ def set_context(self, ti: TaskInstance, **kwargs) -> None:
         full_path = self.handler.baseFilename
         self.log_relative_path = 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)
 
     def close(self) -> None:
         """Close and upload local log file to remote storage Wasb."""