diff --git a/airflow/cli/cli_parser.py b/airflow/cli/cli_parser.py index 3888aa0076d52..c7215228c1165 100644 --- a/airflow/cli/cli_parser.py +++ b/airflow/cli/cli_parser.py @@ -224,6 +224,26 @@ def _check(value): default=ColorMode.AUTO, ) +# DB args +ARG_VERSION_RANGE = Arg( + ("-r", "--range"), + help="Version range(start:end) for offline sql generation. Example: '2.0.2:2.2.3'", + default=None, +) +ARG_REVISION_RANGE = Arg( + ('--revision-range',), + help=( + "Migration revision range(start:end) to use for offline sql generation. " + "Example: 'a13f7613ad25:7b2661a43ba3'" + ), + default=None, +) +ARG_REVISION_RANGE = Arg( + ('--revision-range',), + help='Revision range(start:end) to use for offline sql generation', + default=None, +) + # list_dag_runs ARG_DAG_ID_OPT = Arg(("-d", "--dag-id"), help="The id of the dag") ARG_NO_BACKFILL = Arg( @@ -1274,7 +1294,7 @@ class GroupCommand(NamedTuple): name='upgrade', help="Upgrade the metadata database to latest version", func=lazy_load_command('airflow.cli.commands.db_command.upgradedb'), - args=(), + args=(ARG_VERSION_RANGE, ARG_REVISION_RANGE), ), ActionCommand( name='shell', diff --git a/airflow/cli/commands/db_command.py b/airflow/cli/commands/db_command.py index dd502abc2e5a6..09fe2210cf81a 100644 --- a/airflow/cli/commands/db_command.py +++ b/airflow/cli/commands/db_command.py @@ -45,7 +45,7 @@ def resetdb(args): def upgradedb(args): """Upgrades the metadata database""" print("DB: " + repr(settings.engine.url)) - db.upgradedb() + db.upgradedb(version_range=args.range, revision_range=args.revision_range) print("Upgrades done") diff --git a/airflow/migrations/versions/8646922c8a04_change_default_pool_slots_to_1.py b/airflow/migrations/versions/8646922c8a04_change_default_pool_slots_to_1.py index 2f372a384e9de..9dea4d64ca9a2 100644 --- a/airflow/migrations/versions/8646922c8a04_change_default_pool_slots_to_1.py +++ b/airflow/migrations/versions/8646922c8a04_change_default_pool_slots_to_1.py @@ -26,10 +26,6 @@ import sqlalchemy as sa from alembic import op -from sqlalchemy import Column, Integer, String -from sqlalchemy.ext.declarative import declarative_base - -from airflow.utils.sqlalchemy import UtcDateTime # revision identifiers, used by Alembic. revision = '8646922c8a04' @@ -37,37 +33,14 @@ branch_labels = None depends_on = None -Base = declarative_base() -BATCH_SIZE = 5000 - - -class TaskInstance(Base): # type: ignore - """Minimal model definition for migrations""" - - __tablename__ = "task_instance" - - task_id = Column(String(), primary_key=True) - dag_id = Column(String(), primary_key=True) - execution_date = Column(UtcDateTime, primary_key=True) - pool_slots = Column(Integer, default=1) - def upgrade(): """Change default pool_slots to 1 and make pool_slots not nullable""" - connection = op.get_bind() - sessionmaker = sa.orm.sessionmaker() - session = sessionmaker(bind=connection) - - session.query(TaskInstance).filter(TaskInstance.pool_slots.is_(None)).update( - {TaskInstance.pool_slots: 1}, synchronize_session=False - ) - session.commit() - with op.batch_alter_table("task_instance", schema=None) as batch_op: - batch_op.alter_column("pool_slots", existing_type=sa.Integer, nullable=False) + batch_op.alter_column("pool_slots", existing_type=sa.Integer, nullable=False, server_default='1') def downgrade(): """Unapply Change default pool_slots to 1""" with op.batch_alter_table("task_instance", schema=None) as batch_op: - batch_op.alter_column("pool_slots", existing_type=sa.Integer, nullable=True) + batch_op.alter_column("pool_slots", existing_type=sa.Integer, nullable=True, server_default=None) diff --git a/airflow/models/taskinstance.py b/airflow/models/taskinstance.py index b1043c96605d2..0c8821b3d1498 100644 --- a/airflow/models/taskinstance.py +++ b/airflow/models/taskinstance.py @@ -354,7 +354,7 @@ class TaskInstance(Base, LoggingMixin): unixname = Column(String(1000)) job_id = Column(Integer) pool = Column(String(256), nullable=False) - pool_slots = Column(Integer, default=1, nullable=False) + pool_slots = Column(Integer, default=1, nullable=False, server_default=text("1")) queue = Column(String(256)) priority_weight = Column(Integer) operator = Column(String(1000)) diff --git a/airflow/utils/db.py b/airflow/utils/db.py index 2e45be2dc1906..d452e23fdff90 100644 --- a/airflow/utils/db.py +++ b/airflow/utils/db.py @@ -22,8 +22,9 @@ import sys import time from tempfile import gettempdir -from typing import TYPE_CHECKING, Any, Callable, Iterable, List, Tuple +from typing import TYPE_CHECKING, Any, Callable, Iterable, List, Optional, Tuple +from bcrypt import warnings from sqlalchemy import Table, exc, func, inspect, or_, text from sqlalchemy.orm.session import Session @@ -68,6 +69,21 @@ log = logging.getLogger(__name__) +REVISION_HEADS_MAP = { + "2.0.0": "e959f08ac86c", + "2.0.1": "82b7c48c147f", + "2.0.2": "2e42bb497a22", + "2.1.0": "a13f7613ad25", + "2.1.1": "a13f7613ad25", + "2.1.2": "a13f7613ad25", + "2.1.3": "97cdd93827b8", + "2.1.4": "ccde3e26fe78", + "2.2.0": "7b2661a43ba3", + "2.2.1": "7b2661a43ba3", + "2.2.2": "7b2661a43ba3", + "2.2.3": "be2bfac3da23", +} + def _format_airflow_moved_table_name(source_table, version): return "__".join([settings.AIRFLOW_MOVED_TABLE_PREFIX, version.replace(".", "_"), source_table]) @@ -1001,8 +1017,93 @@ def _check_migration_errors(session: Session = NEW_SESSION) -> Iterable[str]: session.commit() +def _offline_migration(command, config, revision): + log.info("Running offline migrations for revision range %s", revision) + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + logging.disable(logging.CRITICAL) + command.upgrade(config, revision, sql=True) + logging.disable(logging.NOTSET) + + +def _validate_version_range(command, config, version_range): + if ':' not in version_range: + raise AirflowException( + 'Please provide Airflow version range with the format "old_version:new_version"' + ) + lower, upper = version_range.split(':') + + if not REVISION_HEADS_MAP.get(lower) or not REVISION_HEADS_MAP.get(upper): + raise AirflowException('Please provide valid Airflow versions above 2.0.0.') + if REVISION_HEADS_MAP.get(lower) == REVISION_HEADS_MAP.get(upper): + if sys.stdout.isatty(): + size = os.get_terminal_size().columns + else: + size = 0 + print(f"Hey this is your migration script from {lower}, to {upper}, but guess what?".center(size)) + print( + "There is no migration needed as the database has not changed between those versions. " + "You are done.".center(size) + ) + print("""/\\_/\\""".center(size)) + print("""(='_' )""".center(size)) + print("""(,(") (")""".center(size)) + print("""^^^""".center(size)) + return + dbname = settings.engine.dialect.name + if dbname == 'sqlite': + raise AirflowException('SQLite is not supported for offline migration.') + elif dbname == 'mssql' and (lower != '2.2.0' or int(lower.split('.')[1]) < 2): + raise AirflowException( + 'MSSQL is not supported for offline migration in Airflow versions less than 2.2.0.' + ) + revision = f"{REVISION_HEADS_MAP[lower]}:{REVISION_HEADS_MAP[upper]}" + try: + command.history(config, rev_range=revision) + except Exception: + raise AirflowException( + f"Error while checking history for revision range {revision}. " + f"Check that the supplied airflow version is in the format 'old_version:new_version'." + ) + return revision + + +def _validate_revision(command, config, revision_range): + if ':' not in revision_range: + raise AirflowException( + 'Please provide Airflow revision range with the format "old_revision:new_revision"' + ) + dbname = settings.engine.dialect.name + if dbname == 'sqlite': + raise AirflowException('SQLite is not supported for offline migration.') + start_version = '2.0.0' + rev_2_0_0_head = 'e959f08ac86c' + _lowerband, _upperband = revision_range.split(':') + if dbname == 'mssql': + rev_2_2_0_head = '7b2661a43ba3' + head_to_lowerband_range = f"{rev_2_2_0_head}:{_lowerband}" + head_to_upperband_range = f"{rev_2_2_0_head}:{_upperband}" + rev_2_0_0_head = rev_2_2_0_head # for logging purposes + start_version = '2.2.0' + else: + head_to_lowerband_range = f"{rev_2_0_0_head}:{_lowerband}" + head_to_upperband_range = f"{rev_2_0_0_head}:{_upperband}" + for i in [head_to_lowerband_range, head_to_upperband_range]: + try: + command.history(config, rev_range=i) + except Exception: + raise AirflowException( + f"Error while checking history for revision range {i}. " + f"Check that {i.split(':')[1]} is a valid revision. " + f"Supported revision for offline migration is from {rev_2_0_0_head} " + f"which is airflow {start_version} head" + ) + + @provide_session -def upgradedb(session: Session = NEW_SESSION): +def upgradedb( + version_range: Optional[str] = None, revision_range: Optional[str] = None, session: Session = NEW_SESSION +): """Upgrade the database.""" # alembic adds significant import time, so we import it lazily if not settings.SQL_ALCHEMY_CONN: @@ -1012,6 +1113,14 @@ def upgradedb(session: Session = NEW_SESSION): config = _get_alembic_config() config.set_main_option('sqlalchemy.url', settings.SQL_ALCHEMY_CONN.replace('%', '%%')) + if version_range: + revision = _validate_version_range(command, config, version_range) + if not revision: + return + return _offline_migration(command, config, revision) + elif revision_range: + _validate_revision(command, config, revision_range) + return _offline_migration(command, config, revision_range) errors_seen = False for err in _check_migration_errors(session=session): diff --git a/dev/README_RELEASE_AIRFLOW.md b/dev/README_RELEASE_AIRFLOW.md index 2e1f3b2db1431..c9964d98f44b7 100644 --- a/dev/README_RELEASE_AIRFLOW.md +++ b/dev/README_RELEASE_AIRFLOW.md @@ -205,6 +205,7 @@ The Release Candidate artifacts we vote upon should be the exact ones we vote ag - Add a commit that updates `CHANGELOG.md` to add changes from previous version if it has not already added. For now this is done manually, example run `git log --oneline v2-2-test..HEAD --pretty='format:- %s'` and categorize them. - Add section for the release in `UPDATING.md`. If no new entries exist, put "No breaking changes" (e.g. `2.1.4`). +- Update the `REVISION_HEADS_MAP` at airflow/utils/db.py to include the revision head of the release even if there are no migrations. - Commit the version change. - PR from the 'test' branch to the 'stable' branch, and manually merge it once approved. - Check out the 'stable' branch diff --git a/docs/apache-airflow/installation/upgrading.rst b/docs/apache-airflow/installation/upgrading.rst index 929c6044e08d9..6cd45ac04f23a 100644 --- a/docs/apache-airflow/installation/upgrading.rst +++ b/docs/apache-airflow/installation/upgrading.rst @@ -44,6 +44,17 @@ In order to manually upgrade the database you should run the ``airflow db upgrad environment. It can be run either in your virtual environment or in the containers that give you access to Airflow ``CLI`` :doc:`/usage-cli` and the database. +Offline SQL migration scripts +============================= +If you want to run the upgrade script offline, you can use the ``-r`` or ``--revision-range`` flag +to get the SQL statements that would be executed. This feature is supported in Postgres and MySQL +from Airflow 2.0.0 onward and in MSSQL from Airflow 2.2.0 onward. + +Sample usage: + ``airflow db upgrade -r "2.0.0:2.2.0"`` + ``airflow db upgrade --revision-range "e959f08ac86c:142555e44c17"`` + + Migration best practices ======================== diff --git a/tests/cli/commands/test_db_command.py b/tests/cli/commands/test_db_command.py index 8c0ef8a55be89..e94e63a60e01b 100644 --- a/tests/cli/commands/test_db_command.py +++ b/tests/cli/commands/test_db_command.py @@ -53,7 +53,7 @@ def test_cli_check_migrations(self, mock_wait_for_migrations): def test_cli_upgradedb(self, mock_upgradedb): db_command.upgradedb(self.parser.parse_args(['db', 'upgrade'])) - mock_upgradedb.assert_called_once_with() + mock_upgradedb.assert_called_once_with(version_range=None, revision_range=None) @mock.patch("airflow.cli.commands.db_command.execute_interactive") @mock.patch("airflow.cli.commands.db_command.NamedTemporaryFile") diff --git a/tests/utils/test_db.py b/tests/utils/test_db.py index 182f75d162cf4..5d828264f2e41 100644 --- a/tests/utils/test_db.py +++ b/tests/utils/test_db.py @@ -18,8 +18,9 @@ import inspect import re -import unittest +from unittest import mock +import pytest from alembic.autogenerate import compare_metadata from alembic.config import Config from alembic.migration import MigrationContext @@ -27,12 +28,13 @@ from alembic.script import ScriptDirectory from sqlalchemy import MetaData +from airflow.exceptions import AirflowException from airflow.models import Base as airflow_base from airflow.settings import engine -from airflow.utils.db import check_migrations, create_default_connections +from airflow.utils.db import check_migrations, create_default_connections, upgradedb -class TestDb(unittest.TestCase): +class TestDb: def test_database_schema_and_sqlalchemy_model_are_in_sync(self): all_meta_data = MetaData() for (table_name, table) in airflow_base.metadata.tables.items(): @@ -110,3 +112,104 @@ def test_default_connections_sort(self): def test_check_migrations(self): # Should run without error. Can't easily test the behaviour, but we can check it works check_migrations(1) + + @mock.patch('alembic.command') + def test_upgradedb(self, mock_alembic_command): + upgradedb() + mock_alembic_command.upgrade.assert_called_once_with(mock.ANY, 'heads') + + @pytest.mark.parametrize( + 'version, revision', + [('2.0.0:2.2.3', "e959f08ac86c:be2bfac3da23"), ("2.0.2:2.1.4", "2e42bb497a22:ccde3e26fe78")], + ) + def test_offline_upgrade_version(self, version, revision): + with mock.patch('airflow.utils.db.settings.engine.dialect'): + with mock.patch('alembic.command.upgrade') as mock_alembic_upgrade: + upgradedb(version_range=version) + mock_alembic_upgrade.assert_called_once_with(mock.ANY, revision, sql=True) + + @pytest.mark.parametrize( + 'version, revision', + [('2.2.3:2.0.0', "be2bfac3da23:e959f08ac86c"), ("2.1.4:2.0.2", "ccde3e26fe78:2e42bb497a22")], + ) + def test_offline_upgrade_fails_for_migration_incorrect_versions(self, version, revision): + with mock.patch('airflow.utils.db.settings.engine.dialect'): + with pytest.raises(AirflowException) as e: + upgradedb(version) + assert e.exconly() == ( + f"airflow.exceptions.AirflowException: " + f"Error while checking history for revision range {revision}. " + f"Check that the supplied airflow version is in the format 'old_version:new_version'." + ) + + @pytest.mark.parametrize( + 'version, error', + [ + ('2.2.3', 'Please provide Airflow version range with the format "old_version:new_version"'), + ("2.1.2:2.1.5", "Please provide valid Airflow versions above 2.0.0."), + ], + ) + def test_offline_upgrade_fails_for_migration_single_versions_or_not_existing_head(self, version, error): + with pytest.raises(AirflowException) as e: + upgradedb(version) + assert e.exconly() == (f"airflow.exceptions.AirflowException: {error}") + + @pytest.mark.parametrize('revision', ['90d1635d7b86:54bebd308c5f', "e959f08ac86c:587bdf053233"]) + def test_offline_upgrade_revision(self, revision): + with mock.patch('airflow.utils.db.settings.engine.dialect'): + with mock.patch('alembic.command.upgrade') as mock_alembic_upgrade: + upgradedb(revision_range=revision) + mock_alembic_upgrade.assert_called_once_with(mock.ANY, revision, sql=True) + + def test_offline_upgrade_fails_for_migration_less_than_2_0_0_head(self): + rev_2_0_0_head = 'e959f08ac86c' + with mock.patch('airflow.utils.db.settings.engine.dialect'): + with pytest.raises(AirflowException) as e: + upgradedb(revision_range='e1a11ece99cc:54bebd308c5f') + revision = f"{rev_2_0_0_head}:e1a11ece99cc" + assert e.exconly() == ( + f"airflow.exceptions.AirflowException: " + f"Error while checking history for revision range {revision}. " + f"Check that {revision.split(':')[1]} is a valid revision. " + f"Supported revision for offline migration is from {rev_2_0_0_head} " + f"which is airflow 2.0.0 head" + ) + + def test_sqlite_offline_upgrade_raises_with_revision(self): + with mock.patch('airflow.utils.db.settings.engine.dialect') as dialect: + dialect.name = 'sqlite' + with pytest.raises(AirflowException) as e: + upgradedb(revision_range='e1a11ece99cc:54bebd308c5f') + assert e.exconly() == ( + "airflow.exceptions.AirflowException: SQLite is not supported for offline migration." + ) + + def test_sqlite_offline_upgrade_raises_with_version(self): + with mock.patch('airflow.utils.db.settings.engine.dialect') as dialect: + dialect.name = 'sqlite' + with pytest.raises(AirflowException) as e: + upgradedb(revision_range='2.0.0:2.2.3') + assert e.exconly() == ( + "airflow.exceptions.AirflowException: SQLite is not supported for offline migration." + ) + + def test_offline_upgrade_fails_for_migration_less_than_2_2_0_head_for_mssql(self): + rev_2_2_0_head = '7b2661a43ba3' + with mock.patch('airflow.utils.db.settings.engine.dialect') as dialect: + dialect.name = 'mssql' + with pytest.raises(AirflowException) as e: + upgradedb(revision_range='e1a11ece99cc:54bebd308c5f') + revision = f"{rev_2_2_0_head}:e1a11ece99cc" + assert e.exconly() == ( + f"airflow.exceptions.AirflowException: " + f"Error while checking history for revision range {revision}. " + f"Check that {revision.split(':')[1]} is a valid revision. " + f"Supported revision for offline migration is from {rev_2_2_0_head} " + f"which is airflow 2.2.0 head" + ) + + def test_versions_without_migration_donot_raise(self): + with mock.patch('airflow.utils.db.settings.engine.dialect'): + with mock.patch('alembic.command.upgrade') as mock_alembic_upgrade: + upgradedb("2.1.1:2.1.2") + mock_alembic_upgrade.assert_not_called()