Add indexes for CASCADE deletes for task_instance

airflow/migrations/versions/0111_2_3_3_add_indexes_for_cascade_deletes.py

@@ -0,0 +1,94 @@
+#
+# 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 indexes for CASCADE deletes on task_instance
+
+Some databases don't add indexes on the FK columns so we have to add them for performance on CASCADE deletes.
+
+Revision ID: f5fcbda3e651
+Revises: 3c94c427fdf6
+Create Date: 2022-06-15 18:04:54.081789
+
+"""
+
+from alembic import context, op
+
+# revision identifiers, used by Alembic.
+revision = 'f5fcbda3e651'
+down_revision = '3c94c427fdf6'
+branch_labels = None
+depends_on = None
+airflow_version = '2.3.3'
+
+
+def _mysql_tables_where_indexes_already_present(conn):
+    """
+    If user downgraded and is upgrading again, we have to check for existing
+    indexes on mysql because we can't (and don't) drop them as part of the
+    downgrade.
+    """
+    to_check = [
+        ('xcom', 'idx_xcom_task_instance'),
+        ('task_reschedule', 'idx_task_reschedule_dag_run'),
+        ('task_fail', 'idx_task_fail_task_instance'),
+    ]
+    tables = set()
+    for tbl, idx in to_check:
+        if conn.execute(f"show indexes from {tbl} where Key_name = '{idx}'").first():
+            tables.add(tbl)
+    return tables
+
+
+def upgrade():
+    """Apply Add indexes for CASCADE deletes"""
+    conn = op.get_bind()
+    tables_to_skip = set()
+
+    # mysql requires indexes for FKs, so adding had the effect of renaming, and we cannot remove.
+    if conn.dialect.name == 'mysql' and not context.is_offline_mode():
+        tables_to_skip.update(_mysql_tables_where_indexes_already_present(conn))
+
+    if 'task_fail' not in tables_to_skip:
+        with op.batch_alter_table('task_fail', schema=None) as batch_op:
+            batch_op.create_index('idx_task_fail_task_instance', ['dag_id', 'task_id', 'run_id', 'map_index'])
+
+    if 'task_reschedule' not in tables_to_skip:
+        with op.batch_alter_table('task_reschedule', schema=None) as batch_op:
+            batch_op.create_index('idx_task_reschedule_dag_run', ['dag_id', 'run_id'])
+
+    if 'xcom' not in tables_to_skip:
+        with op.batch_alter_table('xcom', schema=None) as batch_op:
+            batch_op.create_index('idx_xcom_task_instance', ['dag_id', 'task_id', 'run_id', 'map_index'])
+
+
+def downgrade():
+    """Unapply Add indexes for CASCADE deletes"""
+    conn = op.get_bind()
+
+    # mysql requires indexes for FKs, so adding had the effect of renaming, and we cannot remove.
+    if conn.dialect.name == 'mysql':
+        return
+
+    with op.batch_alter_table('xcom', schema=None) as batch_op:
+        batch_op.drop_index('idx_xcom_task_instance')
+
+    with op.batch_alter_table('task_reschedule', schema=None) as batch_op:
+        batch_op.drop_index('idx_task_reschedule_dag_run')
+
+    with op.batch_alter_table('task_fail', schema=None) as batch_op:
+        batch_op.drop_index('idx_task_fail_task_instance')

airflow/models/taskfail.py

@@ -17,7 +17,7 @@
 # under the License.
 """Taskfail tracks the failed run durations of each task instance"""
 
-from sqlalchemy import Column, ForeignKeyConstraint, Integer
+from sqlalchemy import Column, ForeignKeyConstraint, Index, Integer
 from sqlalchemy.orm import relationship
 
 from airflow.models.base import Base, StringID
@@ -39,6 +39,7 @@ class TaskFail(Base):
     duration = Column(Integer)
 
     __table_args__ = (
+        Index("idx_task_fail_task_instance", dag_id, task_id, run_id, map_index),
         ForeignKeyConstraint(
             [dag_id, task_id, run_id, map_index],
             [

airflow/models/taskreschedule.py

@@ -61,6 +61,7 @@ class TaskReschedule(Base):
             name="task_reschedule_ti_fkey",
             ondelete="CASCADE",
         ),
+        Index('idx_task_reschedule_dag_run', dag_id, run_id),
         ForeignKeyConstraint(
             [dag_id, run_id],
             ['dag_run.dag_id', 'dag_run.run_id'],

airflow/models/xcom.py

@@ -72,6 +72,7 @@ class BaseXCom(Base, LoggingMixin):
         # but it goes over MySQL's index length limit. So we instead index 'key'
         # separately, and enforce uniqueness with DagRun.id instead.
         Index("idx_xcom_key", key),
+        Index("idx_xcom_task_instance", dag_id, task_id, run_id, map_index),
         ForeignKeyConstraint(
             [dag_id, task_id, run_id, map_index],
             [

docs/apache-airflow/migrations-ref.rst

@@ -27,7 +27,9 @@ Here's the list of all the Database Migrations that are executed via when you ru
 +---------------------------------+-------------------+-------------------+--------------------------------------------------------------+
 | Revision ID                     | Revises ID        | Airflow Version   | Description                                                  |
 +=================================+===================+===================+==============================================================+
-| ``3c94c427fdf6`` (head)         | ``1de7bc13c950``  | ``2.3.2``         | Add cascade to dag_tag foreign key                           |
+| ``f5fcbda3e651`` (head)         | ``3c94c427fdf6``  | ``2.3.3``         | Add indexes for CASCADE deletes on task_instance             |
++---------------------------------+-------------------+-------------------+--------------------------------------------------------------+
+| ``3c94c427fdf6``                | ``1de7bc13c950``  | ``2.3.2``         | Add cascade to dag_tag foreign key                           |
 +---------------------------------+-------------------+-------------------+--------------------------------------------------------------+
 | ``1de7bc13c950``                | ``b1b348e02d07``  | ``2.3.1``         | Add index for ``event`` column in ``log`` table.             |
 +---------------------------------+-------------------+-------------------+--------------------------------------------------------------+