Add indexes on serialized_dag and dag_code

airflow-core/docs/migrations-ref.rst

@@ -45,7 +45,9 @@ Here's the list of all the Database Migrations that are executed via when you ru
 +-------------------------+------------------+-------------------+--------------------------------------------------------------+
 | ``436dc127462c``        | ``5a5d3253e946`` | ``3.4.0``         | Drop span_status column.                                     |
 +-------------------------+------------------+-------------------+--------------------------------------------------------------+
-| ``5a5d3253e946``        | ``d2f4e1b3c5a7`` | ``3.4.0``         | Add GIN index on asset_event.extra for PostgreSQL.           |
+| ``5a5d3253e946``        | ``3c525f44bea8`` | ``3.4.0``         | Add GIN index on asset_event.extra for PostgreSQL.           |
++-------------------------+------------------+-------------------+--------------------------------------------------------------+
+| ``3c525f44bea8``        | ``d2f4e1b3c5a7`` | ``3.3.1``         | Add indexes on serialized_dag and dag_code.                  |
 +-------------------------+------------------+-------------------+--------------------------------------------------------------+
 | ``d2f4e1b3c5a7``        | ``9ff64e1c35d3`` | ``3.3.0``         | Add partition_date to asset_partition_dag_run.               |
 +-------------------------+------------------+-------------------+--------------------------------------------------------------+

airflow-core/src/airflow/migrations/versions/0124_3_3_1_add_indexes_on_serialized_dag_and_dag_.py

@@ -0,0 +1,55 @@
+#
+# 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 on serialized_dag and dag_code.
+
+Revision ID: 3c525f44bea8
+Revises: d2f4e1b3c5a7
+Create Date: 2026-07-04 18:18:11.140047
+
+"""
+
+from __future__ import annotations
+
+from alembic import op
+
+# revision identifiers, used by Alembic.
+revision = "3c525f44bea8"
+down_revision = "d2f4e1b3c5a7"
+branch_labels = None
+depends_on = None
+airflow_version = "3.3.1"
+
+
+def upgrade():
+    """Apply Add indexes on serialized_dag and dag_code."""
+    with op.batch_alter_table("dag_code", schema=None) as batch_op:
+        batch_op.create_index("idx_dag_code_dag_id", ["dag_id"], unique=False)
+
+    with op.batch_alter_table("serialized_dag", schema=None) as batch_op:
+        batch_op.create_index("idx_serialized_dag_dag_id", ["dag_id"], unique=False)
+
+
+def downgrade():
+    """Unapply Add indexes on serialized_dag and dag_code."""
+    with op.batch_alter_table("serialized_dag", schema=None) as batch_op:
+        batch_op.drop_index("idx_serialized_dag_dag_id")
+
+    with op.batch_alter_table("dag_code", schema=None) as batch_op:
+        batch_op.drop_index("idx_dag_code_dag_id")

airflow-core/src/airflow/migrations/versions/0124_3_4_0_add_gin_index_on_asset_event_extra.py → airflow-core/src/airflow/migrations/versions/0125_3_4_0_add_gin_index_on_asset_event_extra.py

@@ -20,7 +20,7 @@
 Add GIN index on asset_event.extra for PostgreSQL.
 
 Revision ID: 5a5d3253e946
-Revises: d2f4e1b3c5a7
+Revises: 3c525f44bea8
 Create Date: 2026-04-01 23:00:00.000000
 
 """
@@ -32,7 +32,7 @@
 
 # revision identifiers, used by Alembic.
 revision = "5a5d3253e946"
-down_revision = "d2f4e1b3c5a7"
+down_revision = "3c525f44bea8"
 branch_labels = None
 depends_on = None
 airflow_version = "3.4.0"

airflow-core/src/airflow/models/dagcode.py

@@ -23,7 +23,7 @@
 
 import sqlalchemy as sa
 import uuid6
-from sqlalchemy import ForeignKey, String, Text, select
+from sqlalchemy import ForeignKey, Index, String, Text, select
 from sqlalchemy.dialects.mysql import MEDIUMTEXT
 from sqlalchemy.orm import Mapped, mapped_column, relationship
 from sqlalchemy.sql.expression import literal
@@ -70,6 +70,7 @@ class DagCode(Base):
         sa.Uuid(), ForeignKey("dag_version.id", ondelete="CASCADE"), nullable=False, unique=True
     )
     dag_version = relationship("DagVersion", back_populates="dag_code", uselist=False)
+    __table_args__ = (Index("idx_dag_code_dag_id", dag_id),)
 
     def __init__(self, dag_version, full_filepath: str, source_code: str | None = None):
         self.dag_version = dag_version

airflow-core/src/airflow/models/serialized_dag.py

@@ -27,7 +27,7 @@
 from uuid import UUID
 
 import uuid6
-from sqlalchemy import JSON, ForeignKey, LargeBinary, String, Uuid, exists, select, tuple_, update
+from sqlalchemy import JSON, ForeignKey, Index, LargeBinary, String, Uuid, exists, select, tuple_, update
 from sqlalchemy.dialects.postgresql import ARRAY, JSONB
 from sqlalchemy.orm import Mapped, backref, foreign, mapped_column, relationship
 from sqlalchemy.sql.expression import func, literal
@@ -343,6 +343,7 @@ class SerializedDagModel(Base):
     )
 
     load_op_links = True
+    __table_args__ = (Index("idx_serialized_dag_dag_id", dag_id),)
 
     def __init__(self, dag: LazyDeserializedDAG) -> None:
         self.dag_id = dag.dag_id

airflow-core/src/airflow/utils/db.py

@@ -117,6 +117,7 @@ class MappedClassProtocol(Protocol):
     "3.1.8": "509b94a1042d",
     "3.2.0": "1d6611b6ab7c",
     "3.3.0": "d2f4e1b3c5a7",
+    "3.3.1": "3c525f44bea8",
     "3.4.0": "7a98f1b7dbd3",
 }
 

airflow-core/tests/unit/migrations/test_0124_add_indexes_on_serialized_dag_and_dag_code.py

@@ -0,0 +1,55 @@
+#
+# 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
+import sqlalchemy as sa
+
+from airflow import settings
+from airflow.utils.db import downgrade, upgradedb
+
+pytestmark = pytest.mark.db_test
+
+_REVISION = "3c525f44bea8"
+_DOWN_REVISION = "d2f4e1b3c5a7"
+
+_EXPECTED_INDEXES = {
+    "idx_serialized_dag_dag_id": ("serialized_dag", "dag_id"),
+    "idx_dag_code_dag_id": ("dag_code", "dag_id"),
+}
+
+
+class TestMigration0124AddDagIdIndexes:
+    @pytest.fixture(autouse=True)
+    def _restore_head(self):
+        yield
+        upgradedb()
+
+    @staticmethod
+    def _get_indexes(table):
+        with settings.engine.connect() as conn:
+            return {ix["name"]: ix["column_names"] for ix in sa.inspect(conn).get_indexes(table)}
+
+    def test_upgrade_creates_indexes_and_downgrade_drops_them(self):
+        downgrade(to_revision=_DOWN_REVISION)
+        for index, (table, _) in _EXPECTED_INDEXES.items():
+            assert index not in self._get_indexes(table)
+
+        upgradedb(to_revision=_REVISION)
+        for index, (table, column) in _EXPECTED_INDEXES.items():
+            assert self._get_indexes(table).get(index) == [column]