From f05d8560a6afe3b4f7f92c0124e5aba3cc3fb96e Mon Sep 17 00:00:00 2001 From: justinpakzad <114518232+justinpakzad@users.noreply.github.com> Date: Mon, 13 Jul 2026 19:41:25 -0400 Subject: [PATCH] Add indexes on serialized_dag and dag_code --- airflow-core/docs/migrations-ref.rst | 4 +- ..._add_indexes_on_serialized_dag_and_dag_.py | 55 +++++++++++++++++++ ...4_0_add_gin_index_on_asset_event_extra.py} | 4 +- ... => 0126_3_4_0_drop_span_status_column.py} | 0 ...y => 0127_3_4_0_add_index_on_asset_uri.py} | 0 ..._4_0_add_idx_asset_event_partition_key.py} | 0 airflow-core/src/airflow/models/dagcode.py | 3 +- .../src/airflow/models/serialized_dag.py | 3 +- airflow-core/src/airflow/utils/db.py | 1 + ..._indexes_on_serialized_dag_and_dag_code.py | 55 +++++++++++++++++++ 10 files changed, 120 insertions(+), 5 deletions(-) create mode 100644 airflow-core/src/airflow/migrations/versions/0124_3_3_1_add_indexes_on_serialized_dag_and_dag_.py rename airflow-core/src/airflow/migrations/versions/{0124_3_4_0_add_gin_index_on_asset_event_extra.py => 0125_3_4_0_add_gin_index_on_asset_event_extra.py} (96%) rename airflow-core/src/airflow/migrations/versions/{0125_3_4_0_drop_span_status_column.py => 0126_3_4_0_drop_span_status_column.py} (100%) rename airflow-core/src/airflow/migrations/versions/{0126_3_4_0_add_index_on_asset_uri.py => 0127_3_4_0_add_index_on_asset_uri.py} (100%) rename airflow-core/src/airflow/migrations/versions/{0127_3_4_0_add_idx_asset_event_partition_key.py => 0128_3_4_0_add_idx_asset_event_partition_key.py} (100%) create mode 100644 airflow-core/tests/unit/migrations/test_0124_add_indexes_on_serialized_dag_and_dag_code.py diff --git a/airflow-core/docs/migrations-ref.rst b/airflow-core/docs/migrations-ref.rst index dfc96be79019e..a165633af1723 100644 --- a/airflow-core/docs/migrations-ref.rst +++ b/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. | +-------------------------+------------------+-------------------+--------------------------------------------------------------+ diff --git a/airflow-core/src/airflow/migrations/versions/0124_3_3_1_add_indexes_on_serialized_dag_and_dag_.py b/airflow-core/src/airflow/migrations/versions/0124_3_3_1_add_indexes_on_serialized_dag_and_dag_.py new file mode 100644 index 0000000000000..992399e0732f9 --- /dev/null +++ b/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") diff --git a/airflow-core/src/airflow/migrations/versions/0124_3_4_0_add_gin_index_on_asset_event_extra.py b/airflow-core/src/airflow/migrations/versions/0125_3_4_0_add_gin_index_on_asset_event_extra.py similarity index 96% rename from airflow-core/src/airflow/migrations/versions/0124_3_4_0_add_gin_index_on_asset_event_extra.py rename to airflow-core/src/airflow/migrations/versions/0125_3_4_0_add_gin_index_on_asset_event_extra.py index bffa7eb0f0348..a6a515ec3a3ec 100644 --- a/airflow-core/src/airflow/migrations/versions/0124_3_4_0_add_gin_index_on_asset_event_extra.py +++ b/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" diff --git a/airflow-core/src/airflow/migrations/versions/0125_3_4_0_drop_span_status_column.py b/airflow-core/src/airflow/migrations/versions/0126_3_4_0_drop_span_status_column.py similarity index 100% rename from airflow-core/src/airflow/migrations/versions/0125_3_4_0_drop_span_status_column.py rename to airflow-core/src/airflow/migrations/versions/0126_3_4_0_drop_span_status_column.py diff --git a/airflow-core/src/airflow/migrations/versions/0126_3_4_0_add_index_on_asset_uri.py b/airflow-core/src/airflow/migrations/versions/0127_3_4_0_add_index_on_asset_uri.py similarity index 100% rename from airflow-core/src/airflow/migrations/versions/0126_3_4_0_add_index_on_asset_uri.py rename to airflow-core/src/airflow/migrations/versions/0127_3_4_0_add_index_on_asset_uri.py diff --git a/airflow-core/src/airflow/migrations/versions/0127_3_4_0_add_idx_asset_event_partition_key.py b/airflow-core/src/airflow/migrations/versions/0128_3_4_0_add_idx_asset_event_partition_key.py similarity index 100% rename from airflow-core/src/airflow/migrations/versions/0127_3_4_0_add_idx_asset_event_partition_key.py rename to airflow-core/src/airflow/migrations/versions/0128_3_4_0_add_idx_asset_event_partition_key.py diff --git a/airflow-core/src/airflow/models/dagcode.py b/airflow-core/src/airflow/models/dagcode.py index d591ddb25b7b3..3c02b6480e8a8 100644 --- a/airflow-core/src/airflow/models/dagcode.py +++ b/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 diff --git a/airflow-core/src/airflow/models/serialized_dag.py b/airflow-core/src/airflow/models/serialized_dag.py index 9dda74d42b5fc..bdf554b57df21 100644 --- a/airflow-core/src/airflow/models/serialized_dag.py +++ b/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 diff --git a/airflow-core/src/airflow/utils/db.py b/airflow-core/src/airflow/utils/db.py index 1076715ee7180..4f467e0e906d5 100644 --- a/airflow-core/src/airflow/utils/db.py +++ b/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", } diff --git a/airflow-core/tests/unit/migrations/test_0124_add_indexes_on_serialized_dag_and_dag_code.py b/airflow-core/tests/unit/migrations/test_0124_add_indexes_on_serialized_dag_and_dag_code.py new file mode 100644 index 0000000000000..f59e26d961c75 --- /dev/null +++ b/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]