amoghrajesh commented on code in PR #65198:
URL: https://github.com/apache/airflow/pull/65198#discussion_r3194776807
##########
providers/google/src/airflow/providers/google/cloud/log/stackdriver_task_handler.py:
##########
@@ -56,6 +66,209 @@
["https://www.googleapis.com/auth/logging.read";,
"https://www.googleapis.com/auth/logging.write";]
)
+LABEL_TASK_ID = "task_id"
+LABEL_DAG_ID = "dag_id"
+LABEL_LOGICAL_DATE = "logical_date" if AIRFLOW_V_3_0_PLUS else "execution_date"
+LABEL_TRY_NUMBER = "try_number"
+
+
[email protected](kw_only=True)
+class StackdriverRemoteLogIO(LoggingMixin):
+ """Remote log IO that streams logs to and reads from Google Cloud
Stackdriver Logging."""
+
+ base_log_folder: Path = attrs.field(converter=Path)
+ delete_local_copy: bool = True
+
+ gcp_key_path: str | None = None
+ scopes: Collection[str] | None = _DEFAULT_SCOPESS
+ gcp_log_name: str = DEFAULT_LOGGER_NAME
+ transport_type: type[Transport] = BackgroundThreadTransport
+ resource: Resource = _GLOBAL_RESOURCE
+ labels: dict[str, str] | None = None
+
+ @cached_property
+ def _credentials_and_project(self) -> tuple[Credentials, str]:
+ credentials, project = get_credentials_and_project_id(
+ key_path=self.gcp_key_path, scopes=self.scopes,
disable_logging=True
+ )
+ return credentials, project
+
+ @cached_property
+ def _client(self) -> gcp_logging.Client:
+ """The Cloud Library API client."""
+ credentials, project = self._credentials_and_project
+ return gcp_logging.Client(
+ credentials=credentials,
+ project=project,
+ client_info=CLIENT_INFO,
+ )
+
+ @cached_property
+ def _logging_service_client(self) -> LoggingServiceV2Client:
+ """The Cloud logging service v2 client."""
+ credentials, _ = self._credentials_and_project
+ return LoggingServiceV2Client(
+ credentials=credentials,
+ client_info=CLIENT_INFO,
+ )
+
+ @cached_property
+ def _transport(self) -> Transport:
+ """Object responsible for sending data to Stackdriver."""
+ return self.transport_type(self._client, self.gcp_log_name)
+
+ @cached_property
+ def processors(self) -> tuple[structlog.typing.Processor, ...]:
+ from datetime import datetime
+ from logging import getLogRecordFactory
Review Comment:
These can go to top level
##########
providers/google/src/airflow/providers/google/cloud/log/stackdriver_task_handler.py:
##########
@@ -375,12 +455,12 @@ def get_external_log_url(self, task_instance:
TaskInstance, try_number: int) ->
:param try_number: task instance try_number to read logs from
:return: URL to the external log collection service
"""
- _, project_id = self._credentials_and_project
+ _, project_id = self.io._credentials_and_project
- ti_labels = self._task_instance_to_labels(task_instance)
- ti_labels[self.LABEL_TRY_NUMBER] = str(try_number)
+ ti_labels = _task_instance_to_labels(task_instance)
+ ti_labels[LABEL_TRY_NUMBER] = str(try_number)
- log_filter = self._prepare_log_filter(ti_labels)
+ log_filter = self.io._prepare_log_filter(ti_labels)
Review Comment:
Same comment as above here
##########
providers/google/tests/unit/google/cloud/log/test_stackdriver_task_handler.py:
##########
@@ -17,15 +17,21 @@
from __future__ import annotations
import logging
+import tempfile
Review Comment:
Need to add missing test for processors.
##########
providers/google/src/airflow/providers/google/cloud/log/stackdriver_task_handler.py:
##########
@@ -225,18 +408,18 @@ def read(
if not metadata:
metadata = {}
- ti_labels = self._task_instance_to_labels(task_instance)
+ ti_labels = _task_instance_to_labels(task_instance)
if try_number is not None:
- ti_labels[self.LABEL_TRY_NUMBER] = str(try_number)
+ ti_labels[LABEL_TRY_NUMBER] = str(try_number)
else:
- del ti_labels[self.LABEL_TRY_NUMBER]
+ del ti_labels[LABEL_TRY_NUMBER]
- log_filter = self._prepare_log_filter(ti_labels)
+ log_filter = self.io._prepare_log_filter(ti_labels)
next_page_token = metadata.get("next_page_token", None)
all_pages = "download_logs" in metadata and metadata["download_logs"]
- messages, end_of_log, next_page_token = self._read_logs(log_filter,
next_page_token, all_pages)
+ messages, end_of_log, next_page_token = self.io._read_logs(log_filter,
next_page_token, all_pages)
Review Comment:
This is still reaching into private internals. The `CloudWatchTaskHandler`
only calls public methods for example. You can make `_prepare_log_filter` and
`_read_logs`
##########
providers/google/src/airflow/providers/google/cloud/log/stackdriver_task_handler.py:
##########
@@ -56,6 +66,209 @@
["https://www.googleapis.com/auth/logging.read";,
"https://www.googleapis.com/auth/logging.write";]
)
+LABEL_TASK_ID = "task_id"
+LABEL_DAG_ID = "dag_id"
+LABEL_LOGICAL_DATE = "logical_date" if AIRFLOW_V_3_0_PLUS else "execution_date"
+LABEL_TRY_NUMBER = "try_number"
+
+
[email protected](kw_only=True)
+class StackdriverRemoteLogIO(LoggingMixin):
+ """Remote log IO that streams logs to and reads from Google Cloud
Stackdriver Logging."""
+
+ base_log_folder: Path = attrs.field(converter=Path)
+ delete_local_copy: bool = True
+
+ gcp_key_path: str | None = None
+ scopes: Collection[str] | None = _DEFAULT_SCOPESS
+ gcp_log_name: str = DEFAULT_LOGGER_NAME
+ transport_type: type[Transport] = BackgroundThreadTransport
+ resource: Resource = _GLOBAL_RESOURCE
+ labels: dict[str, str] | None = None
+
+ @cached_property
+ def _credentials_and_project(self) -> tuple[Credentials, str]:
+ credentials, project = get_credentials_and_project_id(
+ key_path=self.gcp_key_path, scopes=self.scopes,
disable_logging=True
+ )
+ return credentials, project
+
+ @cached_property
+ def _client(self) -> gcp_logging.Client:
+ """The Cloud Library API client."""
+ credentials, project = self._credentials_and_project
+ return gcp_logging.Client(
+ credentials=credentials,
+ project=project,
+ client_info=CLIENT_INFO,
+ )
+
+ @cached_property
+ def _logging_service_client(self) -> LoggingServiceV2Client:
+ """The Cloud logging service v2 client."""
+ credentials, _ = self._credentials_and_project
+ return LoggingServiceV2Client(
+ credentials=credentials,
+ client_info=CLIENT_INFO,
+ )
+
+ @cached_property
+ def _transport(self) -> Transport:
+ """Object responsible for sending data to Stackdriver."""
+ return self.transport_type(self._client, self.gcp_log_name)
+
+ @cached_property
+ def processors(self) -> tuple[structlog.typing.Processor, ...]:
+ from datetime import datetime
+ from logging import getLogRecordFactory
+
+ import structlog.stdlib
+
+ logRecordFactory = getLogRecordFactory()
+ _transport = self._transport
+
+ from airflow.sdk.log import relative_path_from_logger
+
+ def proc(
+ logger: structlog.typing.WrappedLogger,
+ method_name: str,
+ event: structlog.typing.EventDict,
+ ):
+ if not logger or not relative_path_from_logger(logger):
+ return event
+
+ name = event.get("logger_name") or event.get("logger", "")
+ level = structlog.stdlib.NAME_TO_LEVEL.get(method_name.lower(),
logging.INFO)
+ msg = copy.copy(event)
+ created = None
+ if ts := msg.pop("timestamp", None):
+ with contextlib.suppress(Exception):
+ created = datetime.fromisoformat(ts)
+ record = logRecordFactory(
+ name,
+ level,
+ pathname="",
+ lineno=0,
+ msg=msg,
+ args=(),
+ exc_info=None,
+ func=None,
+ sinfo=None,
+ )
+ if created is not None:
+ ct = created.timestamp()
+ record.created = ct
+ record.msecs = int((ct - int(ct)) * 1000) + 0.0
+ _transport.send(
+ record, str(msg.get("event", "")), resource=self.resource,
labels=self.labels or {}
+ )
+ return event
+
+ return (proc,)
+
+ def upload(self, path: os.PathLike | str, ti: RuntimeTI) -> None:
+ """Flush the transport and optionally delete local log files."""
+ self._transport.flush()
+ if self.delete_local_copy:
+ base = self.base_log_folder.resolve()
+ raw = Path(path)
+ local_path = (raw if raw.is_absolute() else base / raw).resolve()
+ try:
+ local_path.relative_to(base)
+ except ValueError:
+ self.log.warning(
+ "Skipping deletion: path %s is outside base_log_folder %s",
+ local_path,
+ base,
+ )
+ return
+ parent = local_path.parent
+ if parent.exists():
+ shutil.rmtree(parent, ignore_errors=True)
+
+ def read(self, relative_path: str, ti: RuntimeTI) -> LogResponse:
+ """Read logs from Stackdriver Logging using task instance labels."""
+ ti_labels = _task_instance_to_labels(ti)
+ log_filter = self._prepare_log_filter(ti_labels)
+ messages, end_of_log, _ = self._read_logs(log_filter,
next_page_token=None, all_pages=True)
+ return [f"Reading remote log from Stackdriver for {relative_path}"],
[messages] if messages else []
+
+ def _prepare_log_filter(self, ti_labels: dict[str, str]) -> str:
+ def escape_label_key(key: str) -> str:
+ return f'"{key}"' if "." in key else key
+
+ def escape_label_value(value: str) -> str:
+ escaped_value = value.replace("\\", "\\\\").replace('"', '\\"')
+ return f'"{escaped_value}"'
+
+ _, project = self._credentials_and_project
+ log_filters = [
+ f"resource.type={escape_label_value(self.resource.type)}",
+ f'logName="projects/{project}/logs/{self.gcp_log_name}"',
+ ]
+
+ for key, value in self.resource.labels.items():
+
log_filters.append(f"resource.labels.{escape_label_key(key)}={escape_label_value(value)}")
+
+ for key, value in ti_labels.items():
+
log_filters.append(f"labels.{escape_label_key(key)}={escape_label_value(value)}")
+ return "\n".join(log_filters)
+
+ def _read_logs(
+ self, log_filter: str, next_page_token: str | None, all_pages: bool
+ ) -> tuple[str, bool, str | None]:
+ messages = []
+ new_messages, next_page_token = self._read_single_logs_page(
+ log_filter=log_filter,
+ page_token=next_page_token,
+ )
+ messages.append(new_messages)
+ if all_pages:
+ while next_page_token:
+ new_messages, next_page_token = self._read_single_logs_page(
+ log_filter=log_filter, page_token=next_page_token
+ )
+ messages.append(new_messages)
+ if not messages:
+ break
Review Comment:
This is dead code. We delete after append. (maybe older bug but we can fix
it now)
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
