From 36196ee7913367265da82d8222052eec4e50754e Mon Sep 17 00:00:00 2001 From: ZHE YOU LIU Date: Tue, 27 Oct 2026 14:59:55 +0000 Subject: [PATCH 1/5] Load the serialized Dag in ti_run only for stub tasks Starting any task deserialized its whole Dag so the route could find out whether the task was a stub and read its arg_bindings. Denormalize the serialized is_stub flag onto task_instance, which ti_run already fetches and locks, and load the Dag only when it is set. execution_timeout is resolved from the same definition, so it is now sent for a lang-SDK task only; #75387 has a Python task send its own to the supervisor. Co-Authored-By: Claude --- airflow-core/docs/migrations-ref.rst | 4 +- .../execution_api/datamodels/taskinstance.py | 5 +- .../execution_api/routes/task_instances.py | 8 ++- ...0149_3_4_0_add_is_stub_to_task_instance.py | 52 +++++++++++++++++++ ...py => 0150_3_5_0_add_session_id_to_job.py} | 4 +- ...5_0_add_dag_run_dag_id_run_after_index.py} | 0 .../src/airflow/models/taskinstance.py | 5 ++ airflow-core/src/airflow/utils/db.py | 2 +- .../versions/head/test_task_instances.py | 49 ++++++++++++++++- .../lang_sdk_fixtures/java_native.json | 3 ++ .../tests/unit/dag_processing/test_dagbag.py | 11 ++++ .../tests/unit/models/test_taskinstance.py | 31 +++++++++++ 12 files changed, 166 insertions(+), 8 deletions(-) create mode 100644 airflow-core/src/airflow/migrations/versions/0149_3_4_0_add_is_stub_to_task_instance.py rename airflow-core/src/airflow/migrations/versions/{0149_3_5_0_add_session_id_to_job.py => 0150_3_5_0_add_session_id_to_job.py} (97%) rename airflow-core/src/airflow/migrations/versions/{0150_3_5_0_add_dag_run_dag_id_run_after_index.py => 0151_3_5_0_add_dag_run_dag_id_run_after_index.py} (100%) diff --git a/airflow-core/docs/migrations-ref.rst b/airflow-core/docs/migrations-ref.rst index af1282ee4707f..608170aa03954 100644 --- a/airflow-core/docs/migrations-ref.rst +++ b/airflow-core/docs/migrations-ref.rst @@ -41,7 +41,9 @@ Here's the list of all the Database Migrations that are executed via when you ru +=========================+==================+===================+==============================================================+ | ``38c8df3fcb9c`` (head) | ``c71065a1b14f`` | ``3.5.0`` | Add index on dag_run (dag_id, run_after). | +-------------------------+------------------+-------------------+--------------------------------------------------------------+ -| ``c71065a1b14f`` | ``ee86eed19e24`` | ``3.5.0`` | Add session_id and registration_id to Job. | +| ``c71065a1b14f`` | ``b4f1a7d2c903`` | ``3.5.0`` | Add session_id and registration_id to Job. | ++-------------------------+------------------+-------------------+--------------------------------------------------------------+ +| ``b4f1a7d2c903`` | ``ee86eed19e24`` | ``3.4.0`` | Add is_stub to task_instance. | +-------------------------+------------------+-------------------+--------------------------------------------------------------+ | ``ee86eed19e24`` | ``9cffe90fa0c0`` | ``3.4.0`` | Add pending_partition_key to asset_partition_dag_run and | | | | | enforce single pending row per key. | 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 f911b4ffab971..e6d771f0f82b0 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 @@ -558,9 +558,10 @@ class TIRunContext(BaseModel): execution_timeout: float | None = None """ - The task's ``execution_timeout`` in seconds, taken from its serialized Dag. + The ``execution_timeout`` in seconds of a task a language SDK runs, from its serialized Dag. - The supervisor enforces the timeout with this value. ``None`` when no timeout is set. + The supervisor enforces the timeout with this value. A Python task keeps its own timer in the + task runner, so the field is unset for one. ``None`` when no timeout is set. """ multi_team: bool = False 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 daa65d41c8706..c891f43ca6ac2 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 @@ -197,6 +197,7 @@ def ti_run( TI.unixname, TI.pid, TI.dag_version_id, + TI.is_stub, # This selects the raw JSON value, bypassing the deserialization -- we want that to happen on the # client column("next_kwargs", JSON), @@ -382,7 +383,9 @@ def ti_run( if first_task_reschedule_start_date is not None: context.first_task_reschedule_start_date = first_task_reschedule_start_date - serialized_task = get_task(dag_bag, ti, session=session) + # Only a lang-SDK (foreign-runtime) task needs its definition here. ti.is_stub is the row's + # copy of the serialized flag, so a regular task starts without loading the serialized Dag. + serialized_task = get_task(dag_bag, ti, session=session) if ti.is_stub else None # Only set for lang-SDK (foreign-runtime) tasks with a captured TaskFlow arg # spec; the route excludes unset fields, keeping regular responses lean. @@ -404,6 +407,9 @@ def ti_run( }, ) + # Resolved here for a lang-SDK task only, since a regular task no longer loads its + # definition. #75387 has a Python task send its own execution_timeout to the supervisor, + # which is where the value for that kind of task belongs. if (execution_timeout := get_execution_timeout(serialized_task)) is not None: context.execution_timeout = execution_timeout diff --git a/airflow-core/src/airflow/migrations/versions/0149_3_4_0_add_is_stub_to_task_instance.py b/airflow-core/src/airflow/migrations/versions/0149_3_4_0_add_is_stub_to_task_instance.py new file mode 100644 index 0000000000000..7f48942112a60 --- /dev/null +++ b/airflow-core/src/airflow/migrations/versions/0149_3_4_0_add_is_stub_to_task_instance.py @@ -0,0 +1,52 @@ +# +# 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. + +""" +Add is_stub to task_instance. + +Revision ID: b4f1a7d2c903 +Revises: ee86eed19e24 +Create Date: 2026-10-27 09:12:44.317602 + +""" + +from __future__ import annotations + +import sqlalchemy as sa +from alembic import op + +# revision identifiers, used by Alembic. +revision = "b4f1a7d2c903" +down_revision = "ee86eed19e24" +branch_labels = None +depends_on = None +airflow_version = "3.4.0" + +TABLE_NAME = "task_instance" + + +def upgrade(): + """Add is_stub to task_instance.""" + with op.batch_alter_table(TABLE_NAME) as batch_op: + batch_op.add_column(sa.Column("is_stub", sa.Boolean(), nullable=True)) + + +def downgrade(): + """Remove is_stub from task_instance.""" + with op.batch_alter_table(TABLE_NAME) as batch_op: + batch_op.drop_column("is_stub") diff --git a/airflow-core/src/airflow/migrations/versions/0149_3_5_0_add_session_id_to_job.py b/airflow-core/src/airflow/migrations/versions/0150_3_5_0_add_session_id_to_job.py similarity index 97% rename from airflow-core/src/airflow/migrations/versions/0149_3_5_0_add_session_id_to_job.py rename to airflow-core/src/airflow/migrations/versions/0150_3_5_0_add_session_id_to_job.py index e595a52b12a43..406271582baf6 100644 --- a/airflow-core/src/airflow/migrations/versions/0149_3_5_0_add_session_id_to_job.py +++ b/airflow-core/src/airflow/migrations/versions/0150_3_5_0_add_session_id_to_job.py @@ -20,7 +20,7 @@ Add session_id and registration_id to Job. Revision ID: c71065a1b14f -Revises: ee86eed19e24 +Revises: b4f1a7d2c903 Create Date: 2026-10-05 08:40:20.000000 """ @@ -33,7 +33,7 @@ from airflow.migrations.utils import disable_sqlite_fkeys revision = "c71065a1b14f" -down_revision = "ee86eed19e24" +down_revision = "b4f1a7d2c903" branch_labels = None depends_on = None airflow_version = "3.5.0" diff --git a/airflow-core/src/airflow/migrations/versions/0150_3_5_0_add_dag_run_dag_id_run_after_index.py b/airflow-core/src/airflow/migrations/versions/0151_3_5_0_add_dag_run_dag_id_run_after_index.py similarity index 100% rename from airflow-core/src/airflow/migrations/versions/0150_3_5_0_add_dag_run_dag_id_run_after_index.py rename to airflow-core/src/airflow/migrations/versions/0151_3_5_0_add_dag_run_dag_id_run_after_index.py diff --git a/airflow-core/src/airflow/models/taskinstance.py b/airflow-core/src/airflow/models/taskinstance.py index 3def6f39fdc6d..a1b48e8bf7a49 100644 --- a/airflow-core/src/airflow/models/taskinstance.py +++ b/airflow-core/src/airflow/models/taskinstance.py @@ -1141,6 +1141,9 @@ class TaskInstance(Base, LoggingMixin, BaseWorkload): priority_weight: Mapped[int | None] = mapped_column(Integer, nullable=True) operator: Mapped[str | None] = mapped_column(String(1000), nullable=True) custom_operator_name: Mapped[str | None] = mapped_column(String(1000), nullable=True) + # Denormalized from the task so ti_run can tell a stub (foreign-runtime) task from a regular one + # without loading the serialized Dag. Null on rows written before the column existed. + is_stub: Mapped[bool | None] = mapped_column(Boolean, nullable=True) queued_dttm: Mapped[datetime | None] = mapped_column(UtcDateTime, nullable=True) scheduled_dttm: Mapped[datetime | None] = mapped_column(UtcDateTime, nullable=True) queued_by_job_id: Mapped[int | None] = mapped_column(Integer, nullable=True) @@ -1385,6 +1388,7 @@ def insert_mapping( "executor_config": task.executor_config, "operator": task.task_type, "custom_operator_name": getattr(task, "operator_name", None), + "is_stub": task.is_stub, "region_index": region_index, "region_id": region_id, "_task_display_property_value": task.task_display_name, @@ -1581,6 +1585,7 @@ def refresh_from_task( self.operator = task.task_type op_name = getattr(task, "operator_name", None) self.custom_operator_name = op_name if isinstance(op_name, str) else "" + self.is_stub = task.is_stub # Re-apply cluster policy here so that task default do not overload previous data task_instance_mutation_hook(self, dag_run=dag_run) diff --git a/airflow-core/src/airflow/utils/db.py b/airflow-core/src/airflow/utils/db.py index 5fb7ac5dfffb1..85af474f8abfe 100644 --- a/airflow-core/src/airflow/utils/db.py +++ b/airflow-core/src/airflow/utils/db.py @@ -117,7 +117,7 @@ class MappedClassProtocol(Protocol): "3.1.8": "509b94a1042d", "3.2.0": "1d6611b6ab7c", "3.3.0": "d2f4e1b3c5a7", - "3.4.0": "ee86eed19e24", + "3.4.0": "b4f1a7d2c903", "3.5.0": "38c8df3fcb9c", } 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 aef15aade0b3f..1417b15f88f03 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 @@ -1109,6 +1109,49 @@ def transform(country: str, extracted: dict, limit: int = 10): ... assert response.status_code == 200 assert "arg_bindings" not in response.json() + def test_ti_run_does_not_load_the_dag_for_a_regular_task(self, client, dag_maker, mocker): + """The row rules the task out, so starting it reads no serialized Dag -- the point of the column.""" + with dag_maker("test_regular_task_arg_bindings", serialized=True): + + @task + def transform(country: str): ... + + transform("uk") + + dr = dag_maker.create_dagrun() + (ti,) = dr.get_task_instances() + assert ti.is_stub is False + ti.set_state(State.QUEUED) + dag_maker.session.flush() + get_dag = mocker.patch.object(DBDagBag, "get_dag", autospec=True) + + response = client.patch(f"/execution/task-instances/{ti.id}/run", json=self.RUN_PAYLOAD) + + assert response.status_code == 200 + assert "arg_bindings" not in response.json() + get_dag.assert_not_called() + + def test_ti_run_treats_an_unset_is_stub_as_a_regular_task(self, client, dag_maker): + """A row written before the column existed carries null and keeps its pre-upgrade behavior.""" + with dag_maker("test_unset_is_stub", serialized=True): + + @task.stub + def transform(country: str): ... + + transform("uk") + + dr = dag_maker.create_dagrun() + (ti,) = dr.get_task_instances() + ti.set_state(State.QUEUED) + dag_maker.session.get(TaskInstance, ti.id).is_stub = None + dag_maker.session.commit() + assert dag_maker.session.scalar(select(TaskInstance.is_stub).where(TaskInstance.id == ti.id)) is None + + response = client.patch(f"/execution/task-instances/{ti.id}/run", json=self.RUN_PAYLOAD) + + assert response.status_code == 200 + assert "arg_bindings" not in response.json() + @mock.patch( "airflow.api_fastapi.execution_api.routes.task_instances.get_arg_bindings", autospec=True, @@ -1160,7 +1203,9 @@ def transform(country: str): ... pytest.param( task.stub(execution_timeout=timedelta(seconds=-30)), 0.0, id="stub-negative-timeout" ), - pytest.param(task(execution_timeout=timedelta(seconds=10)), 10.0, id="python-task"), + # A Python task's timeout stays with the task runner, which sends it to the supervisor, + # so the server resolves one only for a task whose definition it loads anyway. + pytest.param(task(execution_timeout=timedelta(seconds=10)), None, id="python-task"), ], ) def test_ti_run_returns_execution_timeout(self, client, dag_maker, decorator, expected): @@ -2108,6 +2153,8 @@ def test_ti_run_skips_queued_duration_metric_when_the_response_fails( dag_id=str(uuid4()), ) ti.queued_dttm = queued_at + # The route reads a serialized spec only for a stub task, which is what fails below. + ti.is_stub = True session.commit() with mock.patch( diff --git a/airflow-core/tests/unit/dag_processing/lang_sdk_fixtures/java_native.json b/airflow-core/tests/unit/dag_processing/lang_sdk_fixtures/java_native.json index 50314784b10f0..b31c624d458ae 100644 --- a/airflow-core/tests/unit/dag_processing/lang_sdk_fixtures/java_native.json +++ b/airflow-core/tests/unit/dag_processing/lang_sdk_fixtures/java_native.json @@ -27,6 +27,7 @@ "task_type": "DumpExtract", "_task_module": "org.apache.airflow.sdk.execution", "language": "java", + "is_stub": true, "template_fields": [], "retries": 2, "queue": "java", @@ -44,6 +45,7 @@ "task_type": "DumpLoad", "_task_module": "org.apache.airflow.sdk.execution", "language": "java", + "is_stub": true, "template_fields": [], "queue": "java", "trigger_rule": "all_done" @@ -112,6 +114,7 @@ "task_type": "DumpExtract", "_task_module": "org.apache.airflow.sdk.execution", "language": "java", + "is_stub": true, "template_fields": [] } } diff --git a/airflow-core/tests/unit/dag_processing/test_dagbag.py b/airflow-core/tests/unit/dag_processing/test_dagbag.py index d24443f709fe3..72e9da954f4b2 100644 --- a/airflow-core/tests/unit/dag_processing/test_dagbag.py +++ b/airflow-core/tests/unit/dag_processing/test_dagbag.py @@ -1660,6 +1660,17 @@ def test_recorded_runtime_payloads_are_valid(data): DagSerialization.validate_serialized_dag(data) +@pytest.mark.parametrize("data", _recorded_payloads()) +def test_recorded_runtime_payloads_mark_their_tasks_as_stubs(data): + """ti_run derives arg_bindings only for a stub task, so a Dag defined in a language SDK has to say so.""" + dag = DagSerialization.from_dict(data) + for task_id, task in dag.task_dict.items(): + # A Go or Java TriggerDagRun task is the Python operator itself and runs in-process. The + # TypeScript runtime implements its own, so that one stays a stub. + runs_in_python = task.task_type == "TriggerDagRunOperator" + assert task.is_stub is not runs_in_python, task_id + + def _serialize_native_dag(dag_id: str, fileloc: str, relative_fileloc: str) -> dict: with DAG(dag_id, schedule=None) as dag: BaseOperator(task_id="extract") diff --git a/airflow-core/tests/unit/models/test_taskinstance.py b/airflow-core/tests/unit/models/test_taskinstance.py index f5b8116038cb5..25421f08990ae 100644 --- a/airflow-core/tests/unit/models/test_taskinstance.py +++ b/airflow-core/tests/unit/models/test_taskinstance.py @@ -2865,6 +2865,7 @@ def test_refresh_from_db(self, create_task_instance): "priority_weight": 123, "operator": "some_custom_operator", "custom_operator_name": "some_custom_operator", + "is_stub": False, "queued_dttm": run_date + datetime.timedelta(hours=1), "scheduled_dttm": run_date + datetime.timedelta(hours=1), "rendered_map_index": None, @@ -5067,6 +5068,36 @@ def test_make_task_carrier_with_none_carrier(self): assert "traceparent" in carrier +@pytest.mark.parametrize("is_stub", [True, False]) +def test_refresh_from_task_copies_is_stub(is_stub): + """ti_run gates on the row, so the flag has to travel from the task to the TaskInstance.""" + ser_task = create_scheduler_operator(EmptyOperator(task_id="empty")) + ser_task.is_stub = is_stub + + ti = TI(ser_task, run_id=None, dag_version_id=mock.MagicMock()) + + assert ti.is_stub is is_stub + + +@pytest.mark.db_test +@pytest.mark.parametrize( + ("decorator", "expected"), + [pytest.param(task.stub, True, id="stub"), pytest.param(task, False, id="regular")], +) +def test_scheduled_task_instances_carry_is_stub(dag_maker, decorator, expected): + """The scheduler writes the flag when it creates the rows, not only when it refreshes one.""" + with dag_maker("test_scheduled_is_stub", serialized=True): + + @decorator + def transform(country: str): ... + + transform("uk") + + (ti,) = dag_maker.create_dagrun().get_task_instances() + + assert ti.is_stub is expected + + @pytest.mark.db_test def test_insert_mapping_includes_context_carrier(dag_maker, session): """insert_mapping should include a context_carrier with a traceparent derived from the dag run.""" From 940a608b979a38507cea2fa36defd8db5e6b71f3 Mon Sep 17 00:00:00 2001 From: ZHE YOU LIU Date: Tue, 27 Oct 2026 15:15:23 +0000 Subject: [PATCH 2/5] Hold SQLite foreign keys off while the downgrade rebuilds Dropping a column rebuilds the table on SQLite, unlike adding one, so the downgrade copies task_instance rows while other tables still point at it. Co-Authored-By: Claude --- .../versions/0149_3_4_0_add_is_stub_to_task_instance.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/airflow-core/src/airflow/migrations/versions/0149_3_4_0_add_is_stub_to_task_instance.py b/airflow-core/src/airflow/migrations/versions/0149_3_4_0_add_is_stub_to_task_instance.py index 7f48942112a60..3bc0232aa9a20 100644 --- a/airflow-core/src/airflow/migrations/versions/0149_3_4_0_add_is_stub_to_task_instance.py +++ b/airflow-core/src/airflow/migrations/versions/0149_3_4_0_add_is_stub_to_task_instance.py @@ -30,6 +30,8 @@ import sqlalchemy as sa from alembic import op +from airflow.migrations.utils import sqlite_rebuilds + # revision identifiers, used by Alembic. revision = "b4f1a7d2c903" down_revision = "ee86eed19e24" @@ -48,5 +50,7 @@ def upgrade(): def downgrade(): """Remove is_stub from task_instance.""" - with op.batch_alter_table(TABLE_NAME) as batch_op: + # Dropping a column rebuilds the table on SQLite, unlike adding one, so the foreign keys of + # everything that points at task_instance have to be held off while the rows are copied. + with sqlite_rebuilds(op), op.batch_alter_table(TABLE_NAME) as batch_op: batch_op.drop_column("is_stub") From 9058d02849b431008f7dd0edfaa7cccda313c55c Mon Sep 17 00:00:00 2001 From: ZHE YOU LIU Date: Tue, 27 Oct 2026 15:34:02 +0000 Subject: [PATCH 3/5] Pin the bulk-insert path and reword the is_stub comments create_dagrun refreshes every row from its task right after the bulk write, so the existing test passed even when insert_mapping dropped the key. Assert the mapping itself instead, and record what a row keeps when it moves to a new Dag version without being refreshed. Co-Authored-By: Claude --- .../execution_api/datamodels/taskinstance.py | 4 +-- .../execution_api/routes/task_instances.py | 5 ++-- .../src/airflow/models/taskinstance.py | 4 ++- .../lang_sdk_fixtures/java_native.json | 6 ++--- .../tests/unit/models/test_taskinstance.py | 26 ++++++++++++++++++- 5 files changed, 36 insertions(+), 9 deletions(-) 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 e6d771f0f82b0..c90d64da1ec3f 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 @@ -560,8 +560,8 @@ class TIRunContext(BaseModel): """ The ``execution_timeout`` in seconds of a task a language SDK runs, from its serialized Dag. - The supervisor enforces the timeout with this value. A Python task keeps its own timer in the - task runner, so the field is unset for one. ``None`` when no timeout is set. + The supervisor enforces the timeout with this value. A Python task runs its own timer in the task + runner instead, and gets no value here. ``None`` when no timeout is set. """ multi_team: bool = False 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 c891f43ca6ac2..e305a615f9dab 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 @@ -383,8 +383,9 @@ def ti_run( if first_task_reschedule_start_date is not None: context.first_task_reschedule_start_date = first_task_reschedule_start_date - # Only a lang-SDK (foreign-runtime) task needs its definition here. ti.is_stub is the row's - # copy of the serialized flag, so a regular task starts without loading the serialized Dag. + # Avoid loading the serialized Dag on the normal Python task path: only a lang-SDK + # (foreign-runtime) task needs its definition here, and ti.is_stub carries the row's own copy + # of the serialized flag. serialized_task = get_task(dag_bag, ti, session=session) if ti.is_stub else None # Only set for lang-SDK (foreign-runtime) tasks with a captured TaskFlow arg diff --git a/airflow-core/src/airflow/models/taskinstance.py b/airflow-core/src/airflow/models/taskinstance.py index a1b48e8bf7a49..78d1cb1acaf01 100644 --- a/airflow-core/src/airflow/models/taskinstance.py +++ b/airflow-core/src/airflow/models/taskinstance.py @@ -1142,7 +1142,9 @@ class TaskInstance(Base, LoggingMixin, BaseWorkload): operator: Mapped[str | None] = mapped_column(String(1000), nullable=True) custom_operator_name: Mapped[str | None] = mapped_column(String(1000), nullable=True) # Denormalized from the task so ti_run can tell a stub (foreign-runtime) task from a regular one - # without loading the serialized Dag. Null on rows written before the column existed. + # without loading the serialized Dag. Null on rows written before the column existed. A row moved + # to a new Dag version without being refreshed from its task keeps the value it was created with, + # which fails closed: a task that has become a stub loses its arg_bindings rather than running wrong. is_stub: Mapped[bool | None] = mapped_column(Boolean, nullable=True) queued_dttm: Mapped[datetime | None] = mapped_column(UtcDateTime, nullable=True) scheduled_dttm: Mapped[datetime | None] = mapped_column(UtcDateTime, nullable=True) diff --git a/airflow-core/tests/unit/dag_processing/lang_sdk_fixtures/java_native.json b/airflow-core/tests/unit/dag_processing/lang_sdk_fixtures/java_native.json index b31c624d458ae..53f0a7256d77d 100644 --- a/airflow-core/tests/unit/dag_processing/lang_sdk_fixtures/java_native.json +++ b/airflow-core/tests/unit/dag_processing/lang_sdk_fixtures/java_native.json @@ -27,7 +27,7 @@ "task_type": "DumpExtract", "_task_module": "org.apache.airflow.sdk.execution", "language": "java", - "is_stub": true, + "is_stub": true, "template_fields": [], "retries": 2, "queue": "java", @@ -45,7 +45,7 @@ "task_type": "DumpLoad", "_task_module": "org.apache.airflow.sdk.execution", "language": "java", - "is_stub": true, + "is_stub": true, "template_fields": [], "queue": "java", "trigger_rule": "all_done" @@ -114,7 +114,7 @@ "task_type": "DumpExtract", "_task_module": "org.apache.airflow.sdk.execution", "language": "java", - "is_stub": true, + "is_stub": true, "template_fields": [] } } diff --git a/airflow-core/tests/unit/models/test_taskinstance.py b/airflow-core/tests/unit/models/test_taskinstance.py index 25421f08990ae..8087538ff8bc1 100644 --- a/airflow-core/tests/unit/models/test_taskinstance.py +++ b/airflow-core/tests/unit/models/test_taskinstance.py @@ -5085,7 +5085,7 @@ def test_refresh_from_task_copies_is_stub(is_stub): [pytest.param(task.stub, True, id="stub"), pytest.param(task, False, id="regular")], ) def test_scheduled_task_instances_carry_is_stub(dag_maker, decorator, expected): - """The scheduler writes the flag when it creates the rows, not only when it refreshes one.""" + """A scheduled row carries the flag, whichever of the two write paths filled it in.""" with dag_maker("test_scheduled_is_stub", serialized=True): @decorator @@ -5098,6 +5098,30 @@ def transform(country: str): ... assert ti.is_stub is expected +@pytest.mark.db_test +@pytest.mark.parametrize("is_stub", [True, False]) +def test_insert_mapping_carries_is_stub(dag_maker, session, is_stub): + """The scheduler bulk-inserts rows from this mapping, so the flag has to be in it.""" + with dag_maker("test_insert_mapping_is_stub"): + EmptyOperator(task_id="t1") + session.flush() + + op = create_scheduler_operator(dag_maker.dag.get_task("t1")) + op.is_stub = is_stub + dag_run = mock.MagicMock() + dag_run.context_carrier = new_dagrun_trace_carrier() + + mapping = TaskInstance.insert_mapping( + run_id="test_run", + task=op, + map_index=0, + dag_version_id=None, + dag_run=dag_run, + ) + + assert mapping["is_stub"] is is_stub + + @pytest.mark.db_test def test_insert_mapping_includes_context_carrier(dag_maker, session): """insert_mapping should include a context_carrier with a traceparent derived from the dag run.""" From 78d04b4d0d74acd8b2dd5861fcfd8fc83b3bd4c7 Mon Sep 17 00:00:00 2001 From: LIU ZHE YOU Date: Tue, 27 Oct 2026 16:30:58 +0000 Subject: [PATCH 4/5] Mark a Go or Java TriggerDagRun task a stub, as TypeScript does Co-Authored-By: Claude --- .../versions/head/test_task_instances.py | 2 -- .../dag_processing/lang_sdk_fixtures/go_native.json | 2 ++ airflow-core/tests/unit/dag_processing/test_dagbag.py | 8 +++----- go-sdk/airflow/serialize.go | 9 +++++---- go-sdk/airflow/serialize_test.go | 3 ++- .../kotlin/org/apache/airflow/sdk/execution/Serde.kt | 11 ++++++----- .../org/apache/airflow/sdk/TriggerDagRunTest.kt | 5 +++-- ts-sdk/src/coordinator/serde.ts | 8 +++++--- 8 files changed, 26 insertions(+), 22 deletions(-) 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 1417b15f88f03..c40f69fcc6038 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 @@ -2153,8 +2153,6 @@ def test_ti_run_skips_queued_duration_metric_when_the_response_fails( dag_id=str(uuid4()), ) ti.queued_dttm = queued_at - # The route reads a serialized spec only for a stub task, which is what fails below. - ti.is_stub = True session.commit() with mock.patch( diff --git a/airflow-core/tests/unit/dag_processing/lang_sdk_fixtures/go_native.json b/airflow-core/tests/unit/dag_processing/lang_sdk_fixtures/go_native.json index 031f73bebd1a7..8b440945e3bc3 100644 --- a/airflow-core/tests/unit/dag_processing/lang_sdk_fixtures/go_native.json +++ b/airflow-core/tests/unit/dag_processing/lang_sdk_fixtures/go_native.json @@ -1621,6 +1621,7 @@ "ratio": 0.5, "rows": 2 }, + "is_stub": true, "logical_date": "2026-09-30 00:00:00+00:00", "skip_when_already_exists": true, "task_id": "trigger_all", @@ -1649,6 +1650,7 @@ "Triggered DAG": "_link_TriggerDagRunLink" }, "_task_module": "airflow.providers.standard.operators.trigger_dagrun", + "is_stub": true, "logical_date": "NOTSET", "skip_when_already_exists": false, "task_id": "trigger_default", diff --git a/airflow-core/tests/unit/dag_processing/test_dagbag.py b/airflow-core/tests/unit/dag_processing/test_dagbag.py index 72e9da954f4b2..7bcd745c24cd0 100644 --- a/airflow-core/tests/unit/dag_processing/test_dagbag.py +++ b/airflow-core/tests/unit/dag_processing/test_dagbag.py @@ -1662,13 +1662,11 @@ def test_recorded_runtime_payloads_are_valid(data): @pytest.mark.parametrize("data", _recorded_payloads()) def test_recorded_runtime_payloads_mark_their_tasks_as_stubs(data): - """ti_run derives arg_bindings only for a stub task, so a Dag defined in a language SDK has to say so.""" + """ti_run derives arg_bindings only for a stub task, and a language runtime runs every task of a + Dag defined in it, a TriggerDagRun one included, so each has to say so.""" dag = DagSerialization.from_dict(data) for task_id, task in dag.task_dict.items(): - # A Go or Java TriggerDagRun task is the Python operator itself and runs in-process. The - # TypeScript runtime implements its own, so that one stays a stub. - runs_in_python = task.task_type == "TriggerDagRunOperator" - assert task.is_stub is not runs_in_python, task_id + assert task.is_stub, task_id def _serialize_native_dag(dag_id: str, fileloc: str, relative_fileloc: str) -> dict: diff --git a/go-sdk/airflow/serialize.go b/go-sdk/airflow/serialize.go index bd294b85b4f8e..0eec1a44a3b16 100644 --- a/go-sdk/airflow/serialize.go +++ b/go-sdk/airflow/serialize.go @@ -424,9 +424,6 @@ func (d *DagRef) serializeTaskLocked(task *TaskRef) map[string]any { // Python writes template_fields for every operator. A Go task has no template fields, so the list // is empty. data["template_fields"] = []string{} - // is_stub makes the API server send the task its arguments from _arg_bindings when the - // task runs, as it does for a @task.stub task of a Python Dag. - data["is_stub"] = true if len(task.inputs) > 0 { data["_arg_bindings"] = serializeArgBindings(task.inputs) } @@ -437,7 +434,11 @@ func (d *DagRef) serializeTaskLocked(task *TaskRef) map[string]any { data["_can_skip_downstream"] = true } } - // A task from TriggerDagRun takes the queue of the Dag too, because the Go runtime runs it. + // The Go runtime runs every task of a Go Dag, a task from TriggerDagRun included, so every one of + // them is a stub. is_stub makes the API server send the task its arguments from _arg_bindings when + // the task runs, as it does for a @task.stub task of a Python Dag. + data["is_stub"] = true + // The worker picks a runtime by queue, so a task from TriggerDagRun takes the queue of the Dag too. if spec.Queue == "" { spec.Queue = d.spec.Queue } diff --git a/go-sdk/airflow/serialize_test.go b/go-sdk/airflow/serialize_test.go index 4f728ada5d899..4a7bcf96dd84d 100644 --- a/go-sdk/airflow/serialize_test.go +++ b/go-sdk/airflow/serialize_test.go @@ -925,6 +925,7 @@ func TestSerializeWritesATriggerDagRunAsATriggerDagRunOperator(t *testing.T) { assertJSON(t, `{ "task_id": "trigger", "task_type": "TriggerDagRunOperator", + "is_stub": true, "_task_module": "airflow.providers.standard.operators.trigger_dagrun", "ui_color": "#ffefeb", "template_fields": [ @@ -1012,7 +1013,7 @@ func TestSerializeWritesOnlyTheTriggerDagRunOptionsThatAreSet(t *testing.T) { "skip_when_already_exists": false }`, withoutKeys(serializedTask(t, got, "trigger"), "task_id", "task_type", "_task_module", "ui_color", "template_fields", - "template_fields_renderers", "_operator_extra_links", + "template_fields_renderers", "_operator_extra_links", "is_stub", )) // An empty FailedStates means that no Dag run state fails the task. assertJSON(t, `[]`, serializedTask(t, got, "trigger_never_fails")["failed_states"]) diff --git a/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/execution/Serde.kt b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/execution/Serde.kt index c01592b8eea6e..8a5d470426abe 100644 --- a/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/execution/Serde.kt +++ b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/execution/Serde.kt @@ -145,15 +145,16 @@ private fun serializeTask( // value never matches the tuple default it is compared against), so it // is unconditional here too. Java tasks have no template fields. data["template_fields"] = emptyList() - // What marks a task whose arguments Airflow resolves per instance for a - // runtime outside Python, as `@task.stub` does on the Python side. - // `get_arg_bindings` reads nothing without it. - data["is_stub"] = true argBindings(taskId, def)?.let { data["_arg_bindings"] = it } } + // The Java runtime runs every task of a Java Dag, a trigger task included, so + // every one of them is a stub. is_stub marks a task whose arguments Airflow + // resolves per instance for a runtime outside Python, as `@task.stub` does on + // the Python side. `get_arg_bindings` reads nothing without it. + data["is_stub"] = true // Lets NotPreviouslySkippedDep re-skip a cleared downstream, as Python's SkipMixin does. if (def.decider != null) data["_can_skip_downstream"] = true - // A trigger task takes the Dag's queue too, because the Java runtime runs it. + // The worker picks a runtime by queue, so a trigger task takes the Dag's queue too. val config = if (dagQueue == null || "queue" in def.configValues) def.configValues else def.configValues + ("queue" to dagQueue) // Emit only config entries that differ from their schema default, mirroring diff --git a/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/TriggerDagRunTest.kt b/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/TriggerDagRunTest.kt index 56c0bb410b7a1..b7bef8724f752 100644 --- a/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/TriggerDagRunTest.kt +++ b/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/TriggerDagRunTest.kt @@ -168,9 +168,10 @@ internal class TriggerDagRunTest { assertEquals("NOTSET", data["logical_date"]) assertEquals(false, data["wait_for_completion"]) assertEquals(false, data["skip_when_already_exists"]) - // A Java task's markers belong to a task the Java runtime runs a body for. + // The Java runtime runs a trigger task too, so it carries the stub marker. + assertEquals(true, data["is_stub"]) + // `language` marks a task the Java runtime runs a body for, which this is not. assertFalse("language" in data) - assertFalse("is_stub" in data) assertFalse("reset_dag_run" in data) } diff --git a/ts-sdk/src/coordinator/serde.ts b/ts-sdk/src/coordinator/serde.ts index c528980487cdc..b6938fa81d66d 100644 --- a/ts-sdk/src/coordinator/serde.ts +++ b/ts-sdk/src/coordinator/serde.ts @@ -252,9 +252,11 @@ function serializeTask( // matches the tuple default it is compared against. A TypeScript task has // no Jinja templating, so the list is empty rather than absent. template_fields: [], - // What marks a task whose arguments the API server resolves per instance - // and sends to a foreign runtime, as `@task.stub` does on the Python side. - // `get_arg_bindings` reads nothing without it. + // The TypeScript runtime runs every task of a TypeScript Dag, a trigger task + // included, so every one of them is a stub. `is_stub` marks a task whose + // arguments the API server resolves per instance and sends to a foreign + // runtime, as `@task.stub` does on the Python side. `get_arg_bindings` reads + // nothing without it. is_stub: true, }; const label = `task "${taskId}" of Dag "${dagId}"`; From 978a6bca6e475527cd0517ca2ce96b29148613fb Mon Sep 17 00:00:00 2001 From: LIU ZHE YOU Date: Wed, 28 Oct 2026 02:05:55 +0000 Subject: [PATCH 5/5] Inline the table name in the is_stub migration --- .../versions/0149_3_4_0_add_is_stub_to_task_instance.py | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/airflow-core/src/airflow/migrations/versions/0149_3_4_0_add_is_stub_to_task_instance.py b/airflow-core/src/airflow/migrations/versions/0149_3_4_0_add_is_stub_to_task_instance.py index 3bc0232aa9a20..48cc4c5870e96 100644 --- a/airflow-core/src/airflow/migrations/versions/0149_3_4_0_add_is_stub_to_task_instance.py +++ b/airflow-core/src/airflow/migrations/versions/0149_3_4_0_add_is_stub_to_task_instance.py @@ -39,12 +39,10 @@ depends_on = None airflow_version = "3.4.0" -TABLE_NAME = "task_instance" - def upgrade(): """Add is_stub to task_instance.""" - with op.batch_alter_table(TABLE_NAME) as batch_op: + with op.batch_alter_table("task_instance") as batch_op: batch_op.add_column(sa.Column("is_stub", sa.Boolean(), nullable=True)) @@ -52,5 +50,5 @@ def downgrade(): """Remove is_stub from task_instance.""" # Dropping a column rebuilds the table on SQLite, unlike adding one, so the foreign keys of # everything that points at task_instance have to be held off while the rows are copied. - with sqlite_rebuilds(op), op.batch_alter_table(TABLE_NAME) as batch_op: + with sqlite_rebuilds(op), op.batch_alter_table("task_instance") as batch_op: batch_op.drop_column("is_stub")