diff --git a/providers/google/src/airflow/providers/google/cloud/log/stackdriver_task_handler.py b/providers/google/src/airflow/providers/google/cloud/log/stackdriver_task_handler.py index b11a1cd6a4728..145c2564539b3 100644 --- a/providers/google/src/airflow/providers/google/cloud/log/stackdriver_task_handler.py +++ b/providers/google/src/airflow/providers/google/cloud/log/stackdriver_task_handler.py @@ -18,13 +18,21 @@ from __future__ import annotations +import contextlib +import copy import logging +import os +import shutil import warnings from collections.abc import Collection +from datetime import datetime from functools import cached_property +from logging import getLogRecordFactory +from pathlib import Path from typing import TYPE_CHECKING from urllib.parse import urlencode +import attrs from google.cloud import logging as gcp_logging from google.cloud.logging import Resource from google.cloud.logging.handlers.transports import BackgroundThreadTransport, Transport @@ -35,6 +43,7 @@ from airflow.providers.google.cloud.utils.credentials_provider import get_credentials_and_project_id from airflow.providers.google.common.consts import CLIENT_INFO from airflow.providers.google.version_compat import AIRFLOW_V_3_0_PLUS +from airflow.utils.log.logging_mixin import LoggingMixin try: from airflow.sdk.definitions._internal.types import NOTSET, ArgNotSet @@ -45,9 +54,12 @@ from airflow.utils.log.trigger_handler import ctx_indiv_trigger if TYPE_CHECKING: + import structlog.typing from google.auth.credentials import Credentials from airflow.models import TaskInstance + from airflow.sdk.types import RuntimeTaskInstanceProtocol as RuntimeTI + from airflow.utils.log.file_task_handler import LogResponse DEFAULT_LOGGER_NAME = "airflow" _GLOBAL_RESOURCE = Resource(type="global", labels={}) @@ -56,6 +68,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" + + +@attrs.define(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, ...]: + import structlog.stdlib + + from airflow.sdk.log import relative_path_from_logger + + log_record_factory = getLogRecordFactory() + _transport = self.transport + + 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 = log_record_factory( + 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 + + ti = getattr(record, "task_instance", None) + labels: dict[str, str] = {} + if self.labels: + labels.update(self.labels) + if ti: + labels.update(_task_instance_to_labels(ti)) + _transport.send(record, str(msg.get("event", "")), resource=self.resource, labels=labels) + 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) + + end_of_log = True + next_page_token = None + else: + end_of_log = not bool(next_page_token) + return "\n".join(messages), end_of_log, next_page_token + + def _read_single_logs_page(self, log_filter: str, page_token: str | None = None) -> tuple[str, str]: + _, project = self.credentials_and_project + request = ListLogEntriesRequest( + resource_names=[f"projects/{project}"], + filter=log_filter, + page_token=page_token, + order_by="timestamp asc", + page_size=1000, + ) + response = self._logging_service_client.list_log_entries(request=request) + page: ListLogEntriesResponse = next(response.pages) + messages: list[str] = [] + for entry in page.entries: + if "message" in (entry.json_payload or {}): + messages.append(entry.json_payload["message"]) # type: ignore + elif entry.text_payload: + messages.append(entry.text_payload) + return "\n".join(messages), page.next_page_token + + +def _task_instance_to_labels(ti) -> dict[str, str]: + """Convert a task instance to Stackdriver labels.""" + return { + LABEL_TASK_ID: ti.task_id, + LABEL_DAG_ID: ti.dag_id, + LABEL_LOGICAL_DATE: str(ti.logical_date.isoformat()) + if AIRFLOW_V_3_0_PLUS + else str(ti.execution_date.isoformat()), + LABEL_TRY_NUMBER: str(ti.try_number), + } + class StackdriverTaskHandler(logging.Handler): """ @@ -88,10 +303,11 @@ class StackdriverTaskHandler(logging.Handler): :param labels: (Optional) Mapping of labels for the entry. """ - 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" + # Re-export module-level constants for back-compat with external code reading them off the class + LABEL_TASK_ID = LABEL_TASK_ID + LABEL_DAG_ID = LABEL_DAG_ID + LABEL_LOGICAL_DATE = LABEL_LOGICAL_DATE + LABEL_TRY_NUMBER = LABEL_TRY_NUMBER LOG_VIEWER_BASE_URL = "https://console.cloud.google.com/logs/viewer" LOG_NAME = "Google Stackdriver" @@ -120,53 +336,23 @@ def __init__( gcp_log_name = str(name) super().__init__() - self.gcp_key_path: str | None = gcp_key_path - self.scopes: Collection[str] | None = scopes - self.gcp_log_name: str = gcp_log_name - self.transport_type: type[Transport] = transport - self.resource: Resource = resource + self.io = StackdriverRemoteLogIO( + base_log_folder=Path("."), + gcp_key_path=gcp_key_path, + scopes=scopes, + gcp_log_name=gcp_log_name, + transport_type=transport, + resource=resource, + labels=labels, + ) self.labels: dict[str, str] | None = labels + self.resource: Resource = resource self.task_instance_labels: dict[str, str] | None = {} self.task_instance_hostname = "default-hostname" - @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 - - @property - def _client(self) -> gcp_logging.Client: - """The Cloud Library API client.""" - credentials, project = self._credentials_and_project - client = gcp_logging.Client( - credentials=credentials, - project=project, - client_info=CLIENT_INFO, - ) - return client - - @property - def _logging_service_client(self) -> LoggingServiceV2Client: - """The Cloud logging service v2 client.""" - credentials, _ = self._credentials_and_project - client = LoggingServiceV2Client( - credentials=credentials, - client_info=CLIENT_INFO, - ) - return client - - @cached_property - def _transport(self) -> Transport: - """Object responsible for sending data to Stackdriver.""" - # The Transport object is badly defined (no init) but in the docs client/name as constructor - # arguments are a requirement for any class that derives from Transport class, hence ignore: - return self.transport_type(self._client, self.gcp_log_name) - def _get_labels(self, task_instance=None): if task_instance: - ti_labels = self._task_instance_to_labels(task_instance) + ti_labels = _task_instance_to_labels(task_instance) else: ti_labels = self.task_instance_labels labels: dict[str, str] | None @@ -193,7 +379,7 @@ def emit(self, record: logging.LogRecord) -> None: if not AIRFLOW_V_3_0_PLUS and getattr(record, ctx_indiv_trigger.name, None): ti = getattr(record, "task_instance", None) # trigger context labels = self._get_labels(ti) - self._transport.send(record, message, resource=self.resource, labels=labels) + self.io.transport.send(record, message, resource=self.resource, labels=labels) def set_context(self, task_instance: TaskInstance) -> None: """ @@ -201,7 +387,7 @@ def set_context(self, task_instance: TaskInstance) -> None: :param task_instance: Currently executed task """ - self.task_instance_labels = self._task_instance_to_labels(task_instance) + self.task_instance_labels = _task_instance_to_labels(task_instance) self.task_instance_hostname = task_instance.hostname or "default-hostname" def read( @@ -225,18 +411,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) new_metadata: dict[str, str | bool] = {"end_of_log": end_of_log} @@ -245,102 +431,6 @@ def read( return [((self.task_instance_hostname, messages),)], [new_metadata] - def _prepare_log_filter(self, ti_labels: dict[str, str]) -> str: - """ - Prepare the filter that chooses which log entries to fetch. - - More information: - https://cloud.google.com/logging/docs/reference/v2/rest/v2/entries/list#body.request_body.FIELDS.filter - https://cloud.google.com/logging/docs/view/advanced-queries - - :param ti_labels: Task Instance's labels that will be used to search for logs - :return: logs filter - """ - - def escape_label_key(key: str) -> str: - return f'"{key}"' if "." in key else key - - def escale_label_value(value: str) -> str: - escaped_value = value.replace("\\", "\\\\").replace('"', '\\"') - return f'"{escaped_value}"' - - _, project = self._credentials_and_project - log_filters = [ - f"resource.type={escale_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)}={escale_label_value(value)}") - - for key, value in ti_labels.items(): - log_filters.append(f"labels.{escape_label_key(key)}={escale_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]: - """ - Send requests to the Stackdriver service and downloads logs. - - :param log_filter: Filter specifying the logs to be downloaded. - :param next_page_token: The token of the page from which the log download will start. - If None is passed, it will start from the first page. - :param all_pages: If True is passed, all subpages will be downloaded. Otherwise, only the first - page will be downloaded - :return: A token that contains the following items: - * string with logs - * Boolean value describing whether there are more logs, - * token of the next page - """ - 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 - - end_of_log = True - next_page_token = None - else: - end_of_log = not bool(next_page_token) - return "\n".join(messages), end_of_log, next_page_token - - def _read_single_logs_page(self, log_filter: str, page_token: str | None = None) -> tuple[str, str]: - """ - Send requests to the Stackdriver service and downloads single pages with logs. - - :param log_filter: Filter specifying the logs to be downloaded. - :param page_token: The token of the page to be downloaded. If None is passed, the first page will be - downloaded. - :return: Downloaded logs and next page token - """ - _, project = self._credentials_and_project - request = ListLogEntriesRequest( - resource_names=[f"projects/{project}"], - filter=log_filter, - page_token=page_token, - order_by="timestamp asc", - page_size=1000, - ) - response = self._logging_service_client.list_log_entries(request=request) - page: ListLogEntriesResponse = next(response.pages) - messages: list[str] = [] - for entry in page.entries: - if "message" in (entry.json_payload or {}): - messages.append(entry.json_payload["message"]) # type: ignore - elif entry.text_payload: - messages.append(entry.text_payload) - return "\n".join(messages), page.next_page_token - @classmethod def _task_instance_to_labels(cls, ti: TaskInstance) -> dict[str, str]: return { @@ -375,12 +465,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) url_query_string = { "project": project_id, @@ -393,4 +483,4 @@ def get_external_log_url(self, task_instance: TaskInstance, try_number: int) -> return url def close(self) -> None: - self._transport.flush() + self.io.transport.flush() diff --git a/providers/google/tests/unit/google/cloud/log/test_stackdriver_task_handler.py b/providers/google/tests/unit/google/cloud/log/test_stackdriver_task_handler.py index 66c2ebe29fc86..a55b62fa4f430 100644 --- a/providers/google/tests/unit/google/cloud/log/test_stackdriver_task_handler.py +++ b/providers/google/tests/unit/google/cloud/log/test_stackdriver_task_handler.py @@ -18,14 +18,19 @@ import logging from contextlib import nullcontext +from pathlib import Path from unittest import mock +from unittest.mock import PropertyMock from urllib.parse import parse_qs, urlsplit import pytest from google.cloud.logging import Resource from google.cloud.logging_v2.types import ListLogEntriesRequest, ListLogEntriesResponse, LogEntry -from airflow.providers.google.cloud.log.stackdriver_task_handler import StackdriverTaskHandler +from airflow.providers.google.cloud.log.stackdriver_task_handler import ( + StackdriverRemoteLogIO, + StackdriverTaskHandler, +) from airflow.utils import timezone from airflow.utils.state import TaskInstanceState @@ -50,6 +55,232 @@ def clean_stackdriver_handlers(): del handler +class TestStackdriverRemoteLogIO: + @pytest.fixture(autouse=True) + def _setup(self, tmp_path): + self.local_log_location = str(tmp_path / "local/stackdriver/logs") + self.io = StackdriverRemoteLogIO( + base_log_folder=self.local_log_location, + gcp_key_path="KEY_PATH", + gcp_log_name="airflow", + delete_local_copy=True, + ) + + @mock.patch("airflow.providers.google.cloud.log.stackdriver_task_handler.get_credentials_and_project_id") + @mock.patch("airflow.providers.google.cloud.log.stackdriver_task_handler.LoggingServiceV2Client") + def test_read_logs(self, mock_client, mock_get_creds_and_project_id): + mock_client.return_value.list_log_entries.return_value.pages = iter( + [_create_list_log_entries_response_mock(["MSG1", "MSG2"], None)] + ) + mock_get_creds_and_project_id.return_value = ("creds", "project_id") + + ti = mock.MagicMock() + ti.task_id = "test_task" + ti.dag_id = "test_dag" + ti.try_number = 1 + if AIRFLOW_V_3_0_PLUS: + ti.logical_date = timezone.datetime(2016, 1, 1) + else: + ti.execution_date = timezone.datetime(2016, 1, 1) + + messages, logs = self.io.read("dag_id=test_dag/run_id=run1/task_id=test_task/attempt=1.log", ti) + + assert len(messages) == 1 + assert "Stackdriver" in messages[0] + assert logs == ["MSG1\nMSG2"] + + @mock.patch("airflow.providers.google.cloud.log.stackdriver_task_handler.get_credentials_and_project_id") + @mock.patch("airflow.providers.google.cloud.log.stackdriver_task_handler.LoggingServiceV2Client") + def test_read_logs_empty(self, mock_client, mock_get_creds_and_project_id): + mock_client.return_value.list_log_entries.return_value.pages = iter( + [_create_list_log_entries_response_mock([], None)] + ) + mock_get_creds_and_project_id.return_value = ("creds", "project_id") + + ti = mock.MagicMock() + ti.task_id = "test_task" + ti.dag_id = "test_dag" + ti.try_number = 1 + if AIRFLOW_V_3_0_PLUS: + ti.logical_date = timezone.datetime(2016, 1, 1) + else: + ti.execution_date = timezone.datetime(2016, 1, 1) + + messages, logs = self.io.read("test/path", ti) + + assert len(messages) == 1 + assert logs == [] + + @mock.patch("airflow.providers.google.cloud.log.stackdriver_task_handler.get_credentials_and_project_id") + @mock.patch("airflow.providers.google.cloud.log.stackdriver_task_handler.gcp_logging.Client") + def test_credentials(self, mock_client, mock_get_creds_and_project_id): + mock_get_creds_and_project_id.return_value = ("creds", "project_id") + + _ = self.io._client + + mock_get_creds_and_project_id.assert_called_once_with( + disable_logging=True, + key_path="KEY_PATH", + scopes=frozenset( + { + "https://www.googleapis.com/auth/logging.write", + "https://www.googleapis.com/auth/logging.read", + } + ), + ) + mock_client.assert_called_once_with(credentials="creds", client_info=mock.ANY, project="project_id") + + @mock.patch("airflow.providers.google.cloud.log.stackdriver_task_handler.get_credentials_and_project_id") + @mock.patch("airflow.providers.google.cloud.log.stackdriver_task_handler.gcp_logging.Client") + def test_transport_init(self, mock_client, mock_get_creds_and_project_id): + mock_get_creds_and_project_id.return_value = ("creds", "project_id") + + transport_type = mock.MagicMock() + io = StackdriverRemoteLogIO( + base_log_folder=self.local_log_location, + gcp_log_name="test-log", + transport_type=transport_type, + ) + _ = io.transport + transport_type.assert_called_once_with(mock_client.return_value, "test-log") + + @mock.patch("shutil.rmtree") + @mock.patch( + "airflow.providers.google.cloud.log.stackdriver_task_handler.StackdriverRemoteLogIO.transport", + new_callable=PropertyMock, + ) + def test_upload_flushes_transport_and_deletes_local(self, mock_transport_prop, mock_rmtree): + io = StackdriverRemoteLogIO( + base_log_folder=self.local_log_location, + gcp_log_name="airflow", + delete_local_copy=True, + ) + mock_transport = mock.MagicMock() + mock_transport_prop.return_value = mock_transport + + base = Path(self.local_log_location) + base.mkdir(parents=True, exist_ok=True) + log_dir = base / "subdir" + log_dir.mkdir(parents=True, exist_ok=True) + log_file = log_dir / "test.log" + log_file.write_text("log content") + + ti = mock.MagicMock() + io.upload(str(log_file), ti) + + mock_transport.flush.assert_called_once() + mock_rmtree.assert_called_once_with(log_dir.resolve(), ignore_errors=True) + + @mock.patch( + "airflow.providers.google.cloud.log.stackdriver_task_handler.StackdriverRemoteLogIO.transport", + new_callable=PropertyMock, + ) + def test_upload_no_delete(self, mock_transport_prop): + io = StackdriverRemoteLogIO( + base_log_folder=self.local_log_location, + gcp_log_name="airflow", + delete_local_copy=False, + ) + mock_transport = mock.MagicMock() + mock_transport_prop.return_value = mock_transport + + ti = mock.MagicMock() + io.upload("some/path.log", ti) + + mock_transport.flush.assert_called_once() + + @mock.patch("airflow.providers.google.cloud.log.stackdriver_task_handler.get_credentials_and_project_id") + def test_prepare_log_filter(self, mock_get_creds_and_project_id): + mock_get_creds_and_project_id.return_value = ("creds", "project_id") + + ti_labels = { + "task_id": "test_task", + "dag_id": "test_dag", + "try_number": "1", + } + log_filter = self.io.prepare_log_filter(ti_labels) + + assert 'resource.type="global"' in log_filter + assert 'logName="projects/project_id/logs/airflow"' in log_filter + assert 'labels.task_id="test_task"' in log_filter + assert 'labels.dag_id="test_dag"' in log_filter + + @mock.patch("airflow.providers.google.cloud.log.stackdriver_task_handler.get_credentials_and_project_id") + def test_prepare_log_filter_with_custom_resource(self, mock_get_creds_and_project_id): + mock_get_creds_and_project_id.return_value = ("creds", "project_id") + + io = StackdriverRemoteLogIO( + base_log_folder=self.local_log_location, + gcp_log_name="airflow", + resource=Resource( + type="cloud_composer_environment", + labels={ + "environment.name": "test-instance", + "location": "europe-west-3", + }, + ), + ) + log_filter = io.prepare_log_filter({"task_id": "test"}) + + assert 'resource.type="cloud_composer_environment"' in log_filter + assert 'resource.labels."environment.name"="test-instance"' in log_filter + assert 'resource.labels.location="europe-west-3"' in log_filter + + @pytest.mark.skipif(not AIRFLOW_V_3_0_PLUS, reason="airflow.sdk.log only exists in Airflow 3+") + @mock.patch("airflow.providers.google.cloud.log.stackdriver_task_handler.get_credentials_and_project_id") + @mock.patch("airflow.providers.google.cloud.log.stackdriver_task_handler.gcp_logging.Client") + def test_processors_sends_to_transport(self, mock_client, mock_get_creds_and_project_id): + mock_get_creds_and_project_id.return_value = ("creds", "project_id") + + mock_transport_type = mock.MagicMock() + with mock.patch("airflow.sdk.log.relative_path_from_logger", return_value="dag/task/1.log"): + io = StackdriverRemoteLogIO( + base_log_folder=self.local_log_location, + gcp_log_name="airflow", + labels={"env": "test"}, + transport_type=mock_transport_type, + ) + processors = io.processors + assert len(processors) == 1 + + proc = processors[0] + mock_logger = mock.MagicMock() + + event = { + "event": "hello world", + "logger_name": "airflow.task", + "timestamp": "2026-01-15T10:30:00+00:00", + } + result = proc(mock_logger, "info", event) + + assert result is event + mock_transport = mock_transport_type.return_value + mock_transport.send.assert_called_once() + record = mock_transport.send.call_args[0][0] + assert record.levelno == logging.INFO + + @pytest.mark.skipif(not AIRFLOW_V_3_0_PLUS, reason="airflow.sdk.log only exists in Airflow 3+") + @mock.patch("airflow.providers.google.cloud.log.stackdriver_task_handler.get_credentials_and_project_id") + @mock.patch("airflow.providers.google.cloud.log.stackdriver_task_handler.gcp_logging.Client") + def test_processors_skips_non_task_logger(self, mock_client, mock_get_creds_and_project_id): + mock_get_creds_and_project_id.return_value = ("creds", "project_id") + + mock_transport_type = mock.MagicMock() + with mock.patch("airflow.sdk.log.relative_path_from_logger", return_value=None): + io = StackdriverRemoteLogIO( + base_log_folder=self.local_log_location, + gcp_log_name="airflow", + transport_type=mock_transport_type, + ) + proc = io.processors[0] + + event = {"event": "should not be sent"} + result = proc(mock.MagicMock(), "info", event) + + assert result is event + mock_transport_type.return_value.send.assert_not_called() + + @pytest.mark.usefixtures("clean_stackdriver_handlers") @mock.patch("airflow.providers.google.cloud.log.stackdriver_task_handler.get_credentials_and_project_id") @mock.patch("airflow.providers.google.cloud.log.stackdriver_task_handler.gcp_logging.Client") @@ -77,7 +308,6 @@ def test_should_pass_message_to_client(mock_client, mock_get_creds_and_project_i @mock.patch("airflow.providers.google.cloud.log.stackdriver_task_handler.gcp_logging.Client") def test_should_use_configured_log_name(mock_client, mock_get_creds_and_project_id): import importlib - import logging from airflow import settings from airflow.config_templates import airflow_local_settings @@ -85,9 +315,6 @@ def test_should_use_configured_log_name(mock_client, mock_get_creds_and_project_ mock_get_creds_and_project_id.return_value = ("creds", "project_id") try: - # this is needed for Airflow 2.8 and below where default settings are triggering warning on - # extra "name" in the configuration of stackdriver handler. As of Airflow 2.9 this warning is not - # emitted. context_manager = nullcontext() with context_manager: with conf_vars( @@ -99,12 +326,17 @@ def test_should_use_configured_log_name(mock_client, mock_get_creds_and_project_ importlib.reload(airflow_local_settings) settings.configure_logging() + task_log = getattr(airflow_local_settings, "REMOTE_TASK_LOG", None) + if task_log is not None: + # Airflow 3+ uses REMOTE_TASK_LOG instead of handler-based config + assert isinstance(task_log, StackdriverRemoteLogIO) + assert task_log.gcp_log_name == "path" + return + + # Older Airflow: stackdriver is wired as a logging handler logger = logging.getLogger("airflow.task") handler = logger.handlers[0] assert isinstance(handler, StackdriverTaskHandler) - with mock.patch.object(handler, "transport_type") as transport_type_mock: - logger.error("foo") - transport_type_mock.assert_called_once_with(mock_client.return_value, "path") finally: importlib.reload(airflow_local_settings) settings.configure_logging() @@ -398,7 +630,7 @@ def test_should_use_credentials(self, mock_client, mock_get_creds_and_project_id mock_get_creds_and_project_id.return_value = ("creds", "project_id") stackdriver_task_handler = StackdriverTaskHandler(gcp_key_path="KEY_PATH") - client = stackdriver_task_handler._client + client = stackdriver_task_handler.io._client mock_get_creds_and_project_id.assert_called_once_with( disable_logging=True,