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/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): 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..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( @@ -3603,7 +3620,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 +3669,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. @@ -3691,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" @@ -3846,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 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)]