This is an automated email from the ASF dual-hosted git repository.
pankajkoti pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new b9f3a5392f Extend task context logging support for remote logging
using GCP GCS (#32970)
b9f3a5392f is described below
commit b9f3a5392f1b2fb8043ca56b2c8ded20f40e2297
Author: Pankaj Koti <[email protected]>
AuthorDate: Fri Nov 17 20:17:13 2023 +0530
Extend task context logging support for remote logging using GCP GCS
(#32970)
With the addition of taxt context logging feature in PR #32646,
this PR extends the feature to GCP Cloud storage 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/google/cloud/log/gcs_task_handler.py | 18 ++++++++++++++----
1 file changed, 14 insertions(+), 4 deletions(-)
diff --git a/airflow/providers/google/cloud/log/gcs_task_handler.py
b/airflow/providers/google/cloud/log/gcs_task_handler.py
index cebe9a829e..39d0f072a8 100644
--- a/airflow/providers/google/cloud/log/gcs_task_handler.py
+++ b/airflow/providers/google/cloud/log/gcs_task_handler.py
@@ -22,7 +22,7 @@ import os
import shutil
from functools import cached_property
from pathlib import Path
-from typing import Collection
+from typing import TYPE_CHECKING, Collection
# not sure why but mypy complains on missing `storage` but it is clearly there
and is importable
from google.cloud import storage # type: ignore[attr-defined]
@@ -36,6 +36,9 @@ from airflow.providers.google.common.consts import CLIENT_INFO
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
+
_DEFAULT_SCOPESS = frozenset(
[
"https://www.googleapis.com/auth/devstorage.read_write";,
@@ -96,6 +99,7 @@ class GCSTaskHandler(FileTaskHandler, LoggingMixin):
**kwargs,
):
super().__init__(base_log_folder, filename_template)
+ self.handler: logging.FileHandler | None = None
self.remote_base = gcs_log_folder
self.log_relative_path = ""
self.closed = False
@@ -137,15 +141,21 @@ class GCSTaskHandler(FileTaskHandler, LoggingMixin):
project=self.project_id if self.project_id else project_id,
)
- 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)
# Log relative path is used to construct local and remote
# log path to upload log files into GCS and read from the
# remote location.
+ if TYPE_CHECKING:
+ assert self.handler is not 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):
"""Close and upload local log file to remote storage GCS."""
