diff --git a/airflow/cli/commands/db_command.py b/airflow/cli/commands/db_command.py index b7ae3ee7f3b24..776dabcdead7e 100644 --- a/airflow/cli/commands/db_command.py +++ b/airflow/cli/commands/db_command.py @@ -25,7 +25,7 @@ from tempfile import NamedTemporaryFile from typing import TYPE_CHECKING -from packaging.version import parse as parse_version +from packaging.version import InvalidVersion, parse as parse_version from tenacity import Retrying, stop_after_attempt, wait_fixed from airflow import settings @@ -111,16 +111,24 @@ def migratedb(args): if args.from_revision: from_revision = args.from_revision elif args.from_version: - if parse_version(args.from_version) < parse_version("2.0.0"): + 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 than 2.0.0") from_revision = get_version_revision(args.from_version) if not from_revision: raise SystemExit(f"Unknown version {args.from_version!r} supplied as `--from-version`.") if args.to_version: + try: + 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) if not to_revision: - raise SystemExit(f"Upgrading to version {args.to_version} is not supported.") + raise SystemExit(f"Unknown version {args.to_version!r} supplied as `--to-version`.") elif args.to_revision: to_revision = args.to_revision diff --git a/tests/cli/commands/test_db_command.py b/tests/cli/commands/test_db_command.py index 60fd9ed7a0745..1a68432067114 100644 --- a/tests/cli/commands/test_db_command.py +++ b/tests/cli/commands/test_db_command.py @@ -103,7 +103,18 @@ def test_cli_upgrade_success(self, mock_upgradedb, args, called_with): @pytest.mark.parametrize( "args, pattern", [ - pytest.param(["--to-version", "2.1.25"], "not supported", id="bad version"), + pytest.param( + ["--to-revision", "abc", "--to-version", "2.2.0"], + "Cannot supply both", + id="to both version and revision", + ), + pytest.param( + ["--from-revision", "abc", "--from-version", "2.2.0"], + "Cannot supply both", + id="from both version and revision", + ), + pytest.param(["--to-version", "2.1.25"], "Unknown version '2.1.25'", 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`", @@ -115,9 +126,14 @@ def test_cli_upgrade_success(self, mock_upgradedb, args, called_with): id="requires offline", ), pytest.param( - ["--to-revision", "abc", "--from-version", "2.1.25", "--show-sql-only"], - "Unknown version", - id="bad version", + ["--to-revision", "2.2.0", "--from-version", "2.1.25", "--show-sql-only"], + "Unknown version '2.1.25'", + id="unknown from version", + ), + pytest.param( + ["--to-revision", "2.9.0", "--from-version", "abc", "--show-sql-only"], + "Invalid version 'abc'", + id="invalid from version", ), ], )