From dbca09f27c9dfa1a74c3515422135418ab9c80f7 Mon Sep 17 00:00:00 2001 From: Ash Berlin-Taylor Date: Fri, 7 Oct 2022 17:21:53 +0100 Subject: [PATCH 1/2] Speed up downgrade DB tests There are two changes here: 1. Don't reserialize the dags. It's not part of the test so just adds extra time 2. Use `cached_app` in the migrations, rather than creating three copies of the app. To _ensure_ this has no side-effect outside of migrations we also purge the cached app at the end of running migrations, but only if it was already imported. --- .github/workflows/ci.yml | 4 ++-- airflow/cli/cli_parser.py | 9 +++++++++ airflow/cli/commands/db_command.py | 7 ++++++- airflow/migrations/env.py | 7 +++++++ .../versions/0074_2_0_0_resource_based_permissions.py | 6 +++--- .../0078_2_0_1_remove_can_read_permission_on_config_.py | 6 +++--- ...0084_2_1_0_resource_based_permissions_for_default_.py | 6 +++--- airflow/utils/db.py | 7 ++++--- docs/apache-airflow/img/airflow_erd.sha256 | 2 +- tests/cli/commands/test_db_command.py | 2 +- 10 files changed, 39 insertions(+), 17 deletions(-) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 81bdaf81bb108..0c5b34e151704 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -1064,9 +1064,9 @@ ${{ hashFiles('.pre-commit-config.yaml') }}" - name: "Test downgrade migration file ${{ env.BACKEND }}" run: > breeze shell "airflow db reset --skip-init -y && - airflow db upgrade --to-revision heads && + airflow db upgrade --to-revision heads --no-reserialize-dags && airflow db downgrade -r e959f08ac86c -y && - airflow db upgrade" + airflow db upgrade --no-reserialize-dags" env: IMAGE_TAG: ${{ env.IMAGE_TAG_FOR_THE_BUILD }} - name: "Test downgrade ORM ${{ env.BACKEND }}" diff --git a/airflow/cli/cli_parser.py b/airflow/cli/cli_parser.py index bf08fe2d74b86..bed5fc22499be 100644 --- a/airflow/cli/cli_parser.py +++ b/airflow/cli/cli_parser.py @@ -547,6 +547,14 @@ def string_lower_type(val): type=int, default=60, ) +ARG_DB_RESERIALIZE_DAGS = Arg( + ("--no-reserialize-dags",), + # Not intended for user, so dont show in help + help=argparse.SUPPRESS, + action="store_false", + default=True, + dest="reserialize_dags", +) ARG_DB_VERSION__UPGRADE = Arg( ("-n", "--to-version"), help=( @@ -1445,6 +1453,7 @@ class GroupCommand(NamedTuple): ARG_DB_SQL_ONLY, ARG_DB_FROM_REVISION, ARG_DB_FROM_VERSION, + ARG_DB_RESERIALIZE_DAGS, ), ), ActionCommand( diff --git a/airflow/cli/commands/db_command.py b/airflow/cli/commands/db_command.py index 6d981a352063c..e756e035f87ec 100644 --- a/airflow/cli/commands/db_command.py +++ b/airflow/cli/commands/db_command.py @@ -81,7 +81,12 @@ def upgradedb(args): else: print("Generating sql for upgrade -- upgrade commands will *not* be submitted.") - db.upgradedb(to_revision=to_revision, from_revision=from_revision, show_sql_only=args.show_sql_only) + db.upgradedb( + to_revision=to_revision, + from_revision=from_revision, + show_sql_only=args.show_sql_only, + reserialize_dags=args.reserialize_dags, + ) if not args.show_sql_only: print("Upgrades done") diff --git a/airflow/migrations/env.py b/airflow/migrations/env.py index 9dcd29e9ca0b5..dfdc3c0eaacf9 100644 --- a/airflow/migrations/env.py +++ b/airflow/migrations/env.py @@ -18,6 +18,7 @@ from __future__ import annotations import contextlib +import sys from logging.config import fileConfig from alembic import context @@ -114,3 +115,9 @@ def run_migrations_online(): run_migrations_offline() else: run_migrations_online() + +if 'airflow.www.app' in sys.modules: + # Already imported, make sure we clear out any cached app + from airflow.www.app import purge_cached_app + + purge_cached_app() diff --git a/airflow/migrations/versions/0074_2_0_0_resource_based_permissions.py b/airflow/migrations/versions/0074_2_0_0_resource_based_permissions.py index 296d07747f0ea..a913f13b170bd 100644 --- a/airflow/migrations/versions/0074_2_0_0_resource_based_permissions.py +++ b/airflow/migrations/versions/0074_2_0_0_resource_based_permissions.py @@ -27,7 +27,7 @@ import logging from airflow.security import permissions -from airflow.www.app import create_app +from airflow.www.app import cached_app # revision identifiers, used by Alembic. revision = '2c6edca13270' @@ -288,7 +288,7 @@ def remap_permissions(): """Apply Map Airflow permissions.""" - appbuilder = create_app(config={'FAB_UPDATE_PERMS': False}).appbuilder + appbuilder = cached_app(config={'FAB_UPDATE_PERMS': False}).appbuilder for old, new in mapping.items(): (old_resource_name, old_action_name) = old old_permission = appbuilder.sm.get_permission(old_action_name, old_resource_name) @@ -313,7 +313,7 @@ def remap_permissions(): def undo_remap_permissions(): """Unapply Map Airflow permissions""" - appbuilder = create_app(config={'FAB_UPDATE_PERMS': False}).appbuilder + appbuilder = cached_app(config={'FAB_UPDATE_PERMS': False}).appbuilder for old, new in mapping.items(): (new_resource_name, new_action_name) = new[0] new_permission = appbuilder.sm.get_permission(new_action_name, new_resource_name) diff --git a/airflow/migrations/versions/0078_2_0_1_remove_can_read_permission_on_config_.py b/airflow/migrations/versions/0078_2_0_1_remove_can_read_permission_on_config_.py index d6dd2b18f2a0a..23220b89fda08 100644 --- a/airflow/migrations/versions/0078_2_0_1_remove_can_read_permission_on_config_.py +++ b/airflow/migrations/versions/0078_2_0_1_remove_can_read_permission_on_config_.py @@ -27,7 +27,7 @@ import logging from airflow.security import permissions -from airflow.www.app import create_app +from airflow.www.app import cached_app # revision identifiers, used by Alembic. revision = '82b7c48c147f' @@ -42,7 +42,7 @@ def upgrade(): log = logging.getLogger() handlers = log.handlers[:] - appbuilder = create_app(config={'FAB_UPDATE_PERMS': False}).appbuilder + appbuilder = cached_app(config={'FAB_UPDATE_PERMS': False}).appbuilder roles_to_modify = [role for role in appbuilder.sm.get_all_roles() if role.name in ["User", "Viewer"]] can_read_on_config_perm = appbuilder.sm.get_permission( permissions.ACTION_CAN_READ, permissions.RESOURCE_CONFIG @@ -59,7 +59,7 @@ def upgrade(): def downgrade(): """Add can_read action on config resource for User and Viewer role""" - appbuilder = create_app(config={'FAB_UPDATE_PERMS': False}).appbuilder + appbuilder = cached_app(config={'FAB_UPDATE_PERMS': False}).appbuilder roles_to_modify = [role for role in appbuilder.sm.get_all_roles() if role.name in ["User", "Viewer"]] can_read_on_config_perm = appbuilder.sm.get_permission( permissions.ACTION_CAN_READ, permissions.RESOURCE_CONFIG diff --git a/airflow/migrations/versions/0084_2_1_0_resource_based_permissions_for_default_.py b/airflow/migrations/versions/0084_2_1_0_resource_based_permissions_for_default_.py index c7135f191b4a1..73a4f59b62367 100644 --- a/airflow/migrations/versions/0084_2_1_0_resource_based_permissions_for_default_.py +++ b/airflow/migrations/versions/0084_2_1_0_resource_based_permissions_for_default_.py @@ -27,7 +27,7 @@ import logging from airflow.security import permissions -from airflow.www.app import create_app +from airflow.www.app import cached_app # revision identifiers, used by Alembic. revision = 'a13f7613ad25' @@ -140,7 +140,7 @@ def remap_permissions(): """Apply Map Airflow permissions.""" - appbuilder = create_app(config={'FAB_UPDATE_PERMS': False}).appbuilder + appbuilder = cached_app(config={'FAB_UPDATE_PERMS': False}).appbuilder for old, new in mapping.items(): (old_resource_name, old_action_name) = old old_permission = appbuilder.sm.get_permission(old_action_name, old_resource_name) @@ -165,7 +165,7 @@ def remap_permissions(): def undo_remap_permissions(): """Unapply Map Airflow permissions""" - appbuilder = create_app(config={'FAB_UPDATE_PERMS': False}).appbuilder + appbuilder = cached_app(config={'FAB_UPDATE_PERMS': False}).appbuilder for old, new in mapping.items(): (new_resource_name, new_action_name) = new[0] new_permission = appbuilder.sm.get_permission(new_action_name, new_resource_name) diff --git a/airflow/utils/db.py b/airflow/utils/db.py index f796156f8f96d..06a166f1b0a91 100644 --- a/airflow/utils/db.py +++ b/airflow/utils/db.py @@ -825,8 +825,7 @@ def check_and_run_migrations(): sys.exit(1) -@provide_session -def reserialize_dags(*, session: Session = NEW_SESSION) -> None: +def _reserialize_dags(*, session: Session) -> None: from airflow.models.dagbag import DagBag from airflow.models.serialized_dag import SerializedDagModel @@ -1476,6 +1475,7 @@ def upgradedb( to_revision: str | None = None, from_revision: str | None = None, show_sql_only: bool = False, + reserialize_dags: bool = True, session: Session = NEW_SESSION, ): """ @@ -1556,7 +1556,8 @@ def upgradedb( os.environ['AIRFLOW__DATABASE__SQL_ALCHEMY_MAX_SIZE'] = val settings.reconfigure_orm() - reserialize_dags(session=session) + if reserialize_dags: + _reserialize_dags(session=session) add_default_pool_if_not_exists(session=session) synchronize_log_template(session=session) diff --git a/docs/apache-airflow/img/airflow_erd.sha256 b/docs/apache-airflow/img/airflow_erd.sha256 index bf9c204a6f81f..54654ec921b56 100644 --- a/docs/apache-airflow/img/airflow_erd.sha256 +++ b/docs/apache-airflow/img/airflow_erd.sha256 @@ -1 +1 @@ -dca9bf08dd97b5f51c387726a2e9d25996769e85dda0d403e8e088ce222faa09 \ No newline at end of file +3f6af8d3dbaf44cbef18f738e568b9a66375cc1214d6da9e482ebe5c591d24dc \ No newline at end of file diff --git a/tests/cli/commands/test_db_command.py b/tests/cli/commands/test_db_command.py index aeff2481591ca..c97886202f5bf 100644 --- a/tests/cli/commands/test_db_command.py +++ b/tests/cli/commands/test_db_command.py @@ -96,7 +96,7 @@ def test_cli_check_migrations(self, mock_wait_for_migrations): @mock.patch("airflow.cli.commands.db_command.db.upgradedb") def test_cli_upgrade_success(self, mock_upgradedb, args, called_with): db_command.upgradedb(self.parser.parse_args(['db', 'upgrade', *args])) - mock_upgradedb.assert_called_once_with(**called_with) + mock_upgradedb.assert_called_once_with(**called_with, reserialize_dags=True) @pytest.mark.parametrize( 'args, pattern', From 3f3a604032122678c36e58184e0b02a371f1d11f Mon Sep 17 00:00:00 2001 From: Ash Berlin-Taylor Date: Thu, 13 Oct 2022 19:15:45 +0100 Subject: [PATCH 2/2] Skip reserialize on more up/downgrade tests --- .github/workflows/ci.yml | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 0c5b34e151704..8a52b139b4769 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -1072,9 +1072,9 @@ ${{ hashFiles('.pre-commit-config.yaml') }}" - name: "Test downgrade ORM ${{ env.BACKEND }}" run: > breeze shell "airflow db reset -y && - airflow db upgrade && + airflow db upgrade --no-reserialize-dags && airflow db downgrade -r e959f08ac86c -y && - airflow db upgrade" + airflow db upgrade --no-reserialize-dags" env: IMAGE_TAG: ${{ env.IMAGE_TAG_FOR_THE_BUILD }} - name: "Test Offline SQL generation ${{ env.BACKEND }}" @@ -1156,17 +1156,17 @@ ${{ hashFiles('.pre-commit-config.yaml') }}" - name: "Test downgrade migration file ${{ env.BACKEND }}" run: > breeze shell "airflow db reset --skip-init -y && - airflow db upgrade --to-revision heads && + airflow db upgrade --to-revision heads --no-reserialize-dags && airflow db downgrade -r e959f08ac86c -y && - airflow db upgrade" + airflow db upgrade --no-reserialize-dags" env: IMAGE_TAG: ${{ env.IMAGE_TAG_FOR_THE_BUILD }} - name: "Test downgrade ORM ${{ env.BACKEND }}" run: > breeze shell "airflow db reset -y && - airflow db upgrade && + airflow db upgrade --no-reserialize-dags && airflow db downgrade -r e959f08ac86c -y && - airflow db upgrade" + airflow db upgrade --no-reserialize-dags" env: IMAGE_TAG: ${{ env.IMAGE_TAG_FOR_THE_BUILD }} - name: "Tests: ${{needs.build-info.outputs.test-types}}"