You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@airflow.apache.org by "jason810496 (via GitHub)" <gi...@apache.org> on 2026/03/19 08:38:14 UTC
[PR] Improve 0101_3_2_0_ui_improvements_for_deadlines upgrade migration [airflow]
jason810496 opened a new pull request, #63920:
URL: https://github.com/apache/airflow/pull/63920
* closes: https://github.com/apache/airflow/issues/63549
## Why
> Migration 0101_3_2_0_ui_improvements_for_deadlines upgrade is slow on deployments with large deadline and serialized_dag tables. With 10M deadline rows and 100K serialized_dags, the migration took ~16 minutes.
As mentioned in the issue
## What
There're several critical paths to improve
1. We're looping through `serialized_dag.dag_id` (**column without index**) before, now we're iterating through `serialized_dag.id` as batch processing partition key instead
2. Add temporary index for `serialized_dag.dagrun_id` just for migration as we couldn't avoid iterating `serialized_dag.dagrun_id` column.
3. Avoid ser/deser compression for `dag_data`
4. Build dialect-specific filter to skip rows without deadline data at the SQL level
5. Fetch dagrun IDs once per DAG to avoid repeating the expensive dag_run/serialized_dag JOIN for every alert.
---
##### Was generative AI tooling used to co-author this PR?
- [x] Yes (please specify the tool below) Claude Code for discussing and drafting
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Improve 0101_3_2_0_ui_improvements_for_deadlines upgrade migration [airflow]
Posted by "ephraimbuddy (via GitHub)" <gi...@apache.org>.
ephraimbuddy commented on code in PR #63920:
URL: https://github.com/apache/airflow/pull/63920#discussion_r2964311915
##########
airflow-core/tests/unit/migrations/test_0101_ui_improvements_for_deadlines.py:
##########
@@ -0,0 +1,112 @@
+# 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 importlib.util
+from pathlib import Path
+
+from sqlalchemy.exc import OperationalError
+
+from tests_common.test_utils.paths import AIRFLOW_CORE_SOURCES_PATH
+
+MIGRATION_PATH = Path(
+ AIRFLOW_CORE_SOURCES_PATH,
+ "airflow/migrations/versions/0101_3_2_0_ui_improvements_for_deadlines.py",
+)
+
+
+def load_migration_module():
+ spec = importlib.util.spec_from_file_location(
+ "migration_0101_ui_improvements_for_deadlines", MIGRATION_PATH
+ )
+ module = importlib.util.module_from_spec(spec)
+ assert spec.loader is not None
+ spec.loader.exec_module(module)
+ return module
+
+
+def configure_migration(mocker, migration_module, *, dialect_name: str, indexes: list[dict]):
+ conn = mocker.MagicMock()
+ conn.dialect.name = dialect_name
+
+ inspector = mocker.MagicMock()
+ inspector.get_indexes.return_value = indexes
+ inspector.get_pk_constraint.return_value = {"constrained_columns": ["id"]}
+ inspector.get_unique_constraints.return_value = []
+ mocker.patch.object(migration_module.sa, "inspect", return_value=inspector)
+
+ fake_op = mocker.MagicMock()
+ fake_op.get_bind.return_value = conn
+ mocker.patch.object(migration_module, "op", fake_op)
+ return fake_op
Review Comment:
I don't think we need the test file
##########
airflow-core/src/airflow/migrations/versions/0101_3_2_0_ui_improvements_for_deadlines.py:
##########
@@ -74,6 +75,83 @@
DEFAULT_BATCH_SIZE = 1000
ENCODING_TYPE = "deadline_alert"
+
+def _has_matching_index(conn: Connection, table_name: str, columns: Iterable[str]) -> bool:
+ log.debug("Targeting index check", table=table_name, columns=list(columns))
+ target_columns = list(columns)
+ inspector = sa.inspect(conn)
Review Comment:
This won't work in offline mode. We should indicate that it won't work in offline mode so those using the offline mode are aware of it
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Improve 0101_3_2_0_ui_improvements_for_deadlines upgrade migration [airflow]
Posted by "vatsrahul1001 (via GitHub)" <gi...@apache.org>.
vatsrahul1001 commented on PR #63920:
URL: https://github.com/apache/airflow/pull/63920#issuecomment-4095547286
I have verified this, and there is a significant improvement in migration timing. For me, migration now getting completed in 5 mins
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Improve 0101_3_2_0_ui_improvements_for_deadlines upgrade migration [airflow]
Posted by "jason810496 (via GitHub)" <gi...@apache.org>.
jason810496 commented on code in PR #63920:
URL: https://github.com/apache/airflow/pull/63920#discussion_r2964391725
##########
airflow-core/tests/unit/migrations/test_0101_ui_improvements_for_deadlines.py:
##########
@@ -0,0 +1,112 @@
+# 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 importlib.util
+from pathlib import Path
+
+from sqlalchemy.exc import OperationalError
+
+from tests_common.test_utils.paths import AIRFLOW_CORE_SOURCES_PATH
+
+MIGRATION_PATH = Path(
+ AIRFLOW_CORE_SOURCES_PATH,
+ "airflow/migrations/versions/0101_3_2_0_ui_improvements_for_deadlines.py",
+)
+
+
+def load_migration_module():
+ spec = importlib.util.spec_from_file_location(
+ "migration_0101_ui_improvements_for_deadlines", MIGRATION_PATH
+ )
+ module = importlib.util.module_from_spec(spec)
+ assert spec.loader is not None
+ spec.loader.exec_module(module)
+ return module
+
+
+def configure_migration(mocker, migration_module, *, dialect_name: str, indexes: list[dict]):
+ conn = mocker.MagicMock()
+ conn.dialect.name = dialect_name
+
+ inspector = mocker.MagicMock()
+ inspector.get_indexes.return_value = indexes
+ inspector.get_pk_constraint.return_value = {"constrained_columns": ["id"]}
+ inspector.get_unique_constraints.return_value = []
+ mocker.patch.object(migration_module.sa, "inspect", return_value=inspector)
+
+ fake_op = mocker.MagicMock()
+ fake_op.get_bind.return_value = conn
+ mocker.patch.object(migration_module, "op", fake_op)
+ return fake_op
Review Comment:
Sure, I can remove it before merge.
##########
airflow-core/tests/unit/migrations/test_0101_ui_improvements_for_deadlines.py:
##########
@@ -0,0 +1,112 @@
+# 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 importlib.util
+from pathlib import Path
+
+from sqlalchemy.exc import OperationalError
+
+from tests_common.test_utils.paths import AIRFLOW_CORE_SOURCES_PATH
+
+MIGRATION_PATH = Path(
+ AIRFLOW_CORE_SOURCES_PATH,
+ "airflow/migrations/versions/0101_3_2_0_ui_improvements_for_deadlines.py",
+)
+
+
+def load_migration_module():
+ spec = importlib.util.spec_from_file_location(
+ "migration_0101_ui_improvements_for_deadlines", MIGRATION_PATH
+ )
+ module = importlib.util.module_from_spec(spec)
+ assert spec.loader is not None
+ spec.loader.exec_module(module)
+ return module
+
+
+def configure_migration(mocker, migration_module, *, dialect_name: str, indexes: list[dict]):
+ conn = mocker.MagicMock()
+ conn.dialect.name = dialect_name
+
+ inspector = mocker.MagicMock()
+ inspector.get_indexes.return_value = indexes
+ inspector.get_pk_constraint.return_value = {"constrained_columns": ["id"]}
+ inspector.get_unique_constraints.return_value = []
+ mocker.patch.object(migration_module.sa, "inspect", return_value=inspector)
+
+ fake_op = mocker.MagicMock()
+ fake_op.get_bind.return_value = conn
+ mocker.patch.object(migration_module, "op", fake_op)
+ return fake_op
Review Comment:
Sure, I will remove it before merge.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Improve 0101_3_2_0_ui_improvements_for_deadlines upgrade migration [airflow]
Posted by "jason810496 (via GitHub)" <gi...@apache.org>.
jason810496 commented on code in PR #63920:
URL: https://github.com/apache/airflow/pull/63920#discussion_r2964389607
##########
airflow-core/src/airflow/migrations/versions/0101_3_2_0_ui_improvements_for_deadlines.py:
##########
@@ -74,6 +75,83 @@
DEFAULT_BATCH_SIZE = 1000
ENCODING_TYPE = "deadline_alert"
+
+def _has_matching_index(conn: Connection, table_name: str, columns: Iterable[str]) -> bool:
+ log.debug("Targeting index check", table=table_name, columns=list(columns))
+ target_columns = list(columns)
+ inspector = sa.inspect(conn)
Review Comment:
I just checked that we already had warning and return for the offline mode at the beginning of `migrate_existing_deadline_alert_data_from_serialized_dag` upgrade path:
https://github.com/apache/airflow/blob/42ea52cf97804353fd60551df3f33235f2c5d572/airflow-core/src/airflow/migrations/versions/0101_3_2_0_ui_improvements_for_deadlines.py#L467-L477
So we will create the table but not migrate the data in offline mode.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Improve 0101_3_2_0_ui_improvements_for_deadlines upgrade migration [airflow]
Posted by "jason810496 (via GitHub)" <gi...@apache.org>.
jason810496 commented on code in PR #63920:
URL: https://github.com/apache/airflow/pull/63920#discussion_r2973402407
##########
airflow-core/src/airflow/migrations/versions/0101_3_2_0_ui_improvements_for_deadlines.py:
##########
@@ -519,57 +644,124 @@ def _migrate_dag_deadlines(dag_conn: Connection) -> Iterable[str]:
)
if not validate_written_data(
- dag_conn, deadline_alert_id, reference_data, interval_data, callback_data
+ dag_conn,
+ deadline_alert_id,
+ reference_data,
+ interval_data,
+ callback_data,
):
dags_with_errors[dag_id].append(f"Invalid DeadlineAlert data: {serialized_alert}")
continue
yield deadline_alert_id
- dag_conn.execute(
- sa.text("""
- UPDATE deadline
- SET deadline_alert_id = :alert_id
- WHERE dagrun_id IN (
+ # Fetch dagrun IDs once per DAG to avoid repeating the
+ # expensive dag_run/serialized_dag JOIN for every alert.
+ dagrun_ids = [
Review Comment:
Nice catch! I oversaw when resolving the merge conflict during [0b0f25d](https://github.com/apache/airflow/pull/63920/commits/0b0f25d41d91f3ded063263632acd90f3e91a0b8)
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Improve 0101_3_2_0_ui_improvements_for_deadlines upgrade migration [airflow]
Posted by "ephraimbuddy (via GitHub)" <gi...@apache.org>.
ephraimbuddy commented on code in PR #63920:
URL: https://github.com/apache/airflow/pull/63920#discussion_r2973296545
##########
airflow-core/src/airflow/migrations/versions/0101_3_2_0_ui_improvements_for_deadlines.py:
##########
@@ -519,57 +644,124 @@ def _migrate_dag_deadlines(dag_conn: Connection) -> Iterable[str]:
)
if not validate_written_data(
- dag_conn, deadline_alert_id, reference_data, interval_data, callback_data
+ dag_conn,
+ deadline_alert_id,
+ reference_data,
+ interval_data,
+ callback_data,
):
dags_with_errors[dag_id].append(f"Invalid DeadlineAlert data: {serialized_alert}")
continue
yield deadline_alert_id
- dag_conn.execute(
- sa.text("""
- UPDATE deadline
- SET deadline_alert_id = :alert_id
- WHERE dagrun_id IN (
+ # Fetch dagrun IDs once per DAG to avoid repeating the
+ # expensive dag_run/serialized_dag JOIN for every alert.
+ dagrun_ids = [
Review Comment:
The comment here "Fetch dagrun IDs once per DAG" is different from what's implemented below. `serialized_dag_id ` is constant for the whole DAG, so `dagrun_ids` is the same for every alert. You can move this query outside the `for serialized_alert ` loop (e.g. once at the start of `_migrate_dag_deadlines`) and reuse `dagrun_ids ` for all alerts in the same DAG. Would make things faster too
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Improve 0101_3_2_0_ui_improvements_for_deadlines upgrade migration [airflow]
Posted by "vatsrahul1001 (via GitHub)" <gi...@apache.org>.
vatsrahul1001 merged PR #63920:
URL: https://github.com/apache/airflow/pull/63920
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Improve 0101_3_2_0_ui_improvements_for_deadlines upgrade migration [airflow]
Posted by "jason810496 (via GitHub)" <gi...@apache.org>.
jason810496 commented on code in PR #63920:
URL: https://github.com/apache/airflow/pull/63920#discussion_r2964389607
##########
airflow-core/src/airflow/migrations/versions/0101_3_2_0_ui_improvements_for_deadlines.py:
##########
@@ -74,6 +75,83 @@
DEFAULT_BATCH_SIZE = 1000
ENCODING_TYPE = "deadline_alert"
+
+def _has_matching_index(conn: Connection, table_name: str, columns: Iterable[str]) -> bool:
+ log.debug("Targeting index check", table=table_name, columns=list(columns))
+ target_columns = list(columns)
+ inspector = sa.inspect(conn)
Review Comment:
I just checked that we already had warning and return for the offline mode at the beginning of `migrate_existing_deadline_alert_data_from_serialized_dag` upgrade path:
https://github.com/apache/airflow/blob/42ea52cf97804353fd60551df3f33235f2c5d572/airflow-core/src/airflow/migrations/versions/0101_3_2_0_ui_improvements_for_deadlines.py#L467-L477
We will create the table but not migrate the data in offline mode, so I think we're good here, right?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Improve 0101_3_2_0_ui_improvements_for_deadlines upgrade migration [airflow]
Posted by "jason810496 (via GitHub)" <gi...@apache.org>.
jason810496 commented on code in PR #63920:
URL: https://github.com/apache/airflow/pull/63920#discussion_r2973396210
##########
airflow-core/tests/unit/migrations/test_0101_ui_improvements_for_deadlines.py:
##########
@@ -0,0 +1,112 @@
+# 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 importlib.util
+from pathlib import Path
+
+from sqlalchemy.exc import OperationalError
+
+from tests_common.test_utils.paths import AIRFLOW_CORE_SOURCES_PATH
+
+MIGRATION_PATH = Path(
+ AIRFLOW_CORE_SOURCES_PATH,
+ "airflow/migrations/versions/0101_3_2_0_ui_improvements_for_deadlines.py",
+)
+
+
+def load_migration_module():
+ spec = importlib.util.spec_from_file_location(
+ "migration_0101_ui_improvements_for_deadlines", MIGRATION_PATH
+ )
+ module = importlib.util.module_from_spec(spec)
+ assert spec.loader is not None
+ spec.loader.exec_module(module)
+ return module
+
+
+def configure_migration(mocker, migration_module, *, dialect_name: str, indexes: list[dict]):
+ conn = mocker.MagicMock()
+ conn.dialect.name = dialect_name
+
+ inspector = mocker.MagicMock()
+ inspector.get_indexes.return_value = indexes
+ inspector.get_pk_constraint.return_value = {"constrained_columns": ["id"]}
+ inspector.get_unique_constraints.return_value = []
+ mocker.patch.object(migration_module.sa, "inspect", return_value=inspector)
+
+ fake_op = mocker.MagicMock()
+ fake_op.get_bind.return_value = conn
+ mocker.patch.object(migration_module, "op", fake_op)
+ return fake_op
Review Comment:
Removed in [caa6100](https://github.com/apache/airflow/pull/63920/commits/caa61004b83a0a3c4a856a0df9fe0fd73e339450)
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Improve 0101_3_2_0_ui_improvements_for_deadlines upgrade migration [airflow]
Posted by "jason810496 (via GitHub)" <gi...@apache.org>.
jason810496 commented on code in PR #63920:
URL: https://github.com/apache/airflow/pull/63920#discussion_r2973402407
##########
airflow-core/src/airflow/migrations/versions/0101_3_2_0_ui_improvements_for_deadlines.py:
##########
@@ -519,57 +644,124 @@ def _migrate_dag_deadlines(dag_conn: Connection) -> Iterable[str]:
)
if not validate_written_data(
- dag_conn, deadline_alert_id, reference_data, interval_data, callback_data
+ dag_conn,
+ deadline_alert_id,
+ reference_data,
+ interval_data,
+ callback_data,
):
dags_with_errors[dag_id].append(f"Invalid DeadlineAlert data: {serialized_alert}")
continue
yield deadline_alert_id
- dag_conn.execute(
- sa.text("""
- UPDATE deadline
- SET deadline_alert_id = :alert_id
- WHERE dagrun_id IN (
+ # Fetch dagrun IDs once per DAG to avoid repeating the
+ # expensive dag_run/serialized_dag JOIN for every alert.
+ dagrun_ids = [
Review Comment:
Nice catch! I oversaw when resolving the merge conflict, addressed in [0b0f25d](https://github.com/apache/airflow/pull/63920/commits/0b0f25d41d91f3ded063263632acd90f3e91a0b8)
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org