Skip to content

Commit 85ec4b8

Browse files
goingforstudying-ctrlgoingforstudying-ctrl
authored andcommitted
fix(google): extract Cloud Logging labels from AF3 log path
In Airflow 3 the supervisor process runs REMOTE_TASK_LOG handlers, but record.task_instance is never set in supervisor context (it is a task-subprocess concept). When task_instance is missing the proc() closure shipped log entries with empty labels, making Cloud Logging entries unsearchable by dag_id / task_id. Parse dag_id, task_id, and try_number from the structured AF3 log path (dag_id=<x>/run_id=<x>/task_id=<x>/attempt=<N>.log) instead. This requires zero DB access and works regardless of whether the handler runs in a task subprocess or the supervisor. relates to #68240
1 parent 0385bae commit 85ec4b8

2 files changed

Lines changed: 85 additions & 2 deletions

File tree

providers/google/src/airflow/providers/google/cloud/log/stackdriver_task_handler.py

Lines changed: 39 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -139,7 +139,8 @@ def proc(
139139
method_name: str,
140140
event: structlog.typing.EventDict,
141141
):
142-
if not logger or not relative_path_from_logger(logger):
142+
relative = relative_path_from_logger(logger) if logger else None
143+
if not logger or not relative:
143144
return event
144145

145146
name = event.get("logger_name") or event.get("logger", "")
@@ -182,7 +183,10 @@ def proc(
182183
labels[LABEL_TRY_NUMBER] = str(try_number)
183184
if map_index := event.get("map_index"):
184185
labels["map_index"] = str(map_index)
185-
186+
# In AF3 supervisor context record.task_instance is not set.
187+
# Parse labels from the structured log path as additional fallback.
188+
path_labels = _labels_from_path(str(relative))
189+
labels.update(path_labels)
186190
_transport.send(record, str(msg.get("event", "")), resource=self.resource, labels=labels)
187191
return event
188192

@@ -278,6 +282,39 @@ def _read_single_logs_page(self, log_filter: str, page_token: str | None = None)
278282
return "\n".join(messages), page.next_page_token
279283

280284

285+
def _labels_from_path(relative_path: str) -> dict[str, str]:
286+
"""Parse AF3 log path into Stackdriver labels.
287+
288+
AF3's log path template is::
289+
290+
dag_id=<x>/run_id=<x>/task_id=<x>/attempt=<N>.log
291+
292+
All four label fields are extracted with zero DB access. When the path
293+
does not match the expected format the function returns an empty dict
294+
so callers can fall back to other label sources.
295+
"""
296+
# Strip the trailing file extension (.log) and split into segments
297+
stem = relative_path.rsplit(".", 1)[0] if "." in relative_path else relative_path
298+
segments = stem.split("/")
299+
labels: dict[str, str] = {}
300+
for segment in segments:
301+
if "=" not in segment:
302+
continue
303+
key, _, value = segment.partition("=")
304+
if key == "dag_id":
305+
labels[LABEL_DAG_ID] = value
306+
elif key == "task_id":
307+
labels[LABEL_TASK_ID] = value
308+
elif key == "attempt":
309+
labels[LABEL_TRY_NUMBER] = value
310+
elif key == "run_id":
311+
# run_id is NOT a standard Stackdriver label yet, but it is used
312+
# on the write side via the log path template. Store it so the
313+
# read path can filter on it (Bug 2 will wire this up).
314+
pass
315+
return labels
316+
317+
281318
def _task_instance_to_labels(ti) -> dict[str, str]:
282319
"""Convert a task instance to Stackdriver labels."""
283320
return {

providers/google/tests/unit/google/cloud/log/test_stackdriver_task_handler.py

Lines changed: 46 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -325,6 +325,52 @@ def test_processors_skips_non_task_logger(self, mock_client, mock_get_creds_and_
325325
assert result is event
326326
mock_transport_type.return_value.send.assert_not_called()
327327

328+
@pytest.mark.skipif(not AIRFLOW_V_3_0_PLUS, reason="airflow.sdk.log only exists in Airflow 3+")
329+
@mock.patch("airflow.providers.google.cloud.log.stackdriver_task_handler.get_credentials_and_project_id")
330+
@mock.patch("airflow.providers.google.cloud.log.stackdriver_task_handler.gcp_logging.Client")
331+
def test_processors_extracts_labels_from_path_in_supervisor_context(
332+
self, mock_client, mock_get_creds_and_project_id
333+
):
334+
"""Labels are parsed from the AF3 log path when task_instance is not set.
335+
336+
In AF3 supervisor context ``record.task_instance`` is None because it is
337+
a task-subprocess concept. The handler should fall back to extracting
338+
``dag_id``, ``task_id``, and ``try_number`` from the structured log path.
339+
"""
340+
mock_get_creds_and_project_id.return_value = ("creds", "project_id")
341+
342+
mock_transport_type = mock.MagicMock()
343+
af3_path = "dag_id=my_dag/run_id=scheduled__2026-06-08T00:00:00+00:00/task_id=print_date/attempt=3.log"
344+
with mock.patch("airflow.sdk.log.relative_path_from_logger", return_value=af3_path):
345+
io = StackdriverRemoteLogIO(
346+
base_log_folder=self.local_log_location,
347+
gcp_log_name="airflow",
348+
transport_type=mock_transport_type,
349+
)
350+
proc = io.processors[0]
351+
352+
event = {
353+
"event": "task log line",
354+
"logger_name": "airflow.task",
355+
"timestamp": "2026-06-08T10:30:00+00:00",
356+
}
357+
proc(mock.MagicMock(), "info", event)
358+
359+
mock_transport = mock_transport_type.return_value
360+
mock_transport.send.assert_called_once()
361+
labels = mock_transport.send.call_args[1]["labels"]
362+
assert labels["dag_id"] == "my_dag"
363+
assert labels["task_id"] == "print_date"
364+
assert labels["try_number"] == "3"
365+
366+
def test_labels_from_path_returns_empty_on_unexpected_format(self):
367+
"""_labels_from_path returns an empty dict when the path doesn't match."""
368+
from airflow.providers.google.cloud.log.stackdriver_task_handler import _labels_from_path
369+
370+
assert _labels_from_path("") == {}
371+
assert _labels_from_path("no_equals_here.log") == {}
372+
assert _labels_from_path("random/path/without/keys.log") == {}
373+
328374

329375
@pytest.mark.usefixtures("clean_stackdriver_handlers")
330376
@mock.patch("airflow.providers.google.cloud.log.stackdriver_task_handler.get_credentials_and_project_id")

0 commit comments

Comments
 (0)