From 6f779e34519f953b00435cdd7e02f6421ad83ae6 Mon Sep 17 00:00:00 2001 From: Ephraim Anierobi Date: Wed, 19 Jan 2022 18:28:08 +0100 Subject: [PATCH 1/3] Support generating SQL script for upgrades This PR attempts to add support for generating sql scripts for upgrade. Example command: `airflow db upgrade e8d98d8ss99 --sql` `airflow db upgrade --sql` Support offline mode from 2.0.0 upwards Add tests fixup! Add tests fixup! fixup! Add tests add server default to pool_slots Fix single digit revision in offline mode to start from 2.0.0 head fixup! Fix single digit revision in offline mode to start from 2.0.0 head Apply suggestions from code review Separate offline code from online to reduce log noise fixup! Separate offline code from online to reduce log noise Use version instead of revision fixup! Use version instead of revision add back revision range fixup! add back revision range Add instruction to update revision map during release Exclude offline migration for sqlite and mssql less than 2.2.0 Use version range instead of just version in cli parser print nice message for versions without migrations fixup! print nice message for versions without migrations Support generating SQL script for upgrades This PR attempts to add support for generating sql scripts for upgrade. Example command: `airflow db upgrade e8d98d8ss99 --sql` `airflow db upgrade --sql` Support offline mode from 2.0.0 upwards Add tests fixup! Add tests fixup! fixup! Add tests Fix single digit revision in offline mode to start from 2.0.0 head fixup! Fix single digit revision in offline mode to start from 2.0.0 head Separate offline code from online to reduce log noise Use version instead of revision print nice message for versions without migrations fixup! print nice message for versions without migrations Apply suggestions from code review Co-authored-by: Ash Berlin-Taylor use range instead of version-range Fix test Support generating SQL script for upgrades This PR attempts to add support for generating sql scripts for upgrade. Example command: `airflow db upgrade e8d98d8ss99 --sql` `airflow db upgrade --sql` Support offline mode from 2.0.0 upwards Add tests fixup! Add tests fixup! fixup! Add tests Fix single digit revision in offline mode to start from 2.0.0 head fixup! Fix single digit revision in offline mode to start from 2.0.0 head Separate offline code from online to reduce log noise Use version instead of revision add back revision range Use version range instead of just version in cli parser print nice message for versions without migrations fixup! print nice message for versions without migrations Support generating SQL script for upgrades This PR attempts to add support for generating sql scripts for upgrade. Example command: `airflow db upgrade e8d98d8ss99 --sql` `airflow db upgrade --sql` Support offline mode from 2.0.0 upwards Add tests fixup! Add tests fixup! fixup! Add tests Fix single digit revision in offline mode to start from 2.0.0 head fixup! Fix single digit revision in offline mode to start from 2.0.0 head Separate offline code from online to reduce log noise Use version instead of revision print nice message for versions without migrations fixup! print nice message for versions without migrations Apply suggestions from code review Co-authored-by: Ash Berlin-Taylor Apply suggestions from code review Co-authored-by: Kaxil Naik fix conflict --- airflow/cli/cli_parser.py | 22 +++- airflow/cli/commands/db_command.py | 2 +- ...922c8a04_change_default_pool_slots_to_1.py | 31 +---- airflow/models/taskinstance.py | 2 +- airflow/utils/db.py | 113 +++++++++++++++++- dev/README_RELEASE_AIRFLOW.md | 1 + tests/cli/commands/test_db_command.py | 2 +- tests/utils/test_db.py | 107 ++++++++++++++++- 8 files changed, 242 insertions(+), 38 deletions(-) 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/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..870416dd983c5 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,102 @@ 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 the supplied airflow version is in the format 'old_version:new_version'." + ) + + 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() From 3d19303a1fa2b42bf10e5ee12b23fbf3b7b57b5a Mon Sep 17 00:00:00 2001 From: Ephraim Anierobi Date: Wed, 16 Feb 2022 07:44:54 +0100 Subject: [PATCH 2/3] Add documentation --- docs/apache-airflow/installation/upgrading.rst | 11 +++++++++++ 1 file changed, 11 insertions(+) diff --git a/docs/apache-airflow/installation/upgrading.rst b/docs/apache-airflow/installation/upgrading.rst index 929c6044e08d9..de7dd2ca81490 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 onwards and in MSSQL from Airflow 2.2.0 onwards. + +Sample usage: + ``airflow db upgrade -r "2.0.0:2.2.0"`` + ``airflow db upgrade --revision-range "e959f08ac86c:142555e44c17"`` + + Migration best practices ======================== From 0d747bf98ae98f147fe3d2ded6eaba514d982468 Mon Sep 17 00:00:00 2001 From: Ephraim Anierobi Date: Wed, 16 Feb 2022 11:13:58 +0100 Subject: [PATCH 3/3] fix test and doc --- docs/apache-airflow/installation/upgrading.rst | 2 +- tests/utils/test_db.py | 4 +++- 2 files changed, 4 insertions(+), 2 deletions(-) diff --git a/docs/apache-airflow/installation/upgrading.rst b/docs/apache-airflow/installation/upgrading.rst index de7dd2ca81490..6cd45ac04f23a 100644 --- a/docs/apache-airflow/installation/upgrading.rst +++ b/docs/apache-airflow/installation/upgrading.rst @@ -48,7 +48,7 @@ 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 onwards and in MSSQL from Airflow 2.2.0 onwards. +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"`` diff --git a/tests/utils/test_db.py b/tests/utils/test_db.py index 870416dd983c5..5d828264f2e41 100644 --- a/tests/utils/test_db.py +++ b/tests/utils/test_db.py @@ -170,7 +170,9 @@ def test_offline_upgrade_fails_for_migration_less_than_2_0_0_head(self): 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'." + 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):