Better db migrate error messages (#39268)

apache · Apr 26, 2024 · 8188da7 · 8188da7
1 parent 0759e6f
commit 8188da7
Show file tree

Hide file tree

Showing 2 changed files with 31 additions and 7 deletions.
diff --git 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
@@ -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",
             ),
         ],
     )