[MAINTENANCE] Make CLI Check-Config and CLI More Robust (#3562)

* TZLocal being (most likely temporarily) pinned to version 3 (due to a defect in version 4).
great-expectations · Oct 18, 2021 · 863ac93 · 863ac93
1 parent 5be44df
commit 863ac93
Show file tree

Hide file tree

Showing 3 changed files with 87 additions and 27 deletions.
diff --git a/great_expectations/cli/project.py b/great_expectations/cli/project.py
@@ -65,13 +65,24 @@ def do_config_check(target_directory):
             upgrade_message: str = f"""The config_version of your great_expectations.yml -- {float(ge_config_version)} -- is outdated.
 Please consult the 0.13.x migration guide https://docs.greatexpectations.io/en/latest/guides/how_to_guides/migrating_versions.html and
 upgrade your Great Expectations configuration to version {float(CURRENT_GE_CONFIG_VERSION)} in order to take advantage of the latest capabilities.
-            """
+"""
             return (
                 False,
                 upgrade_message,
                 None,
             )
-        return True, None, context
+        elif int(ge_config_version) > CURRENT_GE_CONFIG_VERSION:
+            raise ge_exceptions.UnsupportedConfigVersionError(
+                f"""Invalid config version ({ge_config_version}).\n    The maximum valid version is \
+{CURRENT_GE_CONFIG_VERSION}.
+"""
+            )
+        else:
+            return (
+                True,
+                None,
+                context,
+            )
     except (
         ge_exceptions.InvalidConfigurationYamlError,
         ge_exceptions.InvalidTopLevelConfigKeyError,

diff --git a/great_expectations/cli/toolkit.py b/great_expectations/cli/toolkit.py
@@ -1,4 +1,3 @@
-import datetime
 import json
 import os
 import subprocess
@@ -411,7 +410,12 @@ def load_data_context_with_error_handling(
     """Return a DataContext with good error handling and exit codes."""
     try:
         context: DataContext = DataContext(context_root_dir=directory)
-        if from_cli_upgrade_command:
+        ge_config_version: int = context.get_config().config_version
+        if (
+            from_cli_upgrade_command
+            and int(ge_config_version) < CURRENT_GE_CONFIG_VERSION
+        ):
+            # noinspection PyBroadException
             try:
                 send_usage_message(
                     data_context=context,
@@ -421,29 +425,61 @@ def load_data_context_with_error_handling(
             except Exception:
                 # Don't fail for usage stats
                 pass
-        ge_config_version: int = context.get_config().config_version
-        if (
-            from_cli_upgrade_command
-            and int(ge_config_version) < CURRENT_GE_CONFIG_VERSION
-        ):
-            directory = directory or context.root_directory
-            (
-                increment_version,
-                exception_occurred,
-            ) = upgrade_project_one_version_increment(
-                context_root_dir=directory,
-                ge_config_version=ge_config_version,
-                continuation_message=EXIT_UPGRADE_CONTINUATION_MESSAGE,
-                from_cli_upgrade_command=from_cli_upgrade_command,
-            )
-            if not exception_occurred and increment_version:
-                context = DataContext(context_root_dir=directory)
-                if from_cli_upgrade_command:
-                    send_usage_message(
-                        data_context=context,
-                        event="cli.project.upgrade.end",
-                        success=True,
+
+            if (CURRENT_GE_CONFIG_VERSION - int(ge_config_version)) == 1:
+                directory = directory or context.root_directory
+                (
+                    increment_version,
+                    exception_occurred,
+                ) = upgrade_project_one_version_increment(
+                    context_root_dir=directory,
+                    ge_config_version=ge_config_version,
+                    continuation_message=EXIT_UPGRADE_CONTINUATION_MESSAGE,
+                    from_cli_upgrade_command=from_cli_upgrade_command,
+                )
+                if not exception_occurred and increment_version:
+                    context = DataContext(context_root_dir=directory)
+                    # noinspection PyBroadException
+                    try:
+                        send_usage_message(
+                            data_context=context,
+                            event="cli.project.upgrade.end",
+                            success=True,
+                        )
+                    except Exception:
+                        # Don't fail for usage stats
+                        pass
+            else:
+                directory = directory or DataContext.find_context_root_dir()
+                ge_config_version = DataContext.get_ge_config_version(
+                    context_root_dir=directory
+                )
+                upgrade_helper_class = (
+                    GE_UPGRADE_HELPER_VERSION_MAP.get(int(ge_config_version))
+                    if ge_config_version
+                    else None
+                )
+                if upgrade_helper_class:
+                    upgrade_project(
+                        context_root_dir=directory,
+                        ge_config_version=ge_config_version,
+                        from_cli_upgrade_command=from_cli_upgrade_command,
                     )
+                    context = DataContext(context_root_dir=directory)
+                    # noinspection PyBroadException
+                    try:
+                        send_usage_message(
+                            data_context=context,
+                            event="cli.project.upgrade.end",
+                            success=True,
+                        )
+                    except Exception:
+                        # Don't fail for usage stats
+                        pass
+                else:
+                    error_message: str = f"The upgrade utility for version {int(ge_config_version)} could not be found."
+                    cli_message(string=f"<red>{error_message}</red>")
+                    sys.exit(1)
         return context
     except ge_exceptions.UnsupportedConfigVersionError as err:
         directory = directory or DataContext.find_context_root_dir()
@@ -461,6 +497,17 @@ def load_data_context_with_error_handling(
                 ge_config_version=ge_config_version,
                 from_cli_upgrade_command=from_cli_upgrade_command,
             )
+            context = DataContext(context_root_dir=directory)
+            # noinspection PyBroadException
+            try:
+                send_usage_message(
+                    data_context=context,
+                    event="cli.project.upgrade.end",
+                    success=True,
+                )
+            except Exception:
+                # Don't fail for usage stats
+                pass
         else:
             cli_message(string=f"<red>{err.message}</red>")
             sys.exit(1)
@@ -503,6 +550,7 @@ def upgrade_project(
         "\nWould you like to run the Upgrade Helper to bring your project up-to-date?"
     )
     # This loading of DataContext is optional and just to track if someone exits here
+    # noinspection PyBroadException
     try:
         data_context = DataContext(context_root_dir)
     except Exception:
@@ -553,6 +601,7 @@ def upgrade_project(
             pass
     else:
         cli_message(upgrade_success_message)
+        # noinspection PyBroadException
         try:
             context: DataContext = DataContext(context_root_dir)
             send_usage_message(

diff --git a/requirements.txt b/requirements.txt
@@ -22,4 +22,4 @@ ruamel.yaml>=0.16  # package
 scipy>=0.19.0  # package
 termcolor>=1.1.0  # package
 tqdm>=4.59.0
-tzlocal>=1.2  # package
+tzlocal==3.0  # package