From 2273d89e9e5a218417bd5727a0b9cb0be21b8ed2 Mon Sep 17 00:00:00 2001 From: LIU ZHE YOU Date: Thu, 9 Jul 2026 16:27:42 +0000 Subject: [PATCH 1/5] Make ES and OpenSearch log handlers respect the per-run log ID template The LogTemplate model pins the log_id template to each Dag run so task logs written before a config change stay readable, but the Airflow 3 paths of the Elasticsearch and OpenSearch handlers rendered log_id from the current conf value on both the write (worker) and read (API server) sides, so changing [elasticsearch/opensearch] log_id_template orphaned all previously written logs. Workers cannot query the metadata DB, so the pinned template is delivered to the supervisor through a new optional TIRunContext field populated by the ti_run Execution API endpoint (with a version migration for older clients); the API-server read paths restore the per-run lookup directly. Handlers that do not opt into the new ti_context keyword keep their existing upload signature and behavior. --- .../execution_api/datamodels/taskinstance.py | 6 ++ .../execution_api/routes/task_instances.py | 5 +- .../execution_api/versions/__init__.py | 5 ++ .../execution_api/versions/v2026_09_30.py | 40 +++++++++ .../versions/head/test_task_instances.py | 38 +++++++++ .../versions/v2026_09_30/__init__.py | 16 ++++ .../v2026_09_30/test_task_instances.py | 85 +++++++++++++++++++ .../elasticsearch/log/es_task_handler.py | 85 +++++++++++++++++-- .../elasticsearch/log/test_es_task_handler.py | 72 +++++++++++++++- .../opensearch/log/os_task_handler.py | 83 ++++++++++++++++-- .../opensearch/log/test_os_task_handler.py | 48 +++++++++++ .../airflow/sdk/api/datamodels/_generated.py | 3 +- .../sdk/execution_time/schema/schema.json | 14 ++- .../schema/versions/__init__.py | 3 + .../schema/versions/v2026_09_30.py | 30 +++++++ .../airflow/sdk/execution_time/supervisor.py | 7 +- task-sdk/src/airflow/sdk/log.py | 13 ++- .../execution_time/test_supervisor.py | 47 +++++++++- task-sdk/tests/task_sdk/test_log.py | 74 ++++++++++++++++ 19 files changed, 653 insertions(+), 21 deletions(-) create mode 100644 airflow-core/src/airflow/api_fastapi/execution_api/versions/v2026_09_30.py create mode 100644 airflow-core/tests/unit/api_fastapi/execution_api/versions/v2026_09_30/__init__.py create mode 100644 airflow-core/tests/unit/api_fastapi/execution_api/versions/v2026_09_30/test_task_instances.py create mode 100644 task-sdk/src/airflow/sdk/execution_time/schema/versions/v2026_09_30.py create mode 100644 task-sdk/tests/task_sdk/test_log.py diff --git a/airflow-core/src/airflow/api_fastapi/execution_api/datamodels/taskinstance.py b/airflow-core/src/airflow/api_fastapi/execution_api/datamodels/taskinstance.py index ad051b3e6d340..300d14b1a6bfb 100644 --- a/airflow-core/src/airflow/api_fastapi/execution_api/datamodels/taskinstance.py +++ b/airflow-core/src/airflow/api_fastapi/execution_api/datamodels/taskinstance.py @@ -435,6 +435,12 @@ class TIRunContext(BaseModel): always reflects when the task *first* started, not when it was rescheduled/resumed. """ + log_id_template: str | None = None + """ + Elasticsearch/OpenSearch log id template pinned to this Dag run (``LogTemplate.elasticsearch_id``), + so remote log writers use the template that was in effect when the run was created. + """ + class PrevSuccessfulDagRunResponse(BaseModel): """Schema for response with previous successful DagRun information for Task Template Context.""" diff --git a/airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py b/airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py index c1bac7960234d..47815da5702a4 100644 --- a/airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py +++ b/airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py @@ -76,7 +76,7 @@ require_auth, ) from airflow.configuration import conf -from airflow.exceptions import InvalidPartitionKeyError, TaskNotFound +from airflow.exceptions import AirflowException, InvalidPartitionKeyError, TaskNotFound from airflow.models.asset import AssetActive from airflow.models.base import ID_LEN from airflow.models.dag import DagModel @@ -310,6 +310,9 @@ def ti_run( should_retry=_is_eligible_to_retry(previous_state, ti.try_number, ti.max_tries), ) + with contextlib.suppress(AirflowException): # no LogTemplate row: worker falls back to conf + context.log_id_template = dr.get_log_template(session=session).elasticsearch_id + # Only set if they are non-null if ti.next_method: context.next_method = ti.next_method diff --git a/airflow-core/src/airflow/api_fastapi/execution_api/versions/__init__.py b/airflow-core/src/airflow/api_fastapi/execution_api/versions/__init__.py index dc7035d31e3c9..b9aaa0ea02880 100644 --- a/airflow-core/src/airflow/api_fastapi/execution_api/versions/__init__.py +++ b/airflow-core/src/airflow/api_fastapi/execution_api/versions/__init__.py @@ -51,9 +51,14 @@ AddTeamNameField, AddVariableKeysEndpoint, ) +from airflow.api_fastapi.execution_api.versions.v2026_09_30 import AddLogIdTemplateField bundle = VersionBundle( HeadVersion(), + Version( + "2026-09-30", + AddLogIdTemplateField, + ), Version( "2026-06-30", AddVariableKeysEndpoint, diff --git a/airflow-core/src/airflow/api_fastapi/execution_api/versions/v2026_09_30.py b/airflow-core/src/airflow/api_fastapi/execution_api/versions/v2026_09_30.py new file mode 100644 index 0000000000000..40e4bf3e62299 --- /dev/null +++ b/airflow-core/src/airflow/api_fastapi/execution_api/versions/v2026_09_30.py @@ -0,0 +1,40 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +from __future__ import annotations + +from cadwyn import ( + ResponseInfo, + VersionChange, + convert_response_to_previous_version_for, + schema, +) + +from airflow.api_fastapi.execution_api.datamodels.taskinstance import TIRunContext + + +class AddLogIdTemplateField(VersionChange): + """Add the Dag-run-pinned `log_id_template` field to TIRunContext.""" + + description = __doc__ + + instructions_to_migrate_to_previous_version = (schema(TIRunContext).field("log_id_template").didnt_exist,) + + @convert_response_to_previous_version_for(TIRunContext) # type: ignore[arg-type] + def remove_log_id_template_field(response: ResponseInfo) -> None: # type: ignore[misc] + """Remove log_id_template field for older API versions.""" + response.body.pop("log_id_template", None) diff --git a/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_task_instances.py b/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_task_instances.py index 542ce7eaaf15b..d76cb3f9418e7 100644 --- a/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_task_instances.py +++ b/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_task_instances.py @@ -278,6 +278,7 @@ def test_ti_run_state_to_running( "variables": [], "connections": [], "xcom_keys_to_clear": [], + "log_id_template": "{dag_id}-{task_id}-{run_id}-{map_index}-{try_number}", } # upstream_map_indexes is now computed by Task SDK, not returned by the server in HEAD version assert "upstream_map_indexes" not in result @@ -692,6 +693,7 @@ def test_next_kwargs_still_encoded(self, client, session, create_task_instance, "next_method": "execute_complete", "next_kwargs": expected_next_kwargs, "start_date": None, + "log_id_template": "{dag_id}-{task_id}-{run_id}-{map_index}-{try_number}", } @pytest.mark.parametrize("resume", [True, False]) @@ -766,6 +768,7 @@ def test_next_kwargs_determines_start_date_update(self, client, session, create_ "xcom_keys_to_clear": [], "next_method": "execute_complete", "next_kwargs": expected_next_kwargs, + "log_id_template": "{dag_id}-{task_id}-{run_id}-{map_index}-{try_number}", } session.expunge_all() ti = session.get(TaskInstance, ti.id) @@ -1030,6 +1033,41 @@ def test_ti_run_populates_team_name( assert response.status_code == 200 assert response.json()["dag_run"]["team_name"] == (team_name if expect_team else None) + def test_ti_run_populates_log_id_template( + self, client, session, create_task_instance, create_log_template, time_machine + ): + """``log_id_template`` echoes the LogTemplate row pinned to the Dag run.""" + instant = timezone.parse("2024-09-30T12:00:00Z") + time_machine.move_to(instant, tick=False) + + custom_elasticsearch_id = "{dag_id}-{task_id}-{logical_date}-{try_number}" + create_log_template("{{ ti.dag_id }}.log", custom_elasticsearch_id) + + # The Dag run pins the newest LogTemplate row at creation time. + ti = create_task_instance( + task_id="test_ti_run_populates_log_id_template", + state=State.QUEUED, + dagrun_state=DagRunState.RUNNING, + session=session, + start_date=instant, + dag_id=str(uuid4()), + ) + session.commit() + + response = client.patch( + f"/execution/task-instances/{ti.id}/run", + json={ + "state": "running", + "hostname": "random-hostname", + "unixname": "random-unixname", + "pid": 100, + "start_date": instant.isoformat(), + }, + ) + + assert response.status_code == 200 + assert response.json()["log_id_template"] == custom_elasticsearch_id + def test_ti_run_creates_audit_log(self, client, session, create_task_instance, time_machine): """Test that transitioning to RUNNING creates an audit log record.""" instant_str = "2024-09-30T12:00:00Z" diff --git a/airflow-core/tests/unit/api_fastapi/execution_api/versions/v2026_09_30/__init__.py b/airflow-core/tests/unit/api_fastapi/execution_api/versions/v2026_09_30/__init__.py new file mode 100644 index 0000000000000..13a83393a9124 --- /dev/null +++ b/airflow-core/tests/unit/api_fastapi/execution_api/versions/v2026_09_30/__init__.py @@ -0,0 +1,16 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. diff --git a/airflow-core/tests/unit/api_fastapi/execution_api/versions/v2026_09_30/test_task_instances.py b/airflow-core/tests/unit/api_fastapi/execution_api/versions/v2026_09_30/test_task_instances.py new file mode 100644 index 0000000000000..39d1f4e0372d5 --- /dev/null +++ b/airflow-core/tests/unit/api_fastapi/execution_api/versions/v2026_09_30/test_task_instances.py @@ -0,0 +1,85 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +from __future__ import annotations + +import pytest + +from airflow._shared.timezones import timezone +from airflow.utils.state import DagRunState, State + +from tests_common.test_utils.db import clear_db_runs + +pytestmark = pytest.mark.db_test + +TIMESTAMP_STR = "2024-09-30T12:00:00Z" +TIMESTAMP = timezone.parse(TIMESTAMP_STR) + +RUN_PATCH_BODY = { + "state": "running", + "hostname": "h", + "unixname": "u", + "pid": 1, + "start_date": TIMESTAMP_STR, +} + + +@pytest.fixture +def old_ver_client(client): + """Execution API version immediately before ``log_id_template`` was added.""" + client.headers["Airflow-API-Version"] = "2026-06-30" + return client + + +class TestLogIdTemplateFieldBackwardCompat: + @pytest.fixture(autouse=True) + def _freeze_time(self, time_machine): + time_machine.move_to(TIMESTAMP_STR, tick=False) + + def setup_method(self): + clear_db_runs() + + def teardown_method(self): + clear_db_runs() + + def test_old_version_strips_log_id_template(self, old_ver_client, session, create_task_instance): + ti = create_task_instance( + task_id="test_log_id_template_downgrade", + state=State.QUEUED, + dagrun_state=DagRunState.RUNNING, + session=session, + start_date=TIMESTAMP, + ) + session.commit() + + response = old_ver_client.patch(f"/execution/task-instances/{ti.id}/run", json=RUN_PATCH_BODY) + assert response.status_code == 200 + assert "log_id_template" not in response.json() + + def test_head_version_includes_log_id_template(self, client, session, create_task_instance): + ti = create_task_instance( + task_id="test_log_id_template_head", + state=State.QUEUED, + dagrun_state=DagRunState.RUNNING, + session=session, + start_date=TIMESTAMP, + ) + session.commit() + + response = client.patch(f"/execution/task-instances/{ti.id}/run", json=RUN_PATCH_BODY) + assert response.status_code == 200 + assert response.json()["log_id_template"] == "{dag_id}-{task_id}-{run_id}-{map_index}-{try_number}" diff --git a/providers/elasticsearch/src/airflow/providers/elasticsearch/log/es_task_handler.py b/providers/elasticsearch/src/airflow/providers/elasticsearch/log/es_task_handler.py index fa05b2a0b34cf..8579acd8f9f9e 100644 --- a/providers/elasticsearch/src/airflow/providers/elasticsearch/log/es_task_handler.py +++ b/providers/elasticsearch/src/airflow/providers/elasticsearch/log/es_task_handler.py @@ -51,6 +51,7 @@ from airflow.providers.elasticsearch.version_compat import AIRFLOW_V_3_0_PLUS, AIRFLOW_V_3_2_PLUS from airflow.utils.log.file_task_handler import FileTaskHandler from airflow.utils.log.logging_mixin import ExternalLoggingMixin, LoggingMixin +from airflow.utils.session import create_session from airflow.utils.state import TaskInstanceState if AIRFLOW_V_3_2_PLUS: @@ -76,6 +77,8 @@ EsLogMsgType = list[tuple[str, str]] # type: ignore[assignment,misc] +log = logging.getLogger(__name__) + LOG_LINE_DEFAULTS = {"exc_text": "", "stack_info": ""} # Elasticsearch hosted log type @@ -192,16 +195,59 @@ def _strip_userinfo(url: str) -> str: return parsed._replace(netloc=netloc).geturl() -def _render_log_id(log_id_template: str, ti: TaskInstance | TaskInstanceKey, try_number: int) -> str: +def _render_log_id( + log_id_template: str, + ti: TaskInstance | TaskInstanceKey | RuntimeTI, + try_number: int, + *, + dag_run: Any = None, + json_format: bool = False, +) -> str: + # Templates pinned by older LogTemplate rows may use date placeholders (e.g. the pre-2.3 + # default `{dag_id}-{task_id}-{logical_date}-{try_number}`), and str.format raises KeyError + # on any missing name, so always supply the full set. + def format_date(value: datetime | None) -> str: + if json_format: + return _clean_date(value) + return value.isoformat() if value else "" + + logical_date = format_date(getattr(dag_run, "logical_date", None)) return log_id_template.format( dag_id=ti.dag_id, task_id=ti.task_id, run_id=getattr(ti, "run_id", ""), try_number=try_number, map_index=getattr(ti, "map_index", ""), + logical_date=logical_date, + execution_date=logical_date, + data_interval_start=format_date(getattr(dag_run, "data_interval_start", None)), + data_interval_end=format_date(getattr(dag_run, "data_interval_end", None)), ) +def _resolve_log_id_template( + ti: TaskInstance | TaskInstanceKey | RuntimeTI, default_template: str +) -> tuple[str, DagRun | None]: + """ + Fetch the Dag-run-pinned log id template and the Dag run from the metadata DB. + + Falls back to the configured template when the DB is not reachable (e.g. on workers, which + receive the pinned template through ``TIRunContext`` instead). + """ + try: + with create_session() as session: + dag_run = ti.get_dagrun(session=session) # type: ignore[union-attr] + if USE_PER_RUN_LOG_ID: + return dag_run.get_log_template(session=session).elasticsearch_id, dag_run + return default_template, dag_run + except Exception: + log.warning( + "Could not fetch the Dag-run-pinned log id template; falling back to the configured one", + exc_info=True, + ) + return default_template, None + + def _clean_date(value: datetime | None) -> str: """ Clean up a date value so that it is safe to query in elasticsearch by removing reserved characters. @@ -377,6 +423,10 @@ def _get_es_source_includes(self) -> list[str]: extra_fields = self.json_fields if self.json_format and not AIRFLOW_V_3_0_PLUS else [] return self.io._get_source_includes(extra_fields=extra_fields) + def _render_log_id(self, ti: TaskInstance | TaskInstanceKey, try_number: int) -> str: + log_id_template, dag_run = _resolve_log_id_template(ti, self.log_id_template) + return _render_log_id(log_id_template, ti, try_number, dag_run=dag_run, json_format=self.json_format) + def _read( self, ti: TaskInstance, try_number: int, metadata: LogMetadata | None = None ) -> tuple[EsLogMsgType, LogMetadata]: @@ -407,7 +457,7 @@ def _read( metadata["offset"] = 0 offset = metadata["offset"] - log_id = _render_log_id(self.log_id_template, ti, try_number) + log_id = self._render_log_id(ti, try_number) response = self.io._es_read(log_id, offset, ti, source_includes=self._get_es_source_includes()) # TODO: Can we skip group logs by host ? if response is not None and response.hits: @@ -529,7 +579,7 @@ def set_context(self, ti: TaskInstance, *, identifier: str | None = None) -> Non _clean_date(ti.logical_date) if AIRFLOW_V_3_0_PLUS else _clean_date(ti.execution_date) ), "try_number": str(ti.try_number), - "log_id": _render_log_id(self.log_id_template, ti, ti.try_number), + "log_id": self._render_log_id(ti, ti.try_number), }, ) @@ -595,7 +645,7 @@ 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 """ - log_id = _render_log_id(self.log_id_template, task_instance, try_number) + log_id = self._render_log_id(task_instance, try_number) scheme = "" if "://" in self.frontend else "https://" return scheme + self.frontend.format(log_id=quote(log_id)) @@ -702,8 +752,13 @@ def _get_source_includes(self, extra_fields: Iterable[str] = ()) -> list[str]: ["@timestamp", *TASK_LOG_FIELDS, self.host_field, self.offset_field, *extra_fields] ) - def upload(self, path: os.PathLike | str, ti: RuntimeTI | None = None) -> None: - """Write the log to ElasticSearch.""" + def upload(self, path: os.PathLike | str, ti: RuntimeTI | None = None, *, ti_context: Any = None) -> None: + """ + Write the log to ElasticSearch. + + :param ti_context: the ``TIRunContext`` the supervisor received at task start, carrying the + Dag-run-pinned log id template; when absent, fall back to the configured template. + """ if ti is None: return @@ -714,7 +769,14 @@ def upload(self, path: os.PathLike | str, ti: RuntimeTI | None = None) -> None: else: local_loc = self.base_log_folder.joinpath(path) - log_id = _render_log_id(self.log_id_template, ti, ti.try_number) # type: ignore[arg-type] + log_id_template = getattr(ti_context, "log_id_template", None) or self.log_id_template + log_id = _render_log_id( + log_id_template, + ti, + ti.try_number, + dag_run=getattr(ti_context, "dag_run", None), + json_format=self.json_format, + ) if local_loc.is_file() and self.write_stdout: # Intentionally construct the log_id and offset field @@ -784,7 +846,14 @@ def _write_to_es(self, log_lines: list[dict[str, Any]]) -> bool: return False def read(self, _relative_path: str, ti: RuntimeTI) -> tuple[LogSourceInfo, LogMessages]: - log_id = _render_log_id(self.log_id_template, ti, ti.try_number) # type: ignore[arg-type] + log_id_template, dag_run = _resolve_log_id_template(ti, self.log_id_template) + log_id = _render_log_id( + log_id_template, + ti, + ti.try_number, + dag_run=dag_run, + json_format=self.json_format, + ) self.log.info("Reading log %s from Elasticsearch", log_id) offset = 0 response = self._es_read(log_id, offset, ti, source_includes=self._get_source_includes()) diff --git a/providers/elasticsearch/tests/unit/elasticsearch/log/test_es_task_handler.py b/providers/elasticsearch/tests/unit/elasticsearch/log/test_es_task_handler.py index b923d8b7c6f00..e1ea981fe0ea7 100644 --- a/providers/elasticsearch/tests/unit/elasticsearch/log/test_es_task_handler.py +++ b/providers/elasticsearch/tests/unit/elasticsearch/log/test_es_task_handler.py @@ -23,6 +23,7 @@ import re from io import StringIO from pathlib import Path +from types import SimpleNamespace from unittest import mock from unittest.mock import Mock, patch from urllib.parse import quote @@ -490,6 +491,22 @@ def test_close_reopens_stream(self, ti, stream_state): def test_render_log_id(self, ti): assert _render_log_id(self.es_task_handler.log_id_template, ti, 1) == self.LOG_ID + @pytest.mark.db_test + def test_render_log_id_uses_pinned_template(self, ti): + """The handler renders with the LogTemplate row pinned to the Dag run, not the conf value.""" + logical_date = ti.get_dagrun().logical_date + expected = f"{self.DAG_ID}-{self.TASK_ID}-{logical_date.isoformat()}-{self.TRY_NUM}" + assert self.es_task_handler._render_log_id(ti, self.TRY_NUM) == expected + + @pytest.mark.db_test + def test_remote_log_io_read_uses_pinned_template(self, ti): + with mock.patch.object(self.es_task_handler.io, "_es_read", return_value=None) as mock_es_read: + self.es_task_handler.io.read("", ti) + + logical_date = ti.get_dagrun().logical_date + expected = f"{self.DAG_ID}-{self.TASK_ID}-{logical_date.isoformat()}-{self.TRY_NUM}" + mock_es_read.assert_called_once_with(expected, 0, ti, source_includes=mock.ANY) + def test_clean_date(self): clean_logical_date = _clean_date(datetime(2016, 7, 8, 9, 10, 11, 12)) assert clean_logical_date == "2016_07_08T09_10_11_000012" @@ -517,8 +534,13 @@ def test_get_external_log_url(self, ti, json_format, es_frontend, expected_url): offset_field=self.offset_field, frontend=es_frontend, ) + # The URL renders with the log id template pinned to the Dag run by the ``ti`` fixture + # (`{dag_id}-{task_id}-{logical_date}-{try_number}`), not the configured template. + logical_date = ti.get_dagrun().logical_date + expected_date = _clean_date(logical_date) if json_format else logical_date.isoformat() + expected_log_id = f"{self.DAG_ID}-{self.TASK_ID}-{expected_date}-{self.TRY_NUM}" assert es_task_handler.get_external_log_url(ti, ti.try_number) == expected_url.format( - log_id=quote(self.LOG_ID) + log_id=quote(expected_log_id) ) @pytest.mark.parametrize( @@ -856,6 +878,54 @@ def test_read_returns_missing_log_message_when_es_read_returns_none(self, ti): def test_upload_returns_early_when_ti_is_none(self, tmp_json_file): self.elasticsearch_io.upload(tmp_json_file, ti=None) + def test_upload_uses_pinned_log_id_template_from_ti_context(self, tmp_json_file, ti): + """Documents are stamped with the Dag-run-pinned template delivered via TIRunContext.""" + ti_context = SimpleNamespace( + log_id_template="{dag_id}-{task_id}-{logical_date}-{try_number}", + dag_run=SimpleNamespace( + logical_date=datetime(2016, 1, 1), + data_interval_start=None, + data_interval_end=None, + ), + ) + + with patch.object(self.elasticsearch_io, "_write_to_es", return_value=False) as mock_write: + self.elasticsearch_io.upload(tmp_json_file, ti, ti_context=ti_context) + + expected_log_id = f"{ti.dag_id}-{ti.task_id}-2016-01-01T00:00:00+00:00-{ti.try_number}" + log_lines = mock_write.call_args.args[0] + assert log_lines + assert all(line["log_id"] == expected_log_id for line in log_lines) + + def test_upload_stamps_pinned_log_id_on_stdout(self, tmp_json_file, ti, capsys): + """The write_stdout (filebeat-style) branch uses the pinned template too.""" + self.elasticsearch_io.write_to_es = False + ti_context = SimpleNamespace( + log_id_template="{dag_id}-{task_id}-{logical_date}-{try_number}", + dag_run=SimpleNamespace( + logical_date=datetime(2016, 1, 1), + data_interval_start=None, + data_interval_end=None, + ), + ) + + self.elasticsearch_io.upload(tmp_json_file, ti, ti_context=ti_context) + + expected_log_id = f"{ti.dag_id}-{ti.task_id}-2016-01-01T00:00:00+00:00-{ti.try_number}" + log_entries = [json.loads(line) for line in capsys.readouterr().out.strip().splitlines()] + assert log_entries + assert all(entry["log_id"] == expected_log_id for entry in log_entries) + + def test_upload_falls_back_to_configured_template_without_ti_context(self, tmp_json_file, ti, capsys): + self.elasticsearch_io.write_to_es = False + + self.elasticsearch_io.upload(tmp_json_file, ti) + + expected_log_id = _render_log_id(self.elasticsearch_io.log_id_template, ti, ti.try_number) + log_entries = [json.loads(line) for line in capsys.readouterr().out.strip().splitlines()] + assert log_entries + assert all(entry["log_id"] == expected_log_id for entry in log_entries) + class TestFormatErrorDetail: def test_returns_none_for_empty(self): diff --git a/providers/opensearch/src/airflow/providers/opensearch/log/os_task_handler.py b/providers/opensearch/src/airflow/providers/opensearch/log/os_task_handler.py index 77ed6830e8b1e..21ed4d0e109f1 100644 --- a/providers/opensearch/src/airflow/providers/opensearch/log/os_task_handler.py +++ b/providers/opensearch/src/airflow/providers/opensearch/log/os_task_handler.py @@ -68,6 +68,8 @@ else: from airflow.utils import timezone # type: ignore[attr-defined,no-redef] +log = logging.getLogger(__name__) + USE_PER_RUN_LOG_ID = hasattr(DagRun, "get_log_template") LOG_LINE_DEFAULTS = {"exc_text": "", "stack_info": ""} TASK_LOG_FIELDS = ["timestamp", "event", "level", "chan", "logger", "error_detail", "message", "levelname"] @@ -222,18 +224,70 @@ def _strip_userinfo(url: str) -> str: return parsed._replace(netloc=netloc).geturl() +def _clean_date(value: datetime | None) -> str: + """ + Clean up a date value so that it is safe to query in opensearch by removing reserved characters. + + https://www.elastic.co/guide/en/elasticsearch/reference/current/query-dsl-query-string-query.html#_reserved_characters + """ + if value is None: + return "" + return value.strftime("%Y_%m_%dT%H_%M_%S_%f") + + def _render_log_id( - log_id_template: str, ti: TaskInstance | TaskInstanceKey | RuntimeTI, try_number: int + log_id_template: str, + ti: TaskInstance | TaskInstanceKey | RuntimeTI, + try_number: int, + *, + dag_run: Any = None, + json_format: bool = False, ) -> str: + # Templates pinned by older LogTemplate rows may use date placeholders (e.g. the pre-2.3 + # default `{dag_id}-{task_id}-{logical_date}-{try_number}`), and str.format raises KeyError + # on any missing name, so always supply the full set. + def format_date(value: datetime | None) -> str: + if json_format: + return _clean_date(value) + return value.isoformat() if value else "" + + logical_date = format_date(getattr(dag_run, "logical_date", None)) return log_id_template.format( dag_id=ti.dag_id, task_id=ti.task_id, run_id=getattr(ti, "run_id", ""), try_number=try_number, map_index=getattr(ti, "map_index", ""), + logical_date=logical_date, + execution_date=logical_date, + data_interval_start=format_date(getattr(dag_run, "data_interval_start", None)), + data_interval_end=format_date(getattr(dag_run, "data_interval_end", None)), ) +def _resolve_log_id_template( + ti: TaskInstance | TaskInstanceKey | RuntimeTI, default_template: str +) -> tuple[str, DagRun | None]: + """ + Fetch the Dag-run-pinned log id template and the Dag run from the metadata DB. + + Falls back to the configured template when the DB is not reachable (e.g. on workers, which + receive the pinned template through ``TIRunContext`` instead). + """ + try: + with create_session() as session: + dag_run = ti.get_dagrun(session=session) # type: ignore[union-attr] + if USE_PER_RUN_LOG_ID: + return dag_run.get_log_template(session=session).elasticsearch_id, dag_run + return default_template, dag_run + except Exception: + log.warning( + "Could not fetch the Dag-run-pinned log id template; falling back to the configured one", + exc_info=True, + ) + return default_template, None + + def _resolve_nested(hit: dict[Any, Any], parent_class=None) -> type[Hit]: """ Resolve nested hits from OpenSearch by iteratively navigating the `_nested` field. @@ -864,8 +918,13 @@ def __attrs_post_init__(self): self._doc_type_map: dict[Any, Any] = {} self._doc_type: list[Any] = [] - def upload(self, path: os.PathLike | str, ti: RuntimeTI | None = None) -> None: - """Emit structured task logs to stdout and/or write them directly to OpenSearch.""" + def upload(self, path: os.PathLike | str, ti: RuntimeTI | None = None, *, ti_context: Any = None) -> None: + """ + Emit structured task logs to stdout and/or write them directly to OpenSearch. + + :param ti_context: the ``TIRunContext`` the supervisor received at task start, carrying the + Dag-run-pinned log id template; when absent, fall back to the configured template. + """ if ti is None: return @@ -874,7 +933,14 @@ def upload(self, path: os.PathLike | str, ti: RuntimeTI | None = None) -> None: if not local_loc.is_file(): return - log_id = _render_log_id(self.log_id_template, ti, ti.try_number) # type: ignore[arg-type] + log_id_template = getattr(ti_context, "log_id_template", None) or self.log_id_template + log_id = _render_log_id( + log_id_template, + ti, + ti.try_number, + dag_run=getattr(ti_context, "dag_run", None), + json_format=self.json_format, + ) if self.write_stdout or self.write_to_opensearch: log_lines = self._parse_raw_log(local_loc.read_text(), log_id) @@ -932,7 +998,14 @@ def _write_to_opensearch(self, log_lines: list[dict[str, Any]]) -> bool: return False def read(self, _relative_path: str, ti: RuntimeTI) -> tuple[LogSourceInfo, LogMessages]: - log_id = _render_log_id(self.log_id_template, ti, ti.try_number) # type: ignore[arg-type] + log_id_template, dag_run = _resolve_log_id_template(ti, self.log_id_template) + log_id = _render_log_id( + log_id_template, + ti, + ti.try_number, + dag_run=dag_run, + json_format=self.json_format, + ) self.log.info("Reading log %s from Opensearch", log_id) response = self._os_read(log_id, 0, ti) if response is not None and response.hits: diff --git a/providers/opensearch/tests/unit/opensearch/log/test_os_task_handler.py b/providers/opensearch/tests/unit/opensearch/log/test_os_task_handler.py index f9d9fa8942603..645395e3c6a31 100644 --- a/providers/opensearch/tests/unit/opensearch/log/test_os_task_handler.py +++ b/providers/opensearch/tests/unit/opensearch/log/test_os_task_handler.py @@ -23,6 +23,7 @@ import re from io import StringIO from pathlib import Path +from types import SimpleNamespace from unittest.mock import Mock, patch import pendulum @@ -509,6 +510,14 @@ def test_render_log_id(self, ti): self.os_task_handler.json_format = True assert self.os_task_handler._render_log_id(ti, 1) == self.JSON_LOG_ID + @pytest.mark.db_test + def test_remote_log_io_read_uses_pinned_template(self, ti): + """The read path renders with the LogTemplate row pinned to the Dag run, not the conf value.""" + with patch.object(self.os_task_handler.io, "_os_read", return_value=None) as mock_os_read: + self.os_task_handler.io.read("", ti) + + mock_os_read.assert_called_once_with(self.LOG_ID, 0, ti) + def test_clean_date(self): clean_logical_date = OpensearchTaskHandler._clean_date(datetime(2016, 7, 8, 9, 10, 11, 12)) assert clean_logical_date == "2016_07_08T09_10_11_000012" @@ -760,6 +769,45 @@ def test_upload_returns_early_when_ti_is_none(self, tmp_path): log_file.write_text('{"message": "test"}\n') self.opensearch_io.upload(log_file, ti=None) + def test_upload_uses_pinned_log_id_template_from_ti_context(self, tmp_json_file, ti): + """Documents are stamped with the Dag-run-pinned template delivered via TIRunContext.""" + self.opensearch_io.write_stdout = False + ti_context = SimpleNamespace( + log_id_template="{dag_id}-{task_id}-{logical_date}-{try_number}", + dag_run=SimpleNamespace( + logical_date=datetime(2016, 1, 1), + data_interval_start=None, + data_interval_end=None, + ), + ) + + with patch.object(self.opensearch_io, "_write_to_opensearch", return_value=False) as mock_write: + self.opensearch_io.upload(tmp_json_file, ti, ti_context=ti_context) + + expected_log_id = f"{ti.dag_id}-{ti.task_id}-2016-01-01T00:00:00+00:00-{ti.try_number}" + log_lines = mock_write.call_args.args[0] + assert log_lines + assert all(line["log_id"] == expected_log_id for line in log_lines) + + def test_upload_stamps_pinned_log_id_on_stdout(self, tmp_json_file, ti, capsys): + """The write_stdout (filebeat-style) branch uses the pinned template too.""" + self.opensearch_io.write_to_opensearch = False + ti_context = SimpleNamespace( + log_id_template="{dag_id}-{task_id}-{logical_date}-{try_number}", + dag_run=SimpleNamespace( + logical_date=datetime(2016, 1, 1), + data_interval_start=None, + data_interval_end=None, + ), + ) + + self.opensearch_io.upload(tmp_json_file, ti, ti_context=ti_context) + + expected_log_id = f"{ti.dag_id}-{ti.task_id}-2016-01-01T00:00:00+00:00-{ti.try_number}" + log_entries = [json.loads(line) for line in capsys.readouterr().out.strip().splitlines()] + assert log_entries + assert all(entry["log_id"] == expected_log_id for entry in log_entries) + class TestFormatErrorDetail: def test_returns_none_for_empty(self): diff --git a/task-sdk/src/airflow/sdk/api/datamodels/_generated.py b/task-sdk/src/airflow/sdk/api/datamodels/_generated.py index 72aebeab3e49b..4a684aaa03261 100644 --- a/task-sdk/src/airflow/sdk/api/datamodels/_generated.py +++ b/task-sdk/src/airflow/sdk/api/datamodels/_generated.py @@ -27,7 +27,7 @@ from pydantic import AwareDatetime, BaseModel, ConfigDict, Field, JsonValue, RootModel -API_VERSION: Final[str] = "2026-06-30" +API_VERSION: Final[str] = "2026-09-30" class AssetAliasReferenceAssetEventDagRun(BaseModel): @@ -797,3 +797,4 @@ class TIRunContext(BaseModel): xcom_keys_to_clear: Annotated[list[str] | None, Field(title="Xcom Keys To Clear")] = None should_retry: Annotated[bool | None, Field(title="Should Retry")] = False start_date: Annotated[AwareDatetime | None, Field(title="Start Date")] = None + log_id_template: Annotated[str | None, Field(title="Log Id Template")] = None diff --git a/task-sdk/src/airflow/sdk/execution_time/schema/schema.json b/task-sdk/src/airflow/sdk/execution_time/schema/schema.json index 01e3c8970b5cb..4ad4d6b5313b4 100644 --- a/task-sdk/src/airflow/sdk/execution_time/schema/schema.json +++ b/task-sdk/src/airflow/sdk/execution_time/schema/schema.json @@ -1,6 +1,6 @@ { "$schema": "https://json-schema.org/draft/2020-12/schema", - "api_version": "2026-06-16", + "api_version": "2026-09-30", "description": "Apache Airflow SDK Supervisor Schema", "$defs": { "AssetAliasReferenceAssetEventDagRun": { @@ -4878,6 +4878,18 @@ ], "default": null, "title": "Start Date" + }, + "log_id_template": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Log Id Template" } }, "required": [ diff --git a/task-sdk/src/airflow/sdk/execution_time/schema/versions/__init__.py b/task-sdk/src/airflow/sdk/execution_time/schema/versions/__init__.py index 9491a8993fdc3..ad22cc071369f 100644 --- a/task-sdk/src/airflow/sdk/execution_time/schema/versions/__init__.py +++ b/task-sdk/src/airflow/sdk/execution_time/schema/versions/__init__.py @@ -37,8 +37,11 @@ def get_bundle() -> VersionBundle: """ from cadwyn import HeadVersion, Version, VersionBundle + from airflow.sdk.execution_time.schema.versions.v2026_09_30 import AddLogIdTemplateField + return VersionBundle( HeadVersion(), + Version("2026-09-30", AddLogIdTemplateField), Version("2026-06-16"), ) diff --git a/task-sdk/src/airflow/sdk/execution_time/schema/versions/v2026_09_30.py b/task-sdk/src/airflow/sdk/execution_time/schema/versions/v2026_09_30.py new file mode 100644 index 0000000000000..fdf074210c93a --- /dev/null +++ b/task-sdk/src/airflow/sdk/execution_time/schema/versions/v2026_09_30.py @@ -0,0 +1,30 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +from __future__ import annotations + +from cadwyn import VersionChange, schema + +from airflow.sdk.api.datamodels._generated import TIRunContext + + +class AddLogIdTemplateField(VersionChange): + """Add the Dag-run-pinned `log_id_template` field to TIRunContext (nested in StartupDetails).""" + + description = __doc__ + + instructions_to_migrate_to_previous_version = (schema(TIRunContext).field("log_id_template").didnt_exist,) diff --git a/task-sdk/src/airflow/sdk/execution_time/supervisor.py b/task-sdk/src/airflow/sdk/execution_time/supervisor.py index 447f7c7594de6..401078018c5e5 100644 --- a/task-sdk/src/airflow/sdk/execution_time/supervisor.py +++ b/task-sdk/src/airflow/sdk/execution_time/supervisor.py @@ -171,6 +171,7 @@ from typing_extensions import Self from airflow.executors.workloads import BundleInfo + from airflow.sdk.api.datamodels._generated import TIRunContext from airflow.sdk.bases.secrets_backend import BaseSecretsBackend from airflow.sdk.definitions.connection import Connection from airflow.sdk.types import RuntimeTaskInstanceProtocol as RuntimeTI @@ -1329,6 +1330,9 @@ class ActivitySubprocess(WatchedSubprocess): _should_retry: bool = attrs.field(default=False, init=False) """Whether the task should retry or not as decided by the API server.""" + _ti_context: TIRunContext | None = attrs.field(default=None, init=False) + """The run context received from the API server at task start; reused when uploading remote logs.""" + # After the failure of a heartbeat, we'll increment this counter. If it reaches `MAX_FAILED_HEARTBEATS`, we # will kill theprocess. This is to handle temporary network issues etc. ensuring that the process # does not hang around forever. @@ -1387,6 +1391,7 @@ def _on_child_started( # tell us "no, stop!" for any reason) ti_context = self.client.task_instances.start(ti.id, self.pid, datetime.now(tz=timezone.utc)) self._should_retry = ti_context.should_retry + self._ti_context = ti_context self._last_successful_heartbeat = time.monotonic() except Exception: # On any error kill that subprocess! @@ -1534,7 +1539,7 @@ def _upload_logs(self): try: with _remote_logging_conn(self.client): - upload_to_remote(self.process_log, self.ti) + upload_to_remote(self.process_log, self.ti, ti_context=self._ti_context) except Exception: self.process_log.exception("Failed to upload remote logs", ti_id=self.id, pid=self.pid) diff --git a/task-sdk/src/airflow/sdk/log.py b/task-sdk/src/airflow/sdk/log.py index 5e775cbcac4bd..b4e102b0baad9 100644 --- a/task-sdk/src/airflow/sdk/log.py +++ b/task-sdk/src/airflow/sdk/log.py @@ -17,6 +17,7 @@ # under the License. from __future__ import annotations +import inspect from contextlib import suppress from functools import cache from pathlib import Path @@ -33,6 +34,7 @@ from structlog.typing import EventDict, FilteringBoundLogger, Processor from airflow.sdk._shared.logging.remote import RemoteLogIO + from airflow.sdk.api.datamodels._generated import TIRunContext from airflow.sdk.types import Logger, RuntimeTaskInstanceProtocol as RuntimeTI @@ -224,7 +226,9 @@ def relative_path_from_logger(logger) -> Path | None: return Path(fname).relative_to(base_log_folder) -def upload_to_remote(logger: FilteringBoundLogger, ti: RuntimeTI | None = None): +def upload_to_remote( + logger: FilteringBoundLogger, ti: RuntimeTI | None = None, *, ti_context: TIRunContext | None = None +): raw_logger = getattr(logger, "_logger") handler = load_remote_log_handler() @@ -239,7 +243,12 @@ def upload_to_remote(logger: FilteringBoundLogger, ti: RuntimeTI | None = None): return log_relative_path = relative_path.as_posix() - handler.upload(log_relative_path, ti) + kwargs = {} + # Only handlers that opt in receive the run context (e.g. ES/OpenSearch use the Dag-run-pinned + # log id template from it); older or third-party RemoteLogIO implementations keep working. + if ti_context is not None and "ti_context" in inspect.signature(handler.upload).parameters: + kwargs["ti_context"] = ti_context + handler.upload(log_relative_path, ti, **kwargs) def mask_secret(secret: JsonValue, name: str | None = None) -> None: diff --git a/task-sdk/tests/task_sdk/execution_time/test_supervisor.py b/task-sdk/tests/task_sdk/execution_time/test_supervisor.py index 870b2e6deedca..b2d1a70940178 100644 --- a/task-sdk/tests/task_sdk/execution_time/test_supervisor.py +++ b/task-sdk/tests/task_sdk/execution_time/test_supervisor.py @@ -3603,7 +3603,7 @@ def handle_request(request: httpx.Request) -> httpx.Response: if remote_logging and expected_env: connection_available = {"available": False, "conn_uri": None} - def mock_upload_to_remote(process_log, ti): + def mock_upload_to_remote(process_log, ti, ti_context=None): connection_available["available"] = expected_env in os.environ connection_available["conn_uri"] = os.environ.get(expected_env) @@ -3652,6 +3652,51 @@ def test_log_upload_failures_are_non_fatal(mocker): ) +def test_on_child_started_retains_ti_context(mocker, make_ti_context): + ti_context = make_ti_context() + client = mocker.MagicMock() + client.task_instances.start.return_value = ti_context + proc = ActivitySubprocess( + process_log=mocker.MagicMock(), + id=TI_ID, + pid=12345, + stdin=mocker.MagicMock(), + client=client, + process=mocker.MagicMock(), + ) + mocker.patch.object(ActivitySubprocess, "send_msg") + + proc._on_child_started( + ti=TaskInstance(id=TI_ID, task_id="b", dag_id="c", run_id="d", try_number=1, dag_version_id=uuid7()), + dag_rel_path="test.py", + bundle_info=FAKE_BUNDLE, + sentry_integration="", + ) + + assert proc._ti_context is ti_context + + +def test_upload_logs_forwards_ti_context(mocker, make_ti_context): + """The supervisor passes the retained TIRunContext to the remote log uploader.""" + proc = ActivitySubprocess( + process_log=mocker.MagicMock(), + id=TI_ID, + pid=12345, + stdin=mocker.MagicMock(), + client=mocker.MagicMock(), + process=mocker.MagicMock(), + ) + proc.ti = mocker.MagicMock() + proc._ti_context = make_ti_context() + + mocker.patch("airflow.sdk.execution_time.supervisor._remote_logging_conn") + upload_to_remote = mocker.patch("airflow.sdk.log.upload_to_remote") + + proc._upload_logs() + + upload_to_remote.assert_called_once_with(proc.process_log, proc.ti, ti_context=proc._ti_context) + + def test_logs_uploaded_even_when_state_update_fails(mocker): """`wait()` must upload remote logs even if the final state update raises. diff --git a/task-sdk/tests/task_sdk/test_log.py b/task-sdk/tests/task_sdk/test_log.py new file mode 100644 index 0000000000000..a22644b90c6fe --- /dev/null +++ b/task-sdk/tests/task_sdk/test_log.py @@ -0,0 +1,74 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +from __future__ import annotations + +from pathlib import Path +from unittest import mock + +from airflow.sdk.log import upload_to_remote + + +class LegacyRemoteLogIO: + """A handler whose ``upload`` predates the ``ti_context`` keyword.""" + + def __init__(self): + self.calls: list[tuple] = [] + + def upload(self, path, ti): + self.calls.append((path, ti)) + + +class ContextAwareRemoteLogIO: + """A handler that opts in to receiving the run context.""" + + def __init__(self): + self.calls: list[tuple] = [] + + def upload(self, path, ti, *, ti_context=None): + self.calls.append((path, ti, ti_context)) + + +@mock.patch("airflow.sdk.log.relative_path_from_logger", return_value=Path("dag/task/1.log")) +@mock.patch("airflow.sdk.log.load_remote_log_handler") +class TestUploadToRemoteTIContext: + def test_legacy_handler_does_not_receive_ti_context(self, mock_load_handler, mock_relative_path): + handler = LegacyRemoteLogIO() + mock_load_handler.return_value = handler + ti, ti_context = mock.Mock(), mock.Mock() + + upload_to_remote(mock.MagicMock(), ti, ti_context=ti_context) + + assert handler.calls == [("dag/task/1.log", ti)] + + def test_opted_in_handler_receives_ti_context(self, mock_load_handler, mock_relative_path): + handler = ContextAwareRemoteLogIO() + mock_load_handler.return_value = handler + ti, ti_context = mock.Mock(), mock.Mock() + + upload_to_remote(mock.MagicMock(), ti, ti_context=ti_context) + + assert handler.calls == [("dag/task/1.log", ti, ti_context)] + + def test_no_ti_context_keeps_handler_default(self, mock_load_handler, mock_relative_path): + handler = ContextAwareRemoteLogIO() + mock_load_handler.return_value = handler + ti = mock.Mock() + + upload_to_remote(mock.MagicMock(), ti) + + assert handler.calls == [("dag/task/1.log", ti, None)] From cda35150cb9636ca45ed9f9feba1e1dc5ff3f835 Mon Sep 17 00:00:00 2001 From: LIU ZHE YOU Date: Sat, 11 Jul 2026 06:35:33 +0000 Subject: [PATCH 2/5] Avoid metadata DB access when resolving the log id template for worker-side TIs Worker-side RuntimeTIs have no get_dagrun and cannot reach the metadata DB; the pinned template reaches them through TIRunContext instead. Opening a session just to fail also trips the DB-access guard (AirflowInternalRuntimeError is a BaseException, so the fallback except clause cannot catch it) in non-DB test runs. --- .../src/airflow/providers/elasticsearch/log/es_task_handler.py | 3 +++ .../src/airflow/providers/opensearch/log/os_task_handler.py | 3 +++ 2 files changed, 6 insertions(+) diff --git a/providers/elasticsearch/src/airflow/providers/elasticsearch/log/es_task_handler.py b/providers/elasticsearch/src/airflow/providers/elasticsearch/log/es_task_handler.py index 8579acd8f9f9e..7ac37881b1107 100644 --- a/providers/elasticsearch/src/airflow/providers/elasticsearch/log/es_task_handler.py +++ b/providers/elasticsearch/src/airflow/providers/elasticsearch/log/es_task_handler.py @@ -234,6 +234,9 @@ def _resolve_log_id_template( Falls back to the configured template when the DB is not reachable (e.g. on workers, which receive the pinned template through ``TIRunContext`` instead). """ + if not hasattr(ti, "get_dagrun"): + # A worker-side RuntimeTI: don't even open a session, the metadata DB is out of reach. + return default_template, None try: with create_session() as session: dag_run = ti.get_dagrun(session=session) # type: ignore[union-attr] diff --git a/providers/opensearch/src/airflow/providers/opensearch/log/os_task_handler.py b/providers/opensearch/src/airflow/providers/opensearch/log/os_task_handler.py index 21ed4d0e109f1..31b8add6b583f 100644 --- a/providers/opensearch/src/airflow/providers/opensearch/log/os_task_handler.py +++ b/providers/opensearch/src/airflow/providers/opensearch/log/os_task_handler.py @@ -274,6 +274,9 @@ def _resolve_log_id_template( Falls back to the configured template when the DB is not reachable (e.g. on workers, which receive the pinned template through ``TIRunContext`` instead). """ + if not hasattr(ti, "get_dagrun"): + # A worker-side RuntimeTI: don't even open a session, the metadata DB is out of reach. + return default_template, None try: with create_session() as session: dag_run = ti.get_dagrun(session=session) # type: ignore[union-attr] From e9ac1addd770cec9a1cf2894c1803e0a985fc98e Mon Sep 17 00:00:00 2001 From: LIU ZHE YOU Date: Sat, 11 Jul 2026 06:35:49 +0000 Subject: [PATCH 3/5] Restore the parent package attribute after forced module re-imports in supervisor tests The remote-logging supervisor tests drop airflow.sdk.log (and the logging-config modules) from sys.modules to force a fresh import, but the re-import rebinds the attribute on the parent package and monkeypatch only restores the sys.modules entry. On Python < 3.12 mock.patch resolves dotted targets through getattr on the parent package, so later tests on the same worker patched the stale module object while the code under test imported the restored one, failing test_upload_logs_forwards_ti_context and TestUploadToRemoteTIContext::test_no_ti_context_keeps_handler_default in CI. --- .../execution_time/test_supervisor.py | 35 ++++++++++++++----- 1 file changed, 26 insertions(+), 9 deletions(-) diff --git a/task-sdk/tests/task_sdk/execution_time/test_supervisor.py b/task-sdk/tests/task_sdk/execution_time/test_supervisor.py index b2d1a70940178..8689622e8bad8 100644 --- a/task-sdk/tests/task_sdk/execution_time/test_supervisor.py +++ b/task-sdk/tests/task_sdk/execution_time/test_supervisor.py @@ -3540,6 +3540,23 @@ def noop_handler(request: httpx.Request) -> httpx.Response: assert called == 1 +def _forget_module(monkeypatch, name: str) -> None: + """ + Drop a module from ``sys.modules`` so the code under test re-imports it fresh. + + Besides the ``sys.modules`` entry, pin the attribute on the parent package so monkeypatch + restores it at teardown too: the re-import rebinds it to the fresh module object, and a + binding left stale would make later tests' ``mock.patch(".")`` patch the wrong + module object on Python < 3.12, where dotted patch targets resolve through getattr on the + parent package rather than through ``sys.modules``. + """ + parent_name, _, attr_name = name.rpartition(".") + parent = sys.modules.get(parent_name) + if parent is not None and hasattr(parent, attr_name): + monkeypatch.setattr(parent, attr_name, getattr(parent, attr_name)) + monkeypatch.delitem(sys.modules, name, raising=False) + + @pytest.mark.parametrize( ("remote_logging", "remote_conn", "expected_env"), ( @@ -3555,9 +3572,9 @@ def test_remote_logging_conn(remote_logging, remote_conn, expected_env, monkeypa pytest.importorskip("airflow.providers.amazon", reason="'amazon' provider not installed") # This test is a little bit overly specific to how the logging is currently configured :/ - monkeypatch.delitem(sys.modules, "airflow.logging_config") - monkeypatch.delitem(sys.modules, "airflow.config_templates.airflow_local_settings", raising=False) - monkeypatch.delitem(sys.modules, "airflow.sdk.log", raising=False) + _forget_module(monkeypatch, "airflow.logging_config") + _forget_module(monkeypatch, "airflow.config_templates.airflow_local_settings") + _forget_module(monkeypatch, "airflow.sdk.log") def handle_request(request: httpx.Request) -> httpx.Response: return httpx.Response( @@ -3736,9 +3753,9 @@ def test_remote_logging_conn_sets_process_context(monkeypatch, mocker): from airflow.models.connection import Connection as CoreConnection from airflow.sdk.definitions.connection import Connection as SDKConnection - monkeypatch.delitem(sys.modules, "airflow.logging_config") - monkeypatch.delitem(sys.modules, "airflow.config_templates.airflow_local_settings", raising=False) - monkeypatch.delitem(sys.modules, "airflow.sdk.log", raising=False) + _forget_module(monkeypatch, "airflow.logging_config") + _forget_module(monkeypatch, "airflow.config_templates.airflow_local_settings") + _forget_module(monkeypatch, "airflow.sdk.log") conn_id = "s3_conn_logs" conn_uri = "aws:///?region_name=us-east-1" @@ -3891,9 +3908,9 @@ def test_remote_logging_conn_caches_connection_not_client(monkeypatch): import gc import weakref - monkeypatch.delitem(sys.modules, "airflow.logging_config") - monkeypatch.delitem(sys.modules, "airflow.config_templates.airflow_local_settings", raising=False) - monkeypatch.delitem(sys.modules, "airflow.sdk.log", raising=False) + _forget_module(monkeypatch, "airflow.logging_config") + _forget_module(monkeypatch, "airflow.config_templates.airflow_local_settings") + _forget_module(monkeypatch, "airflow.sdk.log") from airflow.sdk.execution_time import supervisor From e43e13d43e3b895fac3d9b5d558c413f99c1bacd Mon Sep 17 00:00:00 2001 From: LIU ZHE YOU Date: Wed, 15 Jul 2026 02:52:43 +0000 Subject: [PATCH 4/5] Move Elasticsearch and OpenSearch handler changes to stacked PRs Keep this branch scoped to the core plumbing (Execution API and Task SDK) that delivers the per-run log ID template to workers. The Elasticsearch and OpenSearch handler changes re-land in per-provider PRs stacked on this branch, so each provider change can be reviewed and released on its own cadence. --- .../elasticsearch/log/es_task_handler.py | 88 ++----------------- .../elasticsearch/log/test_es_task_handler.py | 72 +-------------- .../opensearch/log/os_task_handler.py | 86 ++---------------- .../opensearch/log/test_os_task_handler.py | 48 ---------- 4 files changed, 14 insertions(+), 280 deletions(-) diff --git a/providers/elasticsearch/src/airflow/providers/elasticsearch/log/es_task_handler.py b/providers/elasticsearch/src/airflow/providers/elasticsearch/log/es_task_handler.py index 7ac37881b1107..fa05b2a0b34cf 100644 --- a/providers/elasticsearch/src/airflow/providers/elasticsearch/log/es_task_handler.py +++ b/providers/elasticsearch/src/airflow/providers/elasticsearch/log/es_task_handler.py @@ -51,7 +51,6 @@ from airflow.providers.elasticsearch.version_compat import AIRFLOW_V_3_0_PLUS, AIRFLOW_V_3_2_PLUS from airflow.utils.log.file_task_handler import FileTaskHandler from airflow.utils.log.logging_mixin import ExternalLoggingMixin, LoggingMixin -from airflow.utils.session import create_session from airflow.utils.state import TaskInstanceState if AIRFLOW_V_3_2_PLUS: @@ -77,8 +76,6 @@ EsLogMsgType = list[tuple[str, str]] # type: ignore[assignment,misc] -log = logging.getLogger(__name__) - LOG_LINE_DEFAULTS = {"exc_text": "", "stack_info": ""} # Elasticsearch hosted log type @@ -195,62 +192,16 @@ def _strip_userinfo(url: str) -> str: return parsed._replace(netloc=netloc).geturl() -def _render_log_id( - log_id_template: str, - ti: TaskInstance | TaskInstanceKey | RuntimeTI, - try_number: int, - *, - dag_run: Any = None, - json_format: bool = False, -) -> str: - # Templates pinned by older LogTemplate rows may use date placeholders (e.g. the pre-2.3 - # default `{dag_id}-{task_id}-{logical_date}-{try_number}`), and str.format raises KeyError - # on any missing name, so always supply the full set. - def format_date(value: datetime | None) -> str: - if json_format: - return _clean_date(value) - return value.isoformat() if value else "" - - logical_date = format_date(getattr(dag_run, "logical_date", None)) +def _render_log_id(log_id_template: str, ti: TaskInstance | TaskInstanceKey, try_number: int) -> str: return log_id_template.format( dag_id=ti.dag_id, task_id=ti.task_id, run_id=getattr(ti, "run_id", ""), try_number=try_number, map_index=getattr(ti, "map_index", ""), - logical_date=logical_date, - execution_date=logical_date, - data_interval_start=format_date(getattr(dag_run, "data_interval_start", None)), - data_interval_end=format_date(getattr(dag_run, "data_interval_end", None)), ) -def _resolve_log_id_template( - ti: TaskInstance | TaskInstanceKey | RuntimeTI, default_template: str -) -> tuple[str, DagRun | None]: - """ - Fetch the Dag-run-pinned log id template and the Dag run from the metadata DB. - - Falls back to the configured template when the DB is not reachable (e.g. on workers, which - receive the pinned template through ``TIRunContext`` instead). - """ - if not hasattr(ti, "get_dagrun"): - # A worker-side RuntimeTI: don't even open a session, the metadata DB is out of reach. - return default_template, None - try: - with create_session() as session: - dag_run = ti.get_dagrun(session=session) # type: ignore[union-attr] - if USE_PER_RUN_LOG_ID: - return dag_run.get_log_template(session=session).elasticsearch_id, dag_run - return default_template, dag_run - except Exception: - log.warning( - "Could not fetch the Dag-run-pinned log id template; falling back to the configured one", - exc_info=True, - ) - return default_template, None - - def _clean_date(value: datetime | None) -> str: """ Clean up a date value so that it is safe to query in elasticsearch by removing reserved characters. @@ -426,10 +377,6 @@ def _get_es_source_includes(self) -> list[str]: extra_fields = self.json_fields if self.json_format and not AIRFLOW_V_3_0_PLUS else [] return self.io._get_source_includes(extra_fields=extra_fields) - def _render_log_id(self, ti: TaskInstance | TaskInstanceKey, try_number: int) -> str: - log_id_template, dag_run = _resolve_log_id_template(ti, self.log_id_template) - return _render_log_id(log_id_template, ti, try_number, dag_run=dag_run, json_format=self.json_format) - def _read( self, ti: TaskInstance, try_number: int, metadata: LogMetadata | None = None ) -> tuple[EsLogMsgType, LogMetadata]: @@ -460,7 +407,7 @@ def _read( metadata["offset"] = 0 offset = metadata["offset"] - log_id = self._render_log_id(ti, try_number) + log_id = _render_log_id(self.log_id_template, ti, try_number) response = self.io._es_read(log_id, offset, ti, source_includes=self._get_es_source_includes()) # TODO: Can we skip group logs by host ? if response is not None and response.hits: @@ -582,7 +529,7 @@ def set_context(self, ti: TaskInstance, *, identifier: str | None = None) -> Non _clean_date(ti.logical_date) if AIRFLOW_V_3_0_PLUS else _clean_date(ti.execution_date) ), "try_number": str(ti.try_number), - "log_id": self._render_log_id(ti, ti.try_number), + "log_id": _render_log_id(self.log_id_template, ti, ti.try_number), }, ) @@ -648,7 +595,7 @@ 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 """ - log_id = self._render_log_id(task_instance, try_number) + log_id = _render_log_id(self.log_id_template, task_instance, try_number) scheme = "" if "://" in self.frontend else "https://" return scheme + self.frontend.format(log_id=quote(log_id)) @@ -755,13 +702,8 @@ def _get_source_includes(self, extra_fields: Iterable[str] = ()) -> list[str]: ["@timestamp", *TASK_LOG_FIELDS, self.host_field, self.offset_field, *extra_fields] ) - def upload(self, path: os.PathLike | str, ti: RuntimeTI | None = None, *, ti_context: Any = None) -> None: - """ - Write the log to ElasticSearch. - - :param ti_context: the ``TIRunContext`` the supervisor received at task start, carrying the - Dag-run-pinned log id template; when absent, fall back to the configured template. - """ + def upload(self, path: os.PathLike | str, ti: RuntimeTI | None = None) -> None: + """Write the log to ElasticSearch.""" if ti is None: return @@ -772,14 +714,7 @@ def upload(self, path: os.PathLike | str, ti: RuntimeTI | None = None, *, ti_con else: local_loc = self.base_log_folder.joinpath(path) - log_id_template = getattr(ti_context, "log_id_template", None) or self.log_id_template - log_id = _render_log_id( - log_id_template, - ti, - ti.try_number, - dag_run=getattr(ti_context, "dag_run", None), - json_format=self.json_format, - ) + log_id = _render_log_id(self.log_id_template, ti, ti.try_number) # type: ignore[arg-type] if local_loc.is_file() and self.write_stdout: # Intentionally construct the log_id and offset field @@ -849,14 +784,7 @@ def _write_to_es(self, log_lines: list[dict[str, Any]]) -> bool: return False def read(self, _relative_path: str, ti: RuntimeTI) -> tuple[LogSourceInfo, LogMessages]: - log_id_template, dag_run = _resolve_log_id_template(ti, self.log_id_template) - log_id = _render_log_id( - log_id_template, - ti, - ti.try_number, - dag_run=dag_run, - json_format=self.json_format, - ) + log_id = _render_log_id(self.log_id_template, ti, ti.try_number) # type: ignore[arg-type] self.log.info("Reading log %s from Elasticsearch", log_id) offset = 0 response = self._es_read(log_id, offset, ti, source_includes=self._get_source_includes()) diff --git a/providers/elasticsearch/tests/unit/elasticsearch/log/test_es_task_handler.py b/providers/elasticsearch/tests/unit/elasticsearch/log/test_es_task_handler.py index e1ea981fe0ea7..b923d8b7c6f00 100644 --- a/providers/elasticsearch/tests/unit/elasticsearch/log/test_es_task_handler.py +++ b/providers/elasticsearch/tests/unit/elasticsearch/log/test_es_task_handler.py @@ -23,7 +23,6 @@ import re from io import StringIO from pathlib import Path -from types import SimpleNamespace from unittest import mock from unittest.mock import Mock, patch from urllib.parse import quote @@ -491,22 +490,6 @@ def test_close_reopens_stream(self, ti, stream_state): def test_render_log_id(self, ti): assert _render_log_id(self.es_task_handler.log_id_template, ti, 1) == self.LOG_ID - @pytest.mark.db_test - def test_render_log_id_uses_pinned_template(self, ti): - """The handler renders with the LogTemplate row pinned to the Dag run, not the conf value.""" - logical_date = ti.get_dagrun().logical_date - expected = f"{self.DAG_ID}-{self.TASK_ID}-{logical_date.isoformat()}-{self.TRY_NUM}" - assert self.es_task_handler._render_log_id(ti, self.TRY_NUM) == expected - - @pytest.mark.db_test - def test_remote_log_io_read_uses_pinned_template(self, ti): - with mock.patch.object(self.es_task_handler.io, "_es_read", return_value=None) as mock_es_read: - self.es_task_handler.io.read("", ti) - - logical_date = ti.get_dagrun().logical_date - expected = f"{self.DAG_ID}-{self.TASK_ID}-{logical_date.isoformat()}-{self.TRY_NUM}" - mock_es_read.assert_called_once_with(expected, 0, ti, source_includes=mock.ANY) - def test_clean_date(self): clean_logical_date = _clean_date(datetime(2016, 7, 8, 9, 10, 11, 12)) assert clean_logical_date == "2016_07_08T09_10_11_000012" @@ -534,13 +517,8 @@ def test_get_external_log_url(self, ti, json_format, es_frontend, expected_url): offset_field=self.offset_field, frontend=es_frontend, ) - # The URL renders with the log id template pinned to the Dag run by the ``ti`` fixture - # (`{dag_id}-{task_id}-{logical_date}-{try_number}`), not the configured template. - logical_date = ti.get_dagrun().logical_date - expected_date = _clean_date(logical_date) if json_format else logical_date.isoformat() - expected_log_id = f"{self.DAG_ID}-{self.TASK_ID}-{expected_date}-{self.TRY_NUM}" assert es_task_handler.get_external_log_url(ti, ti.try_number) == expected_url.format( - log_id=quote(expected_log_id) + log_id=quote(self.LOG_ID) ) @pytest.mark.parametrize( @@ -878,54 +856,6 @@ def test_read_returns_missing_log_message_when_es_read_returns_none(self, ti): def test_upload_returns_early_when_ti_is_none(self, tmp_json_file): self.elasticsearch_io.upload(tmp_json_file, ti=None) - def test_upload_uses_pinned_log_id_template_from_ti_context(self, tmp_json_file, ti): - """Documents are stamped with the Dag-run-pinned template delivered via TIRunContext.""" - ti_context = SimpleNamespace( - log_id_template="{dag_id}-{task_id}-{logical_date}-{try_number}", - dag_run=SimpleNamespace( - logical_date=datetime(2016, 1, 1), - data_interval_start=None, - data_interval_end=None, - ), - ) - - with patch.object(self.elasticsearch_io, "_write_to_es", return_value=False) as mock_write: - self.elasticsearch_io.upload(tmp_json_file, ti, ti_context=ti_context) - - expected_log_id = f"{ti.dag_id}-{ti.task_id}-2016-01-01T00:00:00+00:00-{ti.try_number}" - log_lines = mock_write.call_args.args[0] - assert log_lines - assert all(line["log_id"] == expected_log_id for line in log_lines) - - def test_upload_stamps_pinned_log_id_on_stdout(self, tmp_json_file, ti, capsys): - """The write_stdout (filebeat-style) branch uses the pinned template too.""" - self.elasticsearch_io.write_to_es = False - ti_context = SimpleNamespace( - log_id_template="{dag_id}-{task_id}-{logical_date}-{try_number}", - dag_run=SimpleNamespace( - logical_date=datetime(2016, 1, 1), - data_interval_start=None, - data_interval_end=None, - ), - ) - - self.elasticsearch_io.upload(tmp_json_file, ti, ti_context=ti_context) - - expected_log_id = f"{ti.dag_id}-{ti.task_id}-2016-01-01T00:00:00+00:00-{ti.try_number}" - log_entries = [json.loads(line) for line in capsys.readouterr().out.strip().splitlines()] - assert log_entries - assert all(entry["log_id"] == expected_log_id for entry in log_entries) - - def test_upload_falls_back_to_configured_template_without_ti_context(self, tmp_json_file, ti, capsys): - self.elasticsearch_io.write_to_es = False - - self.elasticsearch_io.upload(tmp_json_file, ti) - - expected_log_id = _render_log_id(self.elasticsearch_io.log_id_template, ti, ti.try_number) - log_entries = [json.loads(line) for line in capsys.readouterr().out.strip().splitlines()] - assert log_entries - assert all(entry["log_id"] == expected_log_id for entry in log_entries) - class TestFormatErrorDetail: def test_returns_none_for_empty(self): diff --git a/providers/opensearch/src/airflow/providers/opensearch/log/os_task_handler.py b/providers/opensearch/src/airflow/providers/opensearch/log/os_task_handler.py index 31b8add6b583f..77ed6830e8b1e 100644 --- a/providers/opensearch/src/airflow/providers/opensearch/log/os_task_handler.py +++ b/providers/opensearch/src/airflow/providers/opensearch/log/os_task_handler.py @@ -68,8 +68,6 @@ else: from airflow.utils import timezone # type: ignore[attr-defined,no-redef] -log = logging.getLogger(__name__) - USE_PER_RUN_LOG_ID = hasattr(DagRun, "get_log_template") LOG_LINE_DEFAULTS = {"exc_text": "", "stack_info": ""} TASK_LOG_FIELDS = ["timestamp", "event", "level", "chan", "logger", "error_detail", "message", "levelname"] @@ -224,73 +222,18 @@ def _strip_userinfo(url: str) -> str: return parsed._replace(netloc=netloc).geturl() -def _clean_date(value: datetime | None) -> str: - """ - Clean up a date value so that it is safe to query in opensearch by removing reserved characters. - - https://www.elastic.co/guide/en/elasticsearch/reference/current/query-dsl-query-string-query.html#_reserved_characters - """ - if value is None: - return "" - return value.strftime("%Y_%m_%dT%H_%M_%S_%f") - - def _render_log_id( - log_id_template: str, - ti: TaskInstance | TaskInstanceKey | RuntimeTI, - try_number: int, - *, - dag_run: Any = None, - json_format: bool = False, + log_id_template: str, ti: TaskInstance | TaskInstanceKey | RuntimeTI, try_number: int ) -> str: - # Templates pinned by older LogTemplate rows may use date placeholders (e.g. the pre-2.3 - # default `{dag_id}-{task_id}-{logical_date}-{try_number}`), and str.format raises KeyError - # on any missing name, so always supply the full set. - def format_date(value: datetime | None) -> str: - if json_format: - return _clean_date(value) - return value.isoformat() if value else "" - - logical_date = format_date(getattr(dag_run, "logical_date", None)) return log_id_template.format( dag_id=ti.dag_id, task_id=ti.task_id, run_id=getattr(ti, "run_id", ""), try_number=try_number, map_index=getattr(ti, "map_index", ""), - logical_date=logical_date, - execution_date=logical_date, - data_interval_start=format_date(getattr(dag_run, "data_interval_start", None)), - data_interval_end=format_date(getattr(dag_run, "data_interval_end", None)), ) -def _resolve_log_id_template( - ti: TaskInstance | TaskInstanceKey | RuntimeTI, default_template: str -) -> tuple[str, DagRun | None]: - """ - Fetch the Dag-run-pinned log id template and the Dag run from the metadata DB. - - Falls back to the configured template when the DB is not reachable (e.g. on workers, which - receive the pinned template through ``TIRunContext`` instead). - """ - if not hasattr(ti, "get_dagrun"): - # A worker-side RuntimeTI: don't even open a session, the metadata DB is out of reach. - return default_template, None - try: - with create_session() as session: - dag_run = ti.get_dagrun(session=session) # type: ignore[union-attr] - if USE_PER_RUN_LOG_ID: - return dag_run.get_log_template(session=session).elasticsearch_id, dag_run - return default_template, dag_run - except Exception: - log.warning( - "Could not fetch the Dag-run-pinned log id template; falling back to the configured one", - exc_info=True, - ) - return default_template, None - - def _resolve_nested(hit: dict[Any, Any], parent_class=None) -> type[Hit]: """ Resolve nested hits from OpenSearch by iteratively navigating the `_nested` field. @@ -921,13 +864,8 @@ def __attrs_post_init__(self): self._doc_type_map: dict[Any, Any] = {} self._doc_type: list[Any] = [] - def upload(self, path: os.PathLike | str, ti: RuntimeTI | None = None, *, ti_context: Any = None) -> None: - """ - Emit structured task logs to stdout and/or write them directly to OpenSearch. - - :param ti_context: the ``TIRunContext`` the supervisor received at task start, carrying the - Dag-run-pinned log id template; when absent, fall back to the configured template. - """ + def upload(self, path: os.PathLike | str, ti: RuntimeTI | None = None) -> None: + """Emit structured task logs to stdout and/or write them directly to OpenSearch.""" if ti is None: return @@ -936,14 +874,7 @@ def upload(self, path: os.PathLike | str, ti: RuntimeTI | None = None, *, ti_con if not local_loc.is_file(): return - log_id_template = getattr(ti_context, "log_id_template", None) or self.log_id_template - log_id = _render_log_id( - log_id_template, - ti, - ti.try_number, - dag_run=getattr(ti_context, "dag_run", None), - json_format=self.json_format, - ) + log_id = _render_log_id(self.log_id_template, ti, ti.try_number) # type: ignore[arg-type] if self.write_stdout or self.write_to_opensearch: log_lines = self._parse_raw_log(local_loc.read_text(), log_id) @@ -1001,14 +932,7 @@ def _write_to_opensearch(self, log_lines: list[dict[str, Any]]) -> bool: return False def read(self, _relative_path: str, ti: RuntimeTI) -> tuple[LogSourceInfo, LogMessages]: - log_id_template, dag_run = _resolve_log_id_template(ti, self.log_id_template) - log_id = _render_log_id( - log_id_template, - ti, - ti.try_number, - dag_run=dag_run, - json_format=self.json_format, - ) + log_id = _render_log_id(self.log_id_template, ti, ti.try_number) # type: ignore[arg-type] self.log.info("Reading log %s from Opensearch", log_id) response = self._os_read(log_id, 0, ti) if response is not None and response.hits: diff --git a/providers/opensearch/tests/unit/opensearch/log/test_os_task_handler.py b/providers/opensearch/tests/unit/opensearch/log/test_os_task_handler.py index 645395e3c6a31..f9d9fa8942603 100644 --- a/providers/opensearch/tests/unit/opensearch/log/test_os_task_handler.py +++ b/providers/opensearch/tests/unit/opensearch/log/test_os_task_handler.py @@ -23,7 +23,6 @@ import re from io import StringIO from pathlib import Path -from types import SimpleNamespace from unittest.mock import Mock, patch import pendulum @@ -510,14 +509,6 @@ def test_render_log_id(self, ti): self.os_task_handler.json_format = True assert self.os_task_handler._render_log_id(ti, 1) == self.JSON_LOG_ID - @pytest.mark.db_test - def test_remote_log_io_read_uses_pinned_template(self, ti): - """The read path renders with the LogTemplate row pinned to the Dag run, not the conf value.""" - with patch.object(self.os_task_handler.io, "_os_read", return_value=None) as mock_os_read: - self.os_task_handler.io.read("", ti) - - mock_os_read.assert_called_once_with(self.LOG_ID, 0, ti) - def test_clean_date(self): clean_logical_date = OpensearchTaskHandler._clean_date(datetime(2016, 7, 8, 9, 10, 11, 12)) assert clean_logical_date == "2016_07_08T09_10_11_000012" @@ -769,45 +760,6 @@ def test_upload_returns_early_when_ti_is_none(self, tmp_path): log_file.write_text('{"message": "test"}\n') self.opensearch_io.upload(log_file, ti=None) - def test_upload_uses_pinned_log_id_template_from_ti_context(self, tmp_json_file, ti): - """Documents are stamped with the Dag-run-pinned template delivered via TIRunContext.""" - self.opensearch_io.write_stdout = False - ti_context = SimpleNamespace( - log_id_template="{dag_id}-{task_id}-{logical_date}-{try_number}", - dag_run=SimpleNamespace( - logical_date=datetime(2016, 1, 1), - data_interval_start=None, - data_interval_end=None, - ), - ) - - with patch.object(self.opensearch_io, "_write_to_opensearch", return_value=False) as mock_write: - self.opensearch_io.upload(tmp_json_file, ti, ti_context=ti_context) - - expected_log_id = f"{ti.dag_id}-{ti.task_id}-2016-01-01T00:00:00+00:00-{ti.try_number}" - log_lines = mock_write.call_args.args[0] - assert log_lines - assert all(line["log_id"] == expected_log_id for line in log_lines) - - def test_upload_stamps_pinned_log_id_on_stdout(self, tmp_json_file, ti, capsys): - """The write_stdout (filebeat-style) branch uses the pinned template too.""" - self.opensearch_io.write_to_opensearch = False - ti_context = SimpleNamespace( - log_id_template="{dag_id}-{task_id}-{logical_date}-{try_number}", - dag_run=SimpleNamespace( - logical_date=datetime(2016, 1, 1), - data_interval_start=None, - data_interval_end=None, - ), - ) - - self.opensearch_io.upload(tmp_json_file, ti, ti_context=ti_context) - - expected_log_id = f"{ti.dag_id}-{ti.task_id}-2016-01-01T00:00:00+00:00-{ti.try_number}" - log_entries = [json.loads(line) for line in capsys.readouterr().out.strip().splitlines()] - assert log_entries - assert all(entry["log_id"] == expected_log_id for entry in log_entries) - class TestFormatErrorDetail: def test_returns_none_for_empty(self): From a6333db6f5b101ba1ebb751418724a483facf65a Mon Sep 17 00:00:00 2001 From: LIU ZHE YOU Date: Wed, 15 Jul 2026 02:54:06 +0000 Subject: [PATCH 5/5] Make OpenSearch log handler respect the per-run log ID template The LogTemplate model pins the log_id template to each Dag run so task logs written before a config change stay readable, but the Airflow 3 paths of the OpenSearch handler rendered log_id from the current conf value on both the write (worker) and read (API server) sides, so changing [opensearch] log_id_template orphaned all previously written logs. The write side now uses the pinned template that the supervisor receives through TIRunContext (added by the core change this branch is stacked on), and the API-server read paths restore the per-run DB lookup with a conf fallback. --- .../opensearch/log/os_task_handler.py | 86 +++++++++++++++++-- .../opensearch/log/test_os_task_handler.py | 48 +++++++++++ 2 files changed, 129 insertions(+), 5 deletions(-) diff --git a/providers/opensearch/src/airflow/providers/opensearch/log/os_task_handler.py b/providers/opensearch/src/airflow/providers/opensearch/log/os_task_handler.py index 77ed6830e8b1e..31b8add6b583f 100644 --- a/providers/opensearch/src/airflow/providers/opensearch/log/os_task_handler.py +++ b/providers/opensearch/src/airflow/providers/opensearch/log/os_task_handler.py @@ -68,6 +68,8 @@ else: from airflow.utils import timezone # type: ignore[attr-defined,no-redef] +log = logging.getLogger(__name__) + USE_PER_RUN_LOG_ID = hasattr(DagRun, "get_log_template") LOG_LINE_DEFAULTS = {"exc_text": "", "stack_info": ""} TASK_LOG_FIELDS = ["timestamp", "event", "level", "chan", "logger", "error_detail", "message", "levelname"] @@ -222,18 +224,73 @@ def _strip_userinfo(url: str) -> str: return parsed._replace(netloc=netloc).geturl() +def _clean_date(value: datetime | None) -> str: + """ + Clean up a date value so that it is safe to query in opensearch by removing reserved characters. + + https://www.elastic.co/guide/en/elasticsearch/reference/current/query-dsl-query-string-query.html#_reserved_characters + """ + if value is None: + return "" + return value.strftime("%Y_%m_%dT%H_%M_%S_%f") + + def _render_log_id( - log_id_template: str, ti: TaskInstance | TaskInstanceKey | RuntimeTI, try_number: int + log_id_template: str, + ti: TaskInstance | TaskInstanceKey | RuntimeTI, + try_number: int, + *, + dag_run: Any = None, + json_format: bool = False, ) -> str: + # Templates pinned by older LogTemplate rows may use date placeholders (e.g. the pre-2.3 + # default `{dag_id}-{task_id}-{logical_date}-{try_number}`), and str.format raises KeyError + # on any missing name, so always supply the full set. + def format_date(value: datetime | None) -> str: + if json_format: + return _clean_date(value) + return value.isoformat() if value else "" + + logical_date = format_date(getattr(dag_run, "logical_date", None)) return log_id_template.format( dag_id=ti.dag_id, task_id=ti.task_id, run_id=getattr(ti, "run_id", ""), try_number=try_number, map_index=getattr(ti, "map_index", ""), + logical_date=logical_date, + execution_date=logical_date, + data_interval_start=format_date(getattr(dag_run, "data_interval_start", None)), + data_interval_end=format_date(getattr(dag_run, "data_interval_end", None)), ) +def _resolve_log_id_template( + ti: TaskInstance | TaskInstanceKey | RuntimeTI, default_template: str +) -> tuple[str, DagRun | None]: + """ + Fetch the Dag-run-pinned log id template and the Dag run from the metadata DB. + + Falls back to the configured template when the DB is not reachable (e.g. on workers, which + receive the pinned template through ``TIRunContext`` instead). + """ + if not hasattr(ti, "get_dagrun"): + # A worker-side RuntimeTI: don't even open a session, the metadata DB is out of reach. + return default_template, None + try: + with create_session() as session: + dag_run = ti.get_dagrun(session=session) # type: ignore[union-attr] + if USE_PER_RUN_LOG_ID: + return dag_run.get_log_template(session=session).elasticsearch_id, dag_run + return default_template, dag_run + except Exception: + log.warning( + "Could not fetch the Dag-run-pinned log id template; falling back to the configured one", + exc_info=True, + ) + return default_template, None + + def _resolve_nested(hit: dict[Any, Any], parent_class=None) -> type[Hit]: """ Resolve nested hits from OpenSearch by iteratively navigating the `_nested` field. @@ -864,8 +921,13 @@ def __attrs_post_init__(self): self._doc_type_map: dict[Any, Any] = {} self._doc_type: list[Any] = [] - def upload(self, path: os.PathLike | str, ti: RuntimeTI | None = None) -> None: - """Emit structured task logs to stdout and/or write them directly to OpenSearch.""" + def upload(self, path: os.PathLike | str, ti: RuntimeTI | None = None, *, ti_context: Any = None) -> None: + """ + Emit structured task logs to stdout and/or write them directly to OpenSearch. + + :param ti_context: the ``TIRunContext`` the supervisor received at task start, carrying the + Dag-run-pinned log id template; when absent, fall back to the configured template. + """ if ti is None: return @@ -874,7 +936,14 @@ def upload(self, path: os.PathLike | str, ti: RuntimeTI | None = None) -> None: if not local_loc.is_file(): return - log_id = _render_log_id(self.log_id_template, ti, ti.try_number) # type: ignore[arg-type] + log_id_template = getattr(ti_context, "log_id_template", None) or self.log_id_template + log_id = _render_log_id( + log_id_template, + ti, + ti.try_number, + dag_run=getattr(ti_context, "dag_run", None), + json_format=self.json_format, + ) if self.write_stdout or self.write_to_opensearch: log_lines = self._parse_raw_log(local_loc.read_text(), log_id) @@ -932,7 +1001,14 @@ def _write_to_opensearch(self, log_lines: list[dict[str, Any]]) -> bool: return False def read(self, _relative_path: str, ti: RuntimeTI) -> tuple[LogSourceInfo, LogMessages]: - log_id = _render_log_id(self.log_id_template, ti, ti.try_number) # type: ignore[arg-type] + log_id_template, dag_run = _resolve_log_id_template(ti, self.log_id_template) + log_id = _render_log_id( + log_id_template, + ti, + ti.try_number, + dag_run=dag_run, + json_format=self.json_format, + ) self.log.info("Reading log %s from Opensearch", log_id) response = self._os_read(log_id, 0, ti) if response is not None and response.hits: diff --git a/providers/opensearch/tests/unit/opensearch/log/test_os_task_handler.py b/providers/opensearch/tests/unit/opensearch/log/test_os_task_handler.py index f9d9fa8942603..645395e3c6a31 100644 --- a/providers/opensearch/tests/unit/opensearch/log/test_os_task_handler.py +++ b/providers/opensearch/tests/unit/opensearch/log/test_os_task_handler.py @@ -23,6 +23,7 @@ import re from io import StringIO from pathlib import Path +from types import SimpleNamespace from unittest.mock import Mock, patch import pendulum @@ -509,6 +510,14 @@ def test_render_log_id(self, ti): self.os_task_handler.json_format = True assert self.os_task_handler._render_log_id(ti, 1) == self.JSON_LOG_ID + @pytest.mark.db_test + def test_remote_log_io_read_uses_pinned_template(self, ti): + """The read path renders with the LogTemplate row pinned to the Dag run, not the conf value.""" + with patch.object(self.os_task_handler.io, "_os_read", return_value=None) as mock_os_read: + self.os_task_handler.io.read("", ti) + + mock_os_read.assert_called_once_with(self.LOG_ID, 0, ti) + def test_clean_date(self): clean_logical_date = OpensearchTaskHandler._clean_date(datetime(2016, 7, 8, 9, 10, 11, 12)) assert clean_logical_date == "2016_07_08T09_10_11_000012" @@ -760,6 +769,45 @@ def test_upload_returns_early_when_ti_is_none(self, tmp_path): log_file.write_text('{"message": "test"}\n') self.opensearch_io.upload(log_file, ti=None) + def test_upload_uses_pinned_log_id_template_from_ti_context(self, tmp_json_file, ti): + """Documents are stamped with the Dag-run-pinned template delivered via TIRunContext.""" + self.opensearch_io.write_stdout = False + ti_context = SimpleNamespace( + log_id_template="{dag_id}-{task_id}-{logical_date}-{try_number}", + dag_run=SimpleNamespace( + logical_date=datetime(2016, 1, 1), + data_interval_start=None, + data_interval_end=None, + ), + ) + + with patch.object(self.opensearch_io, "_write_to_opensearch", return_value=False) as mock_write: + self.opensearch_io.upload(tmp_json_file, ti, ti_context=ti_context) + + expected_log_id = f"{ti.dag_id}-{ti.task_id}-2016-01-01T00:00:00+00:00-{ti.try_number}" + log_lines = mock_write.call_args.args[0] + assert log_lines + assert all(line["log_id"] == expected_log_id for line in log_lines) + + def test_upload_stamps_pinned_log_id_on_stdout(self, tmp_json_file, ti, capsys): + """The write_stdout (filebeat-style) branch uses the pinned template too.""" + self.opensearch_io.write_to_opensearch = False + ti_context = SimpleNamespace( + log_id_template="{dag_id}-{task_id}-{logical_date}-{try_number}", + dag_run=SimpleNamespace( + logical_date=datetime(2016, 1, 1), + data_interval_start=None, + data_interval_end=None, + ), + ) + + self.opensearch_io.upload(tmp_json_file, ti, ti_context=ti_context) + + expected_log_id = f"{ti.dag_id}-{ti.task_id}-2016-01-01T00:00:00+00:00-{ti.try_number}" + log_entries = [json.loads(line) for line in capsys.readouterr().out.strip().splitlines()] + assert log_entries + assert all(entry["log_id"] == expected_log_id for entry in log_entries) + class TestFormatErrorDetail: def test_returns_none_for_empty(self):