Make airflow upgrade_check a command from a separate dist (#12397)

By doing it this way we can add new features to upgrade_check, such as
the adding a flag to make changes automatically, without having to
release a new version of Airflow.

(cherry picked from commit fb90c75)

airflow/bin/cli.py

@@ -72,8 +72,6 @@
 )
 from airflow.ti_deps.dep_context import (DepContext, SCHEDULER_QUEUED_DEPS)
 from airflow.typing_compat import Protocol
-from airflow.upgrade.checker import check_upgrade
-from airflow.upgrade.formatters import (ConsoleFormatter, JSONFormatter)
 from airflow.utils import cli as cli_utils, db
 from airflow.utils.dot_renderer import render_dag
 from airflow.utils.net import get_hostname
@@ -2278,16 +2276,9 @@ def info(args):
 
 
 def upgrade_check(args):
-    if args.save:
-        filename = args.save
-        if not filename.lower().endswith(".json"):
-            print("Only JSON files are supported", file=sys.stderr)
-        formatter = JSONFormatter(args.save)
-    else:
-        formatter = ConsoleFormatter()
-    all_problems = check_upgrade(formatter)
-    if all_problems:
-        sys.exit(1)
+    sys.exit("""
+Please install apache-airflow-upgrade-check distribution from PyPI to perform upgrade checks
+""")
 
 
 class Arg(object):
@@ -3034,9 +3025,10 @@ class CLIFactory(object):
         },
         {
             'name': 'upgrade_check',
-            'help': 'Check if you can upgrade to the new version.',
+            'help': 'Check if you can safely upgrade to the new version.',
             'func': upgrade_check,
-            'args': ('save', ),
+            'from_module': 'airflow.upgrade.checker',
+            'args': (),
         },
     )
     subparsers_dict = {sp['func'].__name__: sp for sp in subparsers}
@@ -3060,6 +3052,14 @@ def error(self, message):
         for sub in subparser_list:
             sub = cls.subparsers_dict[sub]
             sp = subparsers.add_parser(sub['func'].__name__, help=sub['help'])
+            sp.set_defaults(func=sub['func'])
+            if 'from_module' in sub:
+                try:
+                    mod = importlib.import_module(sub['from_module'])
+                    mod.register_arguments(sp)
+                    continue
+                except ImportError:
+                    pass
             for arg in sub['args']:
                 if 'dag_id' in arg and dag_parser:
                     continue
@@ -3068,7 +3068,6 @@ def error(self, message):
                     f: v
                     for f, v in vars(arg).items() if f != 'flags' and v}
                 sp.add_argument(*arg.flags, **kwargs)
-            sp.set_defaults(func=sub['func'])
         return parser
 
 

airflow/upgrade/checker.py

@@ -16,6 +16,8 @@
 # under the License.
 
 from __future__ import absolute_import
+import argparse
+import sys
 from typing import List
 
 from airflow.upgrade.formatters import BaseFormatter
@@ -36,3 +38,32 @@ def check_upgrade(formatter):
         formatter.on_next_rule_status(rule_status)
     formatter.end_checking(all_rule_statuses)
     return all_rule_statuses
+
+
+def register_arguments(subparser):
+    subparser.add_argument(
+        "-s", "--save",
+        help="Saves the result to the indicated file. The file format is determined by the file extension."
+    )
+    subparser.set_defaults(func=run)
+
+
+def run(args):
+    from airflow.upgrade.formatters import (ConsoleFormatter, JSONFormatter)
+    if args.save:
+        filename = args.save
+        if not filename.lower().endswith(".json"):
+            exit("Only JSON files are supported")
+        formatter = JSONFormatter(args.save)
+    else:
+        formatter = ConsoleFormatter()
+    all_problems = check_upgrade(formatter)
+    if all_problems:
+        sys.exit(1)
+
+
+def __main__():
+    parser = argparse.ArgumentParser()
+    register_arguments(parser)
+    args = parser.parse_args()
+    args.func(args)

setup.py

@@ -652,7 +652,7 @@ def do_setup():
         long_description_content_type='text/markdown',
         license='Apache License 2.0',
         version=version,
-        packages=find_packages(exclude=['tests*']),
+        packages=find_packages(exclude=['tests*', 'airflow.upgrade*']),
         package_data={
             '': ['airflow/alembic.ini', "airflow/git_version", "*.ipynb",
                  "airflow/providers/cncf/kubernetes/example_dags/*.yaml"],

tests/upgrade/test_formattes.py

@@ -46,9 +46,10 @@ def test_output(self):
             }
         ]
         parser = cli.CLIFactory.get_parser()
-        with NamedTemporaryFile("w+") as temp:
+        with NamedTemporaryFile("w+", suffix=".json") as temp:
             with pytest.raises(SystemExit):
-                cli.upgrade_check(parser.parse_args(['upgrade_check', '-s', temp.name]))
+                args = parser.parse_args(['upgrade_check', '-s', temp.name])
+                args.func(args)
             content = temp.read()
 
         assert json.loads(content) == expected