diff --git a/airflow/cli/commands/db_command.py b/airflow/cli/commands/db_command.py index 776dabcdead7e..cd7493212b4ee 100644 --- a/airflow/cli/commands/db_command.py +++ b/airflow/cli/commands/db_command.py @@ -72,15 +72,19 @@ def upgradedb(args): migratedb(args) -def get_version_revision(version: str, recursion_limit=10) -> str | None: +def _get_version_revision( + version: str, recursion_limit: int = 10, revision_heads_map: dict[str, str] | None = None +) -> str | None: """ - Recursively search for the revision of the given version. + Recursively search for the revision of the given version in revision_heads_map. - This searches REVISION_HEADS_MAP for the revision of the given version, recursively + This searches given revision_heads_map for the revision of the given version, recursively searching for the previous version if the given version is not found. """ - if version in _REVISION_HEADS_MAP: - return _REVISION_HEADS_MAP[version] + if revision_heads_map is None: + revision_heads_map = _REVISION_HEADS_MAP + if version in revision_heads_map: + return revision_heads_map[version] try: major, minor, patch = map(int, version.split(".")) except ValueError: @@ -90,13 +94,19 @@ def get_version_revision(version: str, recursion_limit=10) -> str | None: if recursion_limit <= 0: # Prevent infinite recursion as I can't imagine 10 successive versions without migration return None - return get_version_revision(new_version, recursion_limit) + return _get_version_revision(new_version, recursion_limit) -@cli_utils.action_cli(check_db=False) -@providers_configuration_loaded -def migratedb(args): - """Migrates the metadata database.""" +def run_db_migrate_command(args, command, revision_heads_map: dict[str, str], reserialize_dags: bool = True): + """ + Run the db migrate command. + + param args: The parsed arguments. + param command: The command to run. + param airflow_db: Whether the command is for the airflow database. + + :meta private: + """ print(f"DB: {settings.engine.url!r}") if args.to_revision and args.to_version: raise SystemExit("Cannot supply both `--to-revision` and `--to-version`.") @@ -112,12 +122,10 @@ def migratedb(args): from_revision = args.from_revision elif args.from_version: try: - parsed_version = parse_version(args.from_version) + parse_version(args.from_version) except InvalidVersion: raise SystemExit(f"Invalid version {args.from_version!r} supplied as `--from-version`.") - if parsed_version < parse_version("2.0.0"): - raise SystemExit("--from-version must be greater or equal to than 2.0.0") - from_revision = get_version_revision(args.from_version) + from_revision = _get_version_revision(args.from_version, revision_heads_map=revision_heads_map) if not from_revision: raise SystemExit(f"Unknown version {args.from_version!r} supplied as `--from-version`.") @@ -126,7 +134,7 @@ def migratedb(args): parse_version(args.to_version) except InvalidVersion: raise SystemExit(f"Invalid version {args.to_version!r} supplied as `--to-version`.") - to_revision = get_version_revision(args.to_version) + to_revision = _get_version_revision(args.to_version, revision_heads_map=revision_heads_map) if not to_revision: raise SystemExit(f"Unknown version {args.to_version!r} supplied as `--to-version`.") elif args.to_revision: @@ -136,21 +144,30 @@ def migratedb(args): print(f"Performing upgrade to the metadata database {settings.engine.url!r}") 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, - reserialize_dags=args.reserialize_dags, - ) + if reserialize_dags: + command( + to_revision=to_revision, + from_revision=from_revision, + show_sql_only=args.show_sql_only, + reserialize_dags=True, + ) + else: + command( + to_revision=to_revision, + from_revision=from_revision, + show_sql_only=args.show_sql_only, + ) if not args.show_sql_only: print("Database migrating done!") -@cli_utils.action_cli(check_db=False) -@providers_configuration_loaded -def downgrade(args): - """Downgrades the metadata database.""" +def run_db_downgrade_command(args, command, revision_heads_map: dict[str, str]): + """ + Run the db downgrade command. + + param args: The parsed arguments. + param command: The command to run. + """ if args.to_revision and args.to_version: raise SystemExit("Cannot supply both `--to-revision` and `--to-version`.") if args.from_version and args.from_revision: @@ -162,14 +179,15 @@ def downgrade(args): if not (args.to_version or args.to_revision): raise SystemExit("Must provide either --to-revision or --to-version.") from_revision = None + to_revision = None if args.from_revision: from_revision = args.from_revision elif args.from_version: - from_revision = get_version_revision(args.from_version) + from_revision = _get_version_revision(args.from_version, revision_heads_map=revision_heads_map) if not from_revision: raise SystemExit(f"Unknown version {args.from_version!r} supplied as `--from-version`.") if args.to_version: - to_revision = get_version_revision(args.to_version) + to_revision = _get_version_revision(args.to_version, revision_heads_map=revision_heads_map) if not to_revision: raise SystemExit(f"Downgrading to version {args.to_version} is not supported.") elif args.to_revision: @@ -188,13 +206,34 @@ def downgrade(args): ).upper() == "Y" ): - db.downgrade(to_revision=to_revision, from_revision=from_revision, show_sql_only=args.show_sql_only) + command(to_revision=to_revision, from_revision=from_revision, show_sql_only=args.show_sql_only) if not args.show_sql_only: print("Downgrade complete") else: raise SystemExit("Cancelled") +@cli_utils.action_cli(check_db=False) +@providers_configuration_loaded +def migratedb(args): + """Migrates the metadata database.""" + if args.from_version: + try: + parsed_version = parse_version(args.from_version) + except InvalidVersion: + raise SystemExit(f"Invalid version {args.from_version!r} supplied as `--from-version`.") + if parsed_version < parse_version("2.0.0"): + raise SystemExit("--from-version must be greater or equal to 2.0.0") + run_db_migrate_command(args, db.upgradedb, _REVISION_HEADS_MAP, reserialize_dags=True) + + +@cli_utils.action_cli(check_db=False) +@providers_configuration_loaded +def downgrade(args): + """Downgrades the metadata database.""" + run_db_downgrade_command(args, db.downgrade, _REVISION_HEADS_MAP) + + @providers_configuration_loaded def check_migrations(args): """Wait for all airflow migrations to complete. Used for launching airflow in k8s.""" diff --git a/airflow/providers/fab/auth_manager/cli_commands/db_command.py b/airflow/providers/fab/auth_manager/cli_commands/db_command.py new file mode 100644 index 0000000000000..8b41cf4216c87 --- /dev/null +++ b/airflow/providers/fab/auth_manager/cli_commands/db_command.py @@ -0,0 +1,52 @@ +# 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 + +from airflow import settings +from airflow.cli.commands.db_command import run_db_downgrade_command, run_db_migrate_command +from airflow.providers.fab.auth_manager.models.db import _REVISION_HEADS_MAP, FABDBManager +from airflow.utils import cli as cli_utils +from airflow.utils.providers_configuration_loader import providers_configuration_loaded + + +@providers_configuration_loaded +def resetdb(args): + """Reset the metadata database.""" + print(f"DB: {settings.engine.url!r}") + if not (args.yes or input("This will drop existing tables if they exist. Proceed? (y/n)").upper() == "Y"): + raise SystemExit("Cancelled") + FABDBManager(settings.Session()).resetdb(skip_init=args.skip_init) + + +@cli_utils.action_cli(check_db=False) +@providers_configuration_loaded +def migratedb(args): + """Migrates the metadata database.""" + session = settings.Session() + upgrade_command = FABDBManager(session).upgradedb + run_db_migrate_command( + args, upgrade_command, revision_heads_map=_REVISION_HEADS_MAP, reserialize_dags=False + ) + + +@cli_utils.action_cli(check_db=False) +@providers_configuration_loaded +def downgrade(args): + """Downgrades the metadata database.""" + session = settings.Session() + dwongrade_command = FABDBManager(session).downgrade + run_db_downgrade_command(args, dwongrade_command, revision_heads_map=_REVISION_HEADS_MAP) diff --git a/airflow/providers/fab/auth_manager/cli_commands/definition.py b/airflow/providers/fab/auth_manager/cli_commands/definition.py index c7be5270d58fe..7f8f1e84e2798 100644 --- a/airflow/providers/fab/auth_manager/cli_commands/definition.py +++ b/airflow/providers/fab/auth_manager/cli_commands/definition.py @@ -19,8 +19,17 @@ import textwrap from airflow.cli.cli_config import ( + ARG_DB_FROM_REVISION, + ARG_DB_FROM_VERSION, + ARG_DB_REVISION__DOWNGRADE, + ARG_DB_REVISION__UPGRADE, + ARG_DB_SKIP_INIT, + ARG_DB_SQL_ONLY, + ARG_DB_VERSION__DOWNGRADE, + ARG_DB_VERSION__UPGRADE, ARG_OUTPUT, ARG_VERBOSE, + ARG_YES, ActionCommand, Arg, lazy_load_command, @@ -243,3 +252,55 @@ func=lazy_load_command("airflow.providers.fab.auth_manager.cli_commands.sync_perm_command.sync_perm"), args=(ARG_INCLUDE_DAGS, ARG_VERBOSE), ) + +DB_COMMANDS = ( + ActionCommand( + name="migrate", + help="Migrates the FAB metadata database to the latest version", + description=( + "Migrate the schema of the FAB metadata database. " + "Create the database if it does not exist " + "To print but not execute commands, use option ``--show-sql-only``. " + "If using options ``--from-revision`` or ``--from-version``, you must also use " + "``--show-sql-only``, because if actually *running* migrations, we should only " + "migrate from the *current* Alembic revision." + ), + func=lazy_load_command("airflow.providers.fab.auth_manager.cli_commands.db_command.migratedb"), + args=( + ARG_DB_REVISION__UPGRADE, + ARG_DB_VERSION__UPGRADE, + ARG_DB_SQL_ONLY, + ARG_DB_FROM_REVISION, + ARG_DB_FROM_VERSION, + ARG_VERBOSE, + ), + ), + ActionCommand( + name="downgrade", + help="Downgrade the schema of the FAB metadata database.", + description=( + "Downgrade the schema of the FAB metadata database. " + "You must provide either `--to-revision` or `--to-version`. " + "To print but not execute commands, use option `--show-sql-only`. " + "If using options `--from-revision` or `--from-version`, you must also use `--show-sql-only`, " + "because if actually *running* migrations, we should only migrate from the *current* Alembic " + "revision." + ), + func=lazy_load_command("airflow.providers.fab.auth_manager.cli_commands.db_command.downgrade"), + args=( + ARG_DB_REVISION__DOWNGRADE, + ARG_DB_VERSION__DOWNGRADE, + ARG_DB_SQL_ONLY, + ARG_YES, + ARG_DB_FROM_REVISION, + ARG_DB_FROM_VERSION, + ARG_VERBOSE, + ), + ), + ActionCommand( + name="reset", + help="Burn down and rebuild the FAB metadata database", + func=lazy_load_command("airflow.providers.fab.auth_manager.cli_commands.db_command.resetdb"), + args=(ARG_YES, ARG_DB_SKIP_INIT, ARG_VERBOSE), + ), +) diff --git a/airflow/providers/fab/auth_manager/fab_auth_manager.py b/airflow/providers/fab/auth_manager/fab_auth_manager.py index a0ec27bc67b00..336437061c2f7 100644 --- a/airflow/providers/fab/auth_manager/fab_auth_manager.py +++ b/airflow/providers/fab/auth_manager/fab_auth_manager.py @@ -22,11 +22,13 @@ from pathlib import Path from typing import TYPE_CHECKING, Container +import packaging.version from connexion import FlaskApi from flask import Blueprint, url_for from sqlalchemy import select from sqlalchemy.orm import Session, joinedload +from airflow import __version__ as airflow_version from airflow.auth.managers.base_auth_manager import BaseAuthManager, ResourceMethod from airflow.auth.managers.models.resource_details import ( AccessView, @@ -47,6 +49,7 @@ from airflow.exceptions import AirflowConfigException, AirflowException from airflow.models import DagModel from airflow.providers.fab.auth_manager.cli_commands.definition import ( + DB_COMMANDS, ROLES_COMMANDS, SYNC_PERM_COMMAND, USERS_COMMANDS, @@ -132,7 +135,7 @@ class FabAuthManager(BaseAuthManager): @staticmethod def get_cli_commands() -> list[CLICommand]: """Vends CLI commands to be included in Airflow CLI.""" - return [ + commands: list[CLICommand] = [ GroupCommand( name="users", help="Manage users", @@ -145,6 +148,12 @@ def get_cli_commands() -> list[CLICommand]: ), SYNC_PERM_COMMAND, # not in a command group ] + # If Airflow version is 3.0.0 or higher, add the fab-db command group + if packaging.version.parse( + packaging.version.parse(airflow_version).base_version + ) >= packaging.version.parse("3.0.0"): + commands.append(GroupCommand(name="fab-db", help="Manage FAB", subcommands=DB_COMMANDS)) + return commands def get_api_endpoints(self) -> None | Blueprint: folder = Path(__file__).parents[0].resolve() # this is airflow/auth/managers/fab/ diff --git a/airflow/providers/fab/auth_manager/models/db.py b/airflow/providers/fab/auth_manager/models/db.py index a971ea29a3f68..f72e1fcc65b8f 100644 --- a/airflow/providers/fab/auth_manager/models/db.py +++ b/airflow/providers/fab/auth_manager/models/db.py @@ -19,17 +19,89 @@ import os import airflow +from airflow import settings +from airflow.exceptions import AirflowException from airflow.providers.fab.auth_manager.models import metadata +from airflow.utils.db import _offline_migration, print_happy_cat from airflow.utils.db_manager import BaseDBManager PACKAGE_DIR = os.path.dirname(airflow.__file__) +_REVISION_HEADS_MAP: dict[str, str] = { + "1.3.0": "6709f7a774b9", +} + class FABDBManager(BaseDBManager): """Manages FAB database.""" metadata = metadata - version_table_name = "fab_alembic_version" + version_table_name = "alembic_version_fab" migration_dir = os.path.join(PACKAGE_DIR, "providers/fab/migrations") alembic_file = os.path.join(PACKAGE_DIR, "providers/fab/alembic.ini") supports_table_dropping = True + + def upgradedb(self, to_revision=None, from_revision=None, show_sql_only=False): + """Upgrade the database.""" + if from_revision and not show_sql_only: + raise AirflowException("`from_revision` only supported with `sql_only=True`.") + + # alembic adds significant import time, so we import it lazily + if not settings.SQL_ALCHEMY_CONN: + raise RuntimeError("The settings.SQL_ALCHEMY_CONN not set. This is a critical assertion.") + from alembic import command + + config = self.get_alembic_config() + + if show_sql_only: + if settings.engine.dialect.name == "sqlite": + raise SystemExit("Offline migration not supported for SQLite.") + if not from_revision: + from_revision = self.get_current_revision() + + if not to_revision: + script = self.get_script_object(config) + to_revision = script.get_current_head() + + if to_revision == from_revision: + print_happy_cat("No migrations to apply; nothing to do.") + return + _offline_migration(command.upgrade, config, f"{from_revision}:{to_revision}") + return # only running sql; our job is done + + if not self.get_current_revision(): + # New DB; initialize and exit + self.initdb() + return + + command.upgrade(config, revision=to_revision or "heads") + + def downgrade(self, to_revision, from_revision=None, show_sql_only=False): + if from_revision and not show_sql_only: + raise ValueError( + "`from_revision` can't be combined with `show_sql_only=False`. When actually " + "applying a downgrade (instead of just generating sql), we always " + "downgrade from current revision." + ) + + if not settings.SQL_ALCHEMY_CONN: + raise RuntimeError("The settings.SQL_ALCHEMY_CONN not set.") + + # alembic adds significant import time, so we import it lazily + from alembic import command + + self.log.info("Attempting downgrade of FAB migration to revision %s", to_revision) + config = self.get_alembic_config() + + if show_sql_only: + self.log.warning("Generating sql scripts for manual migration.") + if not from_revision: + from_revision = self.get_current_revision() + if from_revision is None: + self.log.info("No revision found") + return + revision_range = f"{from_revision}:{to_revision}" + _offline_migration(command.downgrade, config=config, revision=revision_range) + else: + self.log.info("Applying FAB downgrade migrations.") + command.downgrade(config, revision=to_revision, sql=show_sql_only) diff --git a/airflow/providers/fab/migrations/script.py.mako b/airflow/providers/fab/migrations/script.py.mako index 4d0928fcc09ad..c0193ce2b0471 100644 --- a/airflow/providers/fab/migrations/script.py.mako +++ b/airflow/providers/fab/migrations/script.py.mako @@ -30,10 +30,11 @@ import sqlalchemy as sa ${imports if imports else ""} # revision identifiers, used by Alembic. -revision: str = ${repr(up_revision)} -down_revision: Union[str, None] = ${repr(down_revision)} -branch_labels: Union[str, Sequence[str], None] = ${repr(branch_labels)} -depends_on: Union[str, Sequence[str], None] = ${repr(depends_on)} +revision = ${repr(up_revision)} +down_revision = ${repr(down_revision)} +branch_labels = ${repr(branch_labels)} +depends_on = ${repr(depends_on)} +fab_version = None def upgrade() -> None: diff --git a/airflow/providers/fab/migrations/versions/0001_1_3_0_placeholder_migration.py b/airflow/providers/fab/migrations/versions/0001_1_3_0_placeholder_migration.py new file mode 100644 index 0000000000000..685216779dcfd --- /dev/null +++ b/airflow/providers/fab/migrations/versions/0001_1_3_0_placeholder_migration.py @@ -0,0 +1,45 @@ +# +# 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. + +""" +placeholder migration. + +Revision ID: 6709f7a774b9 +Revises: +Create Date: 2024-09-03 17:06:38.040510 + +Note: This is a placeholder migration used to stamp the migration +when we create the migration from the ORM. Otherwise, it will run +without stamping the migration, leading to subsequent changes to +the tables not being migrated. +""" + +from __future__ import annotations + +# revision identifiers, used by Alembic. +revision = "6709f7a774b9" +down_revision = None +branch_labels = None +depends_on = None +fab_version = "1.3.0" + + +def upgrade() -> None: ... + + +def downgrade() -> None: ... diff --git a/airflow/utils/db.py b/airflow/utils/db.py index bc4c60697cead..1dc7275610dfc 100644 --- a/airflow/utils/db.py +++ b/airflow/utils/db.py @@ -89,7 +89,7 @@ class MappedClassProtocol(Protocol): log = logging.getLogger(__name__) -_REVISION_HEADS_MAP = { +_REVISION_HEADS_MAP: dict[str, str] = { "2.7.0": "405de8318b3a", "2.8.0": "10b52ebd31f7", "2.8.1": "88344c1d9134", @@ -819,6 +819,7 @@ def check_migrations(timeout): :return: None """ timeout = timeout or 1 # run the loop at least 1 + external_db_manager = RunDBManager() with _configured_alembic_environment() as env: context = env.get_context() source_heads = None @@ -826,7 +827,7 @@ def check_migrations(timeout): for ticker in range(timeout): source_heads = set(env.script.get_heads()) db_heads = set(context.get_current_heads()) - if source_heads == db_heads: + if source_heads == db_heads and external_db_manager.check_migration(settings.Session()): return time.sleep(1) log.info("Waiting for migrations... %s second(s)", ticker) @@ -1027,6 +1028,11 @@ def _check_migration_errors(session: Session = NEW_SESSION) -> Iterable[str]: def _offline_migration(migration_func: Callable, config, revision): + """ + Run offline migration. + + :meta private: + """ with warnings.catch_warnings(): warnings.simplefilter("ignore") logging.disable(logging.CRITICAL) @@ -1160,6 +1166,13 @@ def upgradedb( os.environ["AIRFLOW__DATABASE__SQL_ALCHEMY_MAX_SIZE"] = "1" settings.reconfigure_orm(pool_class=sqlalchemy.pool.SingletonThreadPool) command.upgrade(config, revision=to_revision or "heads") + current_revision = _get_current_revision(session=session) + with _configured_alembic_environment() as env: + source_heads = env.script.get_heads() + if current_revision == source_heads[0]: + # Only run external DB upgrade migration if user upgraded to heads + external_db_manager = RunDBManager() + external_db_manager.upgradedb(session) finally: if val is None: @@ -1168,8 +1181,6 @@ def upgradedb( os.environ["AIRFLOW__DATABASE__SQL_ALCHEMY_MAX_SIZE"] = val settings.reconfigure_orm() - current_revision = _get_current_revision(session=session) - if reserialize_dags and current_revision != previous_revision: _reserialize_dags(session=session) add_default_pool_if_not_exists(session=session) @@ -1191,7 +1202,7 @@ def resetdb(session: Session = NEW_SESSION, skip_init: bool = False): drop_airflow_models(connection) drop_airflow_moved_tables(connection) external_db_manager = RunDBManager() - external_db_manager.drop_tables(connection) + external_db_manager.drop_tables(session, connection) if not skip_init: initdb(session=session) diff --git a/airflow/utils/db_manager.py b/airflow/utils/db_manager.py index 78d3244cd21fe..2241a646efd48 100644 --- a/airflow/utils/db_manager.py +++ b/airflow/utils/db_manager.py @@ -20,13 +20,16 @@ from typing import TYPE_CHECKING from alembic import command +from sqlalchemy import inspect +from airflow import settings from airflow.configuration import conf from airflow.exceptions import AirflowException from airflow.utils.log.logging_mixin import LoggingMixin from airflow.utils.module_loading import import_string if TYPE_CHECKING: + from alembic.script import ScriptDirectory from sqlalchemy import MetaData @@ -54,14 +57,34 @@ def get_alembic_config(self): config.set_main_option("sqlalchemy.url", settings.SQL_ALCHEMY_CONN.replace("%", "%%")) return config - def get_current_revision(self): + def get_script_object(self, config=None) -> ScriptDirectory: + from alembic.script import ScriptDirectory + + if not config: + config = self.get_alembic_config() + return ScriptDirectory.from_config(config) + + def _get_migration_ctx(self): from alembic.migration import MigrationContext conn = self.session.connection() - migration_ctx = MigrationContext.configure(conn, opts={"version_table": self.version_table_name}) + return MigrationContext.configure(conn, opts={"version_table": self.version_table_name}) - return migration_ctx.get_current_revision() + def get_current_revision(self): + return self._get_migration_ctx().get_current_revision() + + def check_migration(self): + """Check migration done.""" + script_heads = self.get_script_object().get_heads() + db_heads = self.get_current_revision() + if db_heads: + db_heads = {db_heads} + if not db_heads and not script_heads: + return True + if set(script_heads) == db_heads: + return True + return False def _create_db_from_orm(self): """Create database from ORM.""" @@ -70,6 +93,22 @@ def _create_db_from_orm(self): config = self.get_alembic_config() command.stamp(config, "head") + def drop_tables(self, connection): + self.metadata.drop_all(connection) + version = self._get_migration_ctx()._version + if inspect(connection).has_table(version.name): + version.drop(connection) + + def resetdb(self, skip_init=False): + from airflow.utils.db import DBLocks, create_global_lock + + connection = settings.engine.connect() + + with create_global_lock(self.session, lock=DBLocks.MIGRATIONS), connection.begin(): + self.drop_tables(connection) + if not skip_init: + self.initdb() + def initdb(self): """Initialize the database.""" db_exists = self.get_current_revision() @@ -78,14 +117,14 @@ def initdb(self): else: self._create_db_from_orm() - def upgradedb(self, to_version=None, from_version=None, show_sql_only=False): + def upgradedb(self, to_revision=None, from_revision=None, show_sql_only=False): """Upgrade the database.""" self.log.info("Upgrading the %s database", self.__class__.__name__) config = self.get_alembic_config() - command.upgrade(config, revision=to_version or "heads", sql=show_sql_only) + command.upgrade(config, revision=to_revision or "heads", sql=show_sql_only) - def downgradedb(self, to_version, from_version=None, show_sql_only=False): + def downgrade(self, to_version, from_version=None, show_sql_only=False): """Downgrade the database.""" raise NotImplementedError @@ -141,6 +180,14 @@ def _validate(manager: BaseDBManager): if manager.version_table_name == "alembic_version": raise AirflowException(f"{manager}.version_table_name cannot be 'alembic_version'") + def check_migration(self, session): + """Check the external database migration.""" + return_value = [] + for manager in self._managers: + m = manager(session) + return_value.append(m.check_migration) + return all([x() for x in return_value]) + def initdb(self, session): """Initialize the external database managers.""" for manager in self._managers: @@ -153,14 +200,15 @@ def upgradedb(self, session): m = manager(session) m.upgradedb() - def downgradedb(self, session): + def downgrade(self, session): """Downgrade the external database managers.""" for manager in self._managers: m = manager(session) - m.downgradedb() + m.downgrade() - def drop_tables(self, connection): + def drop_tables(self, session, connection): """Drop the external database managers.""" for manager in self._managers: if manager.supports_table_dropping: - manager.metadata.drop_all(connection) + m = manager(session) + m.drop_tables(connection) diff --git a/scripts/ci/pre_commit/version_heads_map.py b/scripts/ci/pre_commit/version_heads_map.py index 4277c4656472d..10a6dee2eaf23 100755 --- a/scripts/ci/pre_commit/version_heads_map.py +++ b/scripts/ci/pre_commit/version_heads_map.py @@ -23,21 +23,27 @@ from pathlib import Path import re2 -from packaging.version import parse as parse_version PROJECT_SOURCE_ROOT_DIR = Path(__file__).resolve().parent.parent.parent.parent DB_FILE = PROJECT_SOURCE_ROOT_DIR / "airflow" / "utils" / "db.py" MIGRATION_PATH = PROJECT_SOURCE_ROOT_DIR / "airflow" / "migrations" / "versions" +FAB_DB_FILE = PROJECT_SOURCE_ROOT_DIR / "airflow" / "providers" / "fab" / "auth_manager" / "models" / "db.py" +FAB_MIGRATION_PATH = PROJECT_SOURCE_ROOT_DIR / "airflow" / "providers" / "fab" / "migrations" / "versions" + sys.path.insert(0, str(Path(__file__).parent.resolve())) # make sure common_precommit_utils is importable -def revision_heads_map(): +def revision_heads_map(migration_path): rh_map = {} pattern = r'revision = "[a-fA-F0-9]+"' - airflow_version_pattern = r'airflow_version = "\d+\.\d+\.\d+"' - filenames = os.listdir(MIGRATION_PATH) + version_pattern = None + if migration_path == MIGRATION_PATH: + version_pattern = r'airflow_version = "\d+\.\d+\.\d+"' + elif migration_path == FAB_MIGRATION_PATH: + version_pattern = r'fab_version = "\d+\.\d+\.\d+"' + filenames = os.listdir(migration_path) def sorting_key(filen): prefix = filen.split("_")[0] @@ -46,43 +52,46 @@ def sorting_key(filen): sorted_filenames = sorted(filenames, key=sorting_key) for filename in sorted_filenames: - if not filename.endswith(".py"): + if not filename.endswith(".py") or filename == "__init__.py": continue - with open(os.path.join(MIGRATION_PATH, filename)) as file: + with open(os.path.join(migration_path, filename)) as file: content = file.read() revision_match = re2.search(pattern, content) - airflow_version_match = re2.search(airflow_version_pattern, content) - if revision_match and airflow_version_match: + _version_match = re2.search(version_pattern, content) + if revision_match and _version_match: revision = revision_match.group(0).split('"')[1] - version = airflow_version_match.group(0).split('"')[1] - if parse_version(version) >= parse_version("2.0.0"): - rh_map[version] = revision + version = _version_match.group(0).split('"')[1] + rh_map[version] = revision return rh_map if __name__ == "__main__": - with open(DB_FILE) as file: - content = file.read() - - pattern = r"_REVISION_HEADS_MAP = {[^}]+\}" - match = re2.search(pattern, content) - if not match: - print( - f"_REVISION_HEADS_MAP not found in {DB_FILE}. If this has been removed intentionally, " - "please update scripts/ci/pre_commit/version_heads_map.py" - ) - sys.exit(1) - - existing_revision_heads_map = match.group(0) - rh_map = revision_heads_map() - updated_revision_heads_map = "_REVISION_HEADS_MAP = {\n" - for k, v in rh_map.items(): - updated_revision_heads_map += f' "{k}": "{v}",\n' - updated_revision_heads_map += "}" - if existing_revision_heads_map != updated_revision_heads_map: - new_content = content.replace(existing_revision_heads_map, updated_revision_heads_map) - - with open(DB_FILE, "w") as file: - file.write(new_content) - print("_REVISION_HEADS_MAP updated in db.py. Please commit the changes.") - sys.exit(1) + paths = [(DB_FILE, MIGRATION_PATH), (FAB_DB_FILE, FAB_MIGRATION_PATH)] + for dbfile, mpath in paths: + with open(dbfile) as file: + content = file.read() + + pattern = r"_REVISION_HEADS_MAP:\s*dict\[\s*str\s*,\s*str\s*\]\s*=\s*\{[^}]*\}" + match = re2.search(pattern, content) + if not match: + print( + f"_REVISION_HEADS_MAP not found in {dbfile}. If this has been removed intentionally, " + "please update scripts/ci/pre_commit/version_heads_map.py" + ) + sys.exit(1) + + existing_revision_heads_map = match.group(0) + rh_map = revision_heads_map(mpath) + updated_revision_heads_map = "_REVISION_HEADS_MAP: dict[str, str] = {\n" + for k, v in rh_map.items(): + updated_revision_heads_map += f' "{k}": "{v}",\n' + updated_revision_heads_map += "}" + if updated_revision_heads_map == "_REVISION_HEADS_MAP: dict[str, str] = {\n}": + updated_revision_heads_map = "_REVISION_HEADS_MAP: dict[str, str] = {}" + if existing_revision_heads_map != updated_revision_heads_map: + new_content = content.replace(existing_revision_heads_map, updated_revision_heads_map) + + with open(dbfile, "w") as file: + file.write(new_content) + print(f"_REVISION_HEADS_MAP updated in {dbfile}. Please commit the changes.") + sys.exit(1) diff --git a/scripts/in_container/run_migration_reference.py b/scripts/in_container/run_migration_reference.py index 204db7a5ea908..fc7d3bd084ca9 100755 --- a/scripts/in_container/run_migration_reference.py +++ b/scripts/in_container/run_migration_reference.py @@ -138,8 +138,7 @@ def get_revisions(app="airflow") -> Iterable[Script]: else: from airflow.providers.fab.auth_manager.models.db import FABDBManager - config = FABDBManager(session="").get_alembic_config() - script = ScriptDirectory.from_config(config) + script = FABDBManager(session="").get_script_object() yield from script.walk_revisions() @@ -234,9 +233,9 @@ def correct_mismatching_revision_nums(revisions: Iterable[Script]): if __name__ == "__main__": apps = ["airflow", "fab"] for app in apps: - console.print("[bright_blue]Updating migration reference") + console.print(f"[bright_blue]Updating migration reference for {app}") revisions = list(reversed(list(get_revisions(app)))) - console.print("[bright_blue]Making sure airflow version updated") + console.print(f"[bright_blue]Making sure {app} version updated") ensure_version(revisions=revisions, app=app) console.print("[bright_blue]Making sure there's no mismatching revision numbers") correct_mismatching_revision_nums(revisions=revisions) diff --git a/tests/always/test_project_structure.py b/tests/always/test_project_structure.py index b846534d8d767..5d1dcc4069f06 100644 --- a/tests/always/test_project_structure.py +++ b/tests/always/test_project_structure.py @@ -169,6 +169,8 @@ def test_providers_modules_should_have_tests(self): modules_files = list(f for f in modules_files if "/_vendor/" not in f) # Exclude __init__.py modules_files = list(f for f in modules_files if not f.endswith("__init__.py")) + # Exclude versions file + modules_files = list(f for f in modules_files if "/versions/" not in f) # Change airflow/ to tests/ expected_test_files = list( f'tests/{f.partition("/")[2]}' for f in modules_files if not f.endswith("__init__.py") diff --git a/tests/providers/fab/auth_manager/cli_commands/test_db_command.py b/tests/providers/fab/auth_manager/cli_commands/test_db_command.py new file mode 100644 index 0000000000000..030b251a55ad6 --- /dev/null +++ b/tests/providers/fab/auth_manager/cli_commands/test_db_command.py @@ -0,0 +1,134 @@ +# 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 + +from unittest import mock + +import pytest + +from airflow.cli import cli_parser + +pytestmark = [pytest.mark.db_test] +try: + from airflow.providers.fab.auth_manager.cli_commands import db_command + from airflow.providers.fab.auth_manager.models.db import FABDBManager + + class TestFABCLiDB: + @classmethod + def setup_class(cls): + cls.parser = cli_parser.get_parser() + + @mock.patch.object(FABDBManager, "resetdb") + def test_cli_resetdb(self, mock_resetdb): + db_command.resetdb(self.parser.parse_args(["fab-db", "reset", "--yes"])) + + mock_resetdb.assert_called_once_with(skip_init=False) + + @mock.patch.object(FABDBManager, "resetdb") + def test_cli_resetdb_skip_init(self, mock_resetdb): + db_command.resetdb(self.parser.parse_args(["fab-db", "reset", "--yes", "--skip-init"])) + mock_resetdb.assert_called_once_with(skip_init=True) + + @pytest.mark.parametrize( + "args, called_with", + [ + ( + [], + dict( + to_revision=None, + from_revision=None, + show_sql_only=False, + ), + ), + ( + ["--show-sql-only"], + dict( + to_revision=None, + from_revision=None, + show_sql_only=True, + ), + ), + ( + ["--to-revision", "abc"], + dict( + to_revision="abc", + from_revision=None, + show_sql_only=False, + ), + ), + ( + ["--to-revision", "abc", "--show-sql-only"], + dict(to_revision="abc", from_revision=None, show_sql_only=True), + ), + ( + ["--to-revision", "abc", "--from-revision", "abc123", "--show-sql-only"], + dict( + to_revision="abc", + from_revision="abc123", + show_sql_only=True, + ), + ), + ], + ) + @mock.patch.object(FABDBManager, "upgradedb") + def test_cli_upgrade_success(self, mock_upgradedb, args, called_with): + db_command.migratedb(self.parser.parse_args(["fab-db", "migrate", *args])) + mock_upgradedb.assert_called_once_with(**called_with) + + @pytest.mark.parametrize( + "args, pattern", + [ + pytest.param( + ["--to-revision", "abc", "--to-version", "1.3.0"], + "Cannot supply both", + id="to both version and revision", + ), + pytest.param( + ["--from-revision", "abc", "--from-version", "1.3.0"], + "Cannot supply both", + id="from both version and revision", + ), + pytest.param(["--to-version", "1.2.0"], "Unknown version '1.2.0'", id="unknown to version"), + pytest.param(["--to-version", "abc"], "Invalid version 'abc'", id="invalid to version"), + pytest.param( + ["--to-revision", "abc", "--from-revision", "abc123"], + "used with `--show-sql-only`", + id="requires offline", + ), + pytest.param( + ["--to-revision", "abc", "--from-version", "1.3.0"], + "used with `--show-sql-only`", + id="requires offline", + ), + pytest.param( + ["--to-revision", "abc", "--from-version", "1.1.25", "--show-sql-only"], + "Unknown version '1.1.25'", + id="unknown from version", + ), + pytest.param( + ["--to-revision", "adaf", "--from-version", "abc", "--show-sql-only"], + "Invalid version 'abc'", + id="invalid from version", + ), + ], + ) + @mock.patch.object(FABDBManager, "upgradedb") + def test_cli_migratedb_failure(self, mock_upgradedb, args, pattern): + with pytest.raises(SystemExit, match=pattern): + db_command.migratedb(self.parser.parse_args(["fab-db", "migrate", *args])) +except (ModuleNotFoundError, ImportError): + pass diff --git a/tests/providers/fab/auth_manager/models/test_db.py b/tests/providers/fab/auth_manager/models/test_db.py index 7b0dc345c7188..528e1cbf099ff 100644 --- a/tests/providers/fab/auth_manager/models/test_db.py +++ b/tests/providers/fab/auth_manager/models/test_db.py @@ -17,6 +17,8 @@ from __future__ import annotations import os +import re +from unittest import mock import pytest from alembic.autogenerate import compare_metadata @@ -35,30 +37,33 @@ from airflow.providers.fab.auth_manager.models.db import FABDBManager class TestFABDBManager: - def setup_method(self, session): + def setup_method(self): self.airflow_dir = os.path.dirname(airflow.__file__) - self.db_manager = FABDBManager(session=session) - def test_version_table_name_set(self): - assert self.db_manager.version_table_name == "fab_alembic_version" + def test_version_table_name_set(self, session): + assert FABDBManager(session=session).version_table_name == "alembic_version_fab" - def test_migration_dir_set(self): - assert self.db_manager.migration_dir == f"{self.airflow_dir}/providers/fab/migrations" + def test_migration_dir_set(self, session): + assert ( + FABDBManager(session=session).migration_dir == f"{self.airflow_dir}/providers/fab/migrations" + ) - def test_alembic_file_set(self): - assert self.db_manager.alembic_file == f"{self.airflow_dir}/providers/fab/alembic.ini" + def test_alembic_file_set(self, session): + assert ( + FABDBManager(session=session).alembic_file == f"{self.airflow_dir}/providers/fab/alembic.ini" + ) - def test_supports_table_dropping_set(self): - assert self.db_manager.supports_table_dropping is True + def test_supports_table_dropping_set(self, session): + assert FABDBManager(session=session).supports_table_dropping is True - def test_database_schema_and_sqlalchemy_model_are_in_sync(self): + def test_database_schema_and_sqlalchemy_model_are_in_sync(self, session): def include_object(_, name, type_, *args): - if type_ == "table" and name not in self.db_manager.metadata.tables: + if type_ == "table" and name not in FABDBManager(session=session).metadata.tables: return False return True all_meta_data = MetaData() - for table_name, table in self.db_manager.metadata.tables.items(): + for table_name, table in FABDBManager(session=session).metadata.tables.items(): all_meta_data._add_table(table_name, table.schema, table) # create diff between database schema and SQLAlchemy model mctx = MigrationContext.configure( @@ -72,5 +77,61 @@ def include_object(_, name, type_, *args): diff = compare_metadata(mctx, all_meta_data) assert not diff, "Database schema and SQLAlchemy model are not in sync: " + str(diff) + + @mock.patch("airflow.providers.fab.auth_manager.models.db._offline_migration") + def test_downgrade_sql_no_from(self, mock_om, session, caplog): + FABDBManager(session=session).downgrade(to_revision="abc", show_sql_only=True, from_revision=None) + actual = mock_om.call_args.kwargs["revision"] + assert re.match(r"[a-z0-9]+:abc", actual) is not None + + @mock.patch("airflow.providers.fab.auth_manager.models.db._offline_migration") + def test_downgrade_sql_with_from(self, mock_om, session): + FABDBManager(session=session).downgrade( + to_revision="abc", show_sql_only=True, from_revision="123" + ) + actual = mock_om.call_args.kwargs["revision"] + assert actual == "123:abc" + + @mock.patch("alembic.command.downgrade") + def test_downgrade_invalid_combo(self, mock_om, session): + """can't combine `sql=False` and `from_revision`""" + with pytest.raises(ValueError, match="can't be combined"): + FABDBManager(session=session).downgrade(to_revision="abc", from_revision="123") + + @mock.patch("alembic.command.downgrade") + def test_downgrade_with_from(self, mock_om, session): + FABDBManager(session=session).downgrade(to_revision="abc") + actual = mock_om.call_args.kwargs["revision"] + assert actual == "abc" + + @mock.patch.object(FABDBManager, "get_current_revision") + def test_sqlite_offline_upgrade_raises_with_revision(self, mock_gcr, session): + with mock.patch( + "airflow.providers.fab.auth_manager.models.db.settings.engine.dialect" + ) as dialect: + dialect.name = "sqlite" + with pytest.raises(SystemExit, match="Offline migration not supported for SQLite"): + FABDBManager(session).upgradedb(from_revision=None, to_revision=None, show_sql_only=True) + + @mock.patch("airflow.utils.db_manager.inspect") + @mock.patch.object(FABDBManager, "metadata") + def test_drop_tables(self, mock_metadata, mock_inspect, session): + manager = FABDBManager(session) + connection = mock.MagicMock() + manager.drop_tables(connection) + mock_metadata.drop_all.assert_called_once_with(connection) + + @pytest.mark.parametrize("skip_init", [True, False]) + @mock.patch.object(FABDBManager, "drop_tables") + @mock.patch.object(FABDBManager, "initdb") + @mock.patch("airflow.utils.db.create_global_lock", new=mock.MagicMock) + def test_resetdb(self, mock_initdb, mock_drop_tables, session, skip_init): + manager = FABDBManager(session) + manager.resetdb(skip_init=skip_init) + mock_drop_tables.assert_called_once() + if skip_init: + mock_initdb.assert_not_called() + else: + mock_initdb.assert_called_once() except ModuleNotFoundError: pass diff --git a/tests/utils/test_db.py b/tests/utils/test_db.py index 2308ed3ca8891..2a197c2e6cf72 100644 --- a/tests/utils/test_db.py +++ b/tests/utils/test_db.py @@ -49,6 +49,7 @@ upgradedb, ) from airflow.utils.db_manager import RunDBManager +from tests.test_utils.config import conf_vars pytestmark = [pytest.mark.db_test, pytest.mark.skip_if_database_isolation_mode] @@ -92,7 +93,7 @@ def test_database_schema_and_sqlalchemy_model_are_in_sync(self): # sqlite sequence is used for autoincrementing columns created with `sqlite_autoincrement` option lambda t: (t[0] == "remove_table" and t[1].name == "sqlite_sequence"), # fab version table - lambda t: (t[0] == "remove_table" and t[1].name == "fab_alembic_version"), + lambda t: (t[0] == "remove_table" and t[1].name == "alembic_version_fab"), ] for ignore in ignores: @@ -131,7 +132,8 @@ def test_check_migrations(self): @mock.patch("alembic.command") def test_upgradedb(self, mock_alembic_command): upgradedb() - mock_alembic_command.upgrade.assert_called_once_with(mock.ANY, revision="heads") + mock_alembic_command.upgrade.assert_called_with(mock.ANY, revision="heads") + assert mock_alembic_command.upgrade.call_count == 2 @pytest.mark.parametrize( "from_revision, to_revision", @@ -200,6 +202,10 @@ def test_downgrade_with_from(self, mock_om): assert actual == "abc" @pytest.mark.parametrize("skip_init", [False, True]) + @conf_vars( + {("database", "external_db_managers"): "airflow.providers.fab.auth_manager.models.db.FABDBManager"} + ) + @mock.patch("airflow.providers.fab.auth_manager.models.db.FABDBManager") @mock.patch("airflow.utils.db.create_global_lock", new=MagicMock) @mock.patch("airflow.utils.db.drop_airflow_models") @mock.patch("airflow.utils.db.drop_airflow_moved_tables") @@ -211,6 +217,7 @@ def test_resetdb( mock_init, mock_drop_moved, mock_drop_airflow, + mock_fabdb_manager, skip_init, ): session_mock = MagicMock() diff --git a/tests/utils/test_db_manager.py b/tests/utils/test_db_manager.py index d6fc91f93934a..1c8a6c6c7dfc8 100644 --- a/tests/utils/test_db_manager.py +++ b/tests/utils/test_db_manager.py @@ -23,7 +23,7 @@ from airflow.exceptions import AirflowException from airflow.models import Base -from airflow.utils.db import downgrade, initdb, upgradedb +from airflow.utils.db import downgrade, initdb from airflow.utils.db_manager import BaseDBManager, RunDBManager from tests.test_utils.config import conf_vars @@ -57,29 +57,26 @@ def test_defining_table_same_name_as_airflow_table_name_raises(self): run_db_manager.validate() metadata._remove_table("dag_run", None) - @mock.patch.object(RunDBManager, "downgradedb") + @mock.patch.object(RunDBManager, "downgrade") @mock.patch.object(RunDBManager, "upgradedb") @mock.patch.object(RunDBManager, "initdb") def test_init_db_calls_rundbmanager(self, mock_initdb, mock_upgrade_db, mock_downgrade_db, session): initdb(session=session) mock_initdb.assert_called() mock_initdb.assert_called_once_with(session) - mock_upgrade_db.assert_not_called() mock_downgrade_db.assert_not_called() - @mock.patch.object(RunDBManager, "downgradedb") + @mock.patch.object(RunDBManager, "downgrade") @mock.patch.object(RunDBManager, "upgradedb") @mock.patch.object(RunDBManager, "initdb") @mock.patch("alembic.command") - def test_upgradedb_or_downgrade_dont_call_rundbmanager( + def test_downgrade_dont_call_rundbmanager( self, mock_alembic_command, mock_initdb, mock_upgrade_db, mock_downgrade_db, session ): - upgradedb(session=session) - mock_alembic_command.upgrade.assert_called_once_with(mock.ANY, revision="heads") downgrade(to_revision="base") mock_alembic_command.downgrade.assert_called_once_with(mock.ANY, revision="base", sql=False) - mock_initdb.assert_not_called() mock_upgrade_db.assert_not_called() + mock_initdb.assert_not_called() mock_downgrade_db.assert_not_called() @conf_vars( @@ -96,12 +93,12 @@ def test_rundbmanager_calls_dbmanager_methods(self, mock_fabdb_manager, session) # upgradedb ext_db.upgradedb(session=session) fabdb_manager.upgradedb.assert_called_once() - # downgradedb - ext_db.downgradedb(session=session) - mock_fabdb_manager.return_value.downgradedb.assert_called_once() + # downgrade + ext_db.downgrade(session=session) + mock_fabdb_manager.return_value.downgrade.assert_called_once() connection = mock.MagicMock() - ext_db.drop_tables(connection) - mock_fabdb_manager.metadata.drop_all.assert_called_once_with(connection) + ext_db.drop_tables(session, connection) + mock_fabdb_manager.return_value.drop_tables.assert_called_once_with(connection) class MockDBManager(BaseDBManager): @@ -127,8 +124,14 @@ def test_create_db_from_orm_called_from_init( @mock.patch.object(BaseDBManager, "get_alembic_config") @mock.patch("alembic.command.upgrade") - def test_upgradedb(self, mock_alembic_cmd, mock_alembic_config, session, caplog): + def test_upgrade(self, mock_alembic_cmd, mock_alembic_config, session, caplog): manager = MockDBManager(session) manager.upgradedb() mock_alembic_cmd.assert_called_once() assert "Upgrading the MockDBManager database" in caplog.text + + @mock.patch.object(BaseDBManager, "get_script_object") + @mock.patch.object(BaseDBManager, "get_current_revision") + def test_check_migration(self, mock_script_obj, mock_current_revision, session): + manager = MockDBManager(session) + manager.check_migration() # just ensure this can be called