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..f481edfff16f9 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,76 @@ 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_populates_latest_log_id_template( + self, client, session, create_task_instance, create_log_template, time_machine + ): + """Each Dag run pins whichever LogTemplate row is newest when it is created.""" + instant = timezone.parse("2024-09-30T12:00:00Z") + time_machine.move_to(instant, tick=False) + + for i in range(3): + elasticsearch_id = f"{{dag_id}}-{{task_id}}-{{logical_date}}-{{try_number}}-v{i}" + create_log_template(f"{{{{ ti.dag_id }}}}-v{i}.log", elasticsearch_id) + + ti = create_task_instance( + task_id="test_ti_run_populates_latest_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"] == 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/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)]