Add index on task_reschedule ti_id (#60931)

guan404ming · ephraimbuddy · commit e7467dc8a344 · 2026-02-17T11:10:00.000+01:00
(cherry picked from commit 14e811c)
diff --git a/airflow-core/docs/img/airflow_erd.sha256 b/airflow-core/docs/img/airflow_erd.sha256
@@ -1 +1 @@
-4827b81c01c9a7993d9796da4990f2a950c2069687db4b0e9cba0c0ae851bd9c
+4afcc3a71762abf79a5fb3808aab1d52e7a6f170a74d1002e9663f2f35d00d6b
diff --git a/airflow-core/docs/img/airflow_erd.svg b/airflow-core/docs/img/airflow_erd.svg
diff --git a/airflow-core/docs/migrations-ref.rst b/airflow-core/docs/migrations-ref.rst
@@ -39,7 +39,9 @@ Here's the list of all the Database Migrations that are executed via when you ru
 +-------------------------+------------------+-------------------+--------------------------------------------------------------+
 | Revision ID             | Revises ID       | Airflow Version   | Description                                                  |
 +=========================+==================+===================+==============================================================+
-| ``cc92b33c6709`` (head) | ``eaf332f43c7c`` | ``3.1.0``         | Add backward compatibility for serialized DAG format v3 to   |
+| ``82dbd68e6171`` (head) | ``cc92b33c6709`` | ``3.1.8``         | Add index to task_reschedule ti_id .                         |
++-------------------------+------------------+-------------------+--------------------------------------------------------------+
+| ``cc92b33c6709``        | ``eaf332f43c7c`` | ``3.1.0``         | Add backward compatibility for serialized DAG format v3 to   |
 |                         |                  |                   | v2.                                                          |
 +-------------------------+------------------+-------------------+--------------------------------------------------------------+
 | ``eaf332f43c7c``        | ``a3c7f2b18d4e`` | ``3.1.0``         | add last_parse_duration to dag model.                        |
diff --git a/airflow-core/src/airflow/migrations/versions/0086_3_1_8_add_index_to_task_reschedule_ti_id.py b/airflow-core/src/airflow/migrations/versions/0086_3_1_8_add_index_to_task_reschedule_ti_id.py
@@ -0,0 +1,58 @@
+#
+# 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 index to task_reschedule ti_id .
+
+Revision ID: 82dbd68e6171
+Revises: cc92b33c6709
+Create Date: 2026-01-22 16:25:42.164449
+
+"""
+
+from __future__ import annotations
+
+from alembic import op
+
+# revision identifiers, used by Alembic.
+revision = "82dbd68e6171"
+down_revision = "cc92b33c6709"
+branch_labels = None
+depends_on = None
+airflow_version = "3.1.8"
+
+
+def upgrade():
+    """Add index to task_reschedule ti_id."""
+    with op.batch_alter_table("task_reschedule", schema=None) as batch_op:
+        batch_op.create_index("idx_task_reschedule_ti_id", ["ti_id"], unique=False)
+
+
+def downgrade():
+    """Remove index from task_reschedule ti_id."""
+    dialect_name = op.get_context().dialect.name
+    if dialect_name == "mysql":
+        with op.batch_alter_table("task_reschedule", schema=None) as batch_op:
+            batch_op.drop_constraint("task_reschedule_ti_fkey", type_="foreignkey")
+            batch_op.drop_index("idx_task_reschedule_ti_id")
+            batch_op.create_foreign_key(
+                "task_reschedule_ti_fkey", "task_instance", ["ti_id"], ["id"], ondelete="CASCADE"
+            )
+    else:
+        with op.batch_alter_table("task_reschedule", schema=None) as batch_op:
+            batch_op.drop_index("idx_task_reschedule_ti_id")
diff --git a/airflow-core/src/airflow/models/taskreschedule.py b/airflow-core/src/airflow/models/taskreschedule.py
@@ -25,6 +25,7 @@
 from sqlalchemy import (
     Column,
     ForeignKey,
+    Index,
     Integer,
     String,
     asc,
@@ -60,6 +61,8 @@ class TaskReschedule(Base):
     duration = Column(Integer, nullable=False)
     reschedule_date = Column(UtcDateTime, nullable=False)
 
+    __table_args__ = (Index("idx_task_reschedule_ti_id", ti_id),)
+
     task_instance = relationship(
         "TaskInstance", primaryjoin="TaskReschedule.ti_id == foreign(TaskInstance.id)", uselist=False
     )
diff --git a/airflow-core/src/airflow/utils/db.py b/airflow-core/src/airflow/utils/db.py
@@ -111,6 +111,7 @@ class MappedClassProtocol(Protocol):
     "3.0.0": "29ce7909c52b",
     "3.0.3": "fe199e1abd77",
     "3.1.0": "cc92b33c6709",
+    "3.1.8": "82dbd68e6171",
 }
 
 

Original file line number	Diff line number	Diff line change
`@@ -1 +1 @@`
`1`		`-4827b81c01c9a7993d9796da4990f2a950c2069687db4b0e9cba0c0ae851bd9c`
	`1`	`+4afcc3a71762abf79a5fb3808aab1d52e7a6f170a74d1002e9663f2f35d00d6b`
Original file line number	Diff line number	Diff line change
`@@ -111,6 +111,7 @@ class MappedClassProtocol(Protocol):`
`111`	`111`	`"3.0.0": "29ce7909c52b",`
`112`	`112`	`"3.0.3": "fe199e1abd77",`
`113`	`113`	`"3.1.0": "cc92b33c6709",`
	`114`	`+ "3.1.8": "82dbd68e6171",`
`114`	`115`	`}`
`115`	`116`
`116`	`117`