Load the serialized Dag in ti_run only for stub tasks (#75795)

* 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.

* 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.

* 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.

* Mark a Go or Java TriggerDagRun task a stub, as TypeScript does

* Inline the table name in the is_stub migration

---------

Author: jason810496

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.                          |

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 runs its own timer in the task
+    runner instead, and gets no value here. ``None`` when no timeout is set.
     """
 
     multi_team: bool = False

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,10 @@ 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)
+        # 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
         # spec; the route excludes unset fields, keeping regular responses lean.
@@ -404,6 +408,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
 

airflow-core/src/airflow/migrations/versions/0149_3_4_0_add_is_stub_to_task_instance.py

@@ -0,0 +1,54 @@
+#
+# 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
+
+from airflow.migrations.utils import sqlite_rebuilds
+
+# revision identifiers, used by Alembic.
+revision = "b4f1a7d2c903"
+down_revision = "ee86eed19e24"
+branch_labels = None
+depends_on = None
+airflow_version = "3.4.0"
+
+
+def upgrade():
+    """Add is_stub to task_instance."""
+    with op.batch_alter_table("task_instance") as batch_op:
+        batch_op.add_column(sa.Column("is_stub", sa.Boolean(), nullable=True))
+
+
+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("task_instance") as batch_op:
+        batch_op.drop_column("is_stub")

…ions/0149_3_5_0_add_session_id_to_job.py …ions/0150_3_5_0_add_session_id_to_job.pyairflow-core/src/airflow/migrations/versions/0149_3_5_0_add_session_id_to_job.py renamed to airflow-core/src/airflow/migrations/versions/0150_3_5_0_add_session_id_to_job.py airflow-core/src/airflow/migrations/versions/0149_3_5_0_add_session_id_to_job.py renamed to 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"

…_0_add_dag_run_dag_id_run_after_index.py …_0_add_dag_run_dag_id_run_after_index.pyairflow-core/src/airflow/migrations/versions/0150_3_5_0_add_dag_run_dag_id_run_after_index.py renamed to airflow-core/src/airflow/migrations/versions/0151_3_5_0_add_dag_run_dag_id_run_after_index.py airflow-core/src/airflow/migrations/versions/0150_3_5_0_add_dag_run_dag_id_run_after_index.py renamed to airflow-core/src/airflow/migrations/versions/0151_3_5_0_add_dag_run_dag_id_run_after_index.py

@@ -1141,6 +1141,11 @@ 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. 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)
     queued_by_job_id: Mapped[int | None] = mapped_column(Integer, nullable=True)
@@ -1385,6 +1390,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 +1587,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)
 

airflow-core/src/airflow/models/taskinstance.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",
 }
 

airflow-core/src/airflow/utils/db.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):

airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_task_instances.py

@@ -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",