Update upgrade command in preparation for release

src/_nebari/constants.py

@@ -1,3 +1,5 @@
+CURRENT_RELEASE = "2023.7.1"
+
 TERRAFORM_VERSION = "1.0.5"
 
 # 04-kubernetes-ingress
@@ -6,11 +8,10 @@
 HIGHEST_SUPPORTED_K8S_VERSION = "1.24.13"
 DEFAULT_GKE_RELEASE_CHANNEL = "UNSPECIFIED"
 
-DEFAULT_NEBARI_DASK_VERSION = "2023.5.1"
-DEFAULT_NEBARI_IMAGE_TAG = "2023.5.1"
+DEFAULT_NEBARI_DASK_VERSION = CURRENT_RELEASE
+DEFAULT_NEBARI_IMAGE_TAG = CURRENT_RELEASE
+DEFAULT_NEBARI_WORKFLOW_CONTROLLER_IMAGE_TAG = CURRENT_RELEASE
 
 DEFAULT_CONDA_STORE_IMAGE_TAG = "v0.4.14"
 
-DEFAULT_NEBARI_WORKFLOW_CONTROLLER_IMAGE_TAG = "update_nwc-05c3b99-20230512"
-
 LATEST_SUPPORTED_PYTHON_VERSION = "3.10"

src/_nebari/upgrade.py

@@ -16,6 +16,11 @@
 
 logger = logging.getLogger(__name__)
 
+NEBARI_WORKFLOW_CONTROLLER_DOCS = (
+    "https://www.nebari.dev/docs/how-tos/using-argo/#jupyterflow-override-beta"
+)
+ARGO_JUPYTER_SCHEDULER_REPO = "https://github.com/nebari-dev/argo-jupyter-scheduler"
+
 
 def do_upgrade(config_filename, attempt_fixes=False):
     config = load_yaml(config_filename)
@@ -34,32 +39,35 @@ def do_upgrade(config_filename, attempt_fixes=False):
     except (ValidationError, ValueError) as e:
         if is_version_accepted(config.get("nebari_version", "")):
             # There is an unrelated validation problem
-            print(
-                f"Your config file {config_filename} appears to be already up-to-date for Nebari version {__version__} but there is another validation error.\n"
+            rich.print(
+                f"Your config file [purple]{config_filename}[/purple] appears to be already up-to-date for Nebari version [green]{__version__}[/green] but there is another validation error.\n"
             )
             raise e
 
     start_version = config.get("nebari_version", "")
+    print("start_version: ", start_version)
 
     UpgradeStep.upgrade(
         config, start_version, __version__, config_filename, attempt_fixes
     )
+    print(config.get("nebari_version"))
+    print(config.get("default_images"))
 
     # Backup old file
     backup_config_file(config_filename, f".{start_version or 'old'}")
 
     with config_filename.open("wt") as f:
         yaml.dump(config, f)
 
-    print(
-        f"Saving new config file {config_filename} ready for Nebari version {__version__}"
+    rich.print(
+        f"Saving new config file [purple]{config_filename}[/purple] ready for Nebari version [green]{__version__}[/green]"
     )
 
     ci_cd = config.get("ci_cd", {}).get("type", "")
     if ci_cd in ("github-actions", "gitlab-ci"):
-        print(
-            f"\nSince you are using ci_cd {ci_cd} you also need to re-render the workflows and re-commit the files to your Git repo:\n"
-            f"   nebari render -c {config_filename}\n"
+        rich.print(
+            f"\nSince you are using ci_cd [green]{ci_cd}[/green] you also need to re-render the workflows and re-commit the files to your Git repo:\n"
+            f"   nebari render -c [purple]{config_filename}[/purple]\n"
         )
 
 
@@ -89,6 +97,7 @@ def upgrade(
         """
         starting_ver = rounded_ver_parse(start_version or "0.0.0")
         finish_ver = rounded_ver_parse(finish_version)
+        print("finish_ver: ", finish_ver)
 
         if finish_ver < starting_ver:
             raise ValueError(
@@ -106,6 +115,8 @@ def upgrade(
             key=rounded_ver_parse,
         )
 
+        print("step_versions: ", step_versions)
+
         current_start_version = start_version
         for stepcls in [cls._steps[str(v)] for v in step_versions]:
             step = stepcls()
@@ -144,9 +155,8 @@ def upgrade_step(self, config, start_version, config_filename, *args, **kwargs):
         __rounded_finish_version__ = ".".join(
             [str(c) for c in rounded_ver_parse(finish_version)]
         )
-
-        print(
-            f"\n---> Starting upgrade from {start_version or 'old version'} to {finish_version}\n"
+        rich.print(
+            f"\n---> Starting upgrade from [green]{start_version or 'old version'}[/green] to [green]{finish_version}[/green]\n"
         )
 
         # Set the new version
@@ -155,51 +165,97 @@ def upgrade_step(self, config, start_version, config_filename, *args, **kwargs):
         assert self.version != start_version
 
         if self.requires_nebari_version_field():
-            print(f"Setting nebari_version to {self.version}")
+            rich.print(f"Setting nebari_version to [green]{self.version}[/green]")
             config["nebari_version"] = self.version
 
-        # Update images
-        start_version_regex = start_version.replace(".", "\\.")
-        if start_version == "":
-            print("Looking for any previous image version")
-            start_version_regex = "0\\.[0-3]\\.[0-9]{1,2}"
-        docker_image_regex = re.compile(
-            f"^([A-Za-z0-9_-]+/[A-Za-z0-9_-]+):v{start_version_regex}$"
-        )
+        def contains_image_and_tag(s: str) -> bool:
+            # match on `quay.io/nebari/nebari-<...>:YYYY.MM.XX``
+            pattern = r"^quay\.io\/nebari\/nebari-(jupyterhub|jupyterlab|dask-worker):\d{4}\.\d+\.\d+$"
+            return bool(re.match(pattern, s))
 
-        def _new_docker_image(
-            v,
-        ):
-            m = docker_image_regex.match(v)
+        def replace_image_tag_legacy(image, start_version, new_version):
+            start_version_regex = start_version.replace(".", "\\.")
+            if not start_version:
+                start_version_regex = "0\\.[0-3]\\.[0-9]{1,2}"
+
+            docker_image_regex = re.compile(
+                f"^([A-Za-z0-9_-]+/[A-Za-z0-9_-]+):v{start_version_regex}$"
+            )
+
+            m = docker_image_regex.match(image)
             if m:
-                return ":".join([m.groups()[0], f"v{__rounded_finish_version__}"])
+                return ":".join([m.groups()[0], f"v{new_version}"])
             return None
 
+        def replace_image_tag(s: str, new_version: str, config_path: str) -> str:
+            legacy_replacement = replace_image_tag_legacy(s, start_version, new_version)
+            if legacy_replacement:
+                return legacy_replacement
+
+            if not contains_image_and_tag(s):
+                return s
+            image_name, current_tag = s.split(":")
+            if current_tag == new_version:
+                return s
+            loc = f"{config_path}: {image_name}"
+            response = Prompt.ask(
+                f"\nDo you want to replace current tag [green]{current_tag}[/green] with [green]{new_version}[/green] for:\n[purple]{loc}[/purple]? [Y/n] ",
+                default="Y",
+            )
+            if response.lower() in ["y", "yes", ""]:
+                return s.replace(current_tag, new_version)
+            else:
+                return s
+
+        def set_nested_item(config: dict, config_path: list, value: str):
+            config_path = config_path.split(".")
+            for k in config_path[:-1]:
+                try:
+                    k = int(k)
+                except ValueError:
+                    pass
+                config = config[k]
+            try:
+                config_path[-1] = int(config_path[-1])
+            except ValueError:
+                pass
+            config[config_path[-1]] = value
+
+        def update_image_tag(config, config_path, current_image, new_version):
+            new_image = replace_image_tag(current_image, new_version, config_path)
+            if new_image != current_image:
+                set_nested_item(config, config_path, new_image)
+
+            return config
+
+        # update default_images
         for k, v in config.get("default_images", {}).items():
-            newimage = _new_docker_image(v)
-            if newimage:
-                print(f"In default_images: {k}: upgrading {v} to {newimage}")
-                config["default_images"][k] = newimage
+            config_path = f"default_images.{k}"
+            config = update_image_tag(
+                config, config_path, v, __rounded_finish_version__
+            )
 
+        # update profiles.jupyterlab images
         for i, v in enumerate(config.get("profiles", {}).get("jupyterlab", [])):
-            oldimage = v.get("kubespawner_override", {}).get("image", "")
-            newimage = _new_docker_image(oldimage)
-            if newimage:
-                print(
-                    f"In profiles: jupyterlab: [{i}]: upgrading {oldimage} to {newimage}"
+            current_image = v.get("kubespawner_override", {}).get("image", None)
+            if current_image:
+                config = update_image_tag(
+                    config,
+                    f"profiles.jupyterlab.{i}.kubespawner_override.image",
+                    current_image,
+                    __rounded_finish_version__,
                 )
-                config["profiles"]["jupyterlab"][i]["kubespawner_override"][
-                    "image"
-                ] = newimage
 
+        # update profiles.dask_worker images
         for k, v in config.get("profiles", {}).get("dask_worker", {}).items():
-            oldimage = v.get("image", "")
-            newimage = _new_docker_image(oldimage)
-            if newimage:
-                print(
-                    f"In profiles: dask_worker: {k}: upgrading {oldimage} to {newimage}"
+            current_image = v.get("kubespawner_override", {}).get("image", None)
+            if current_image:
+                config = update_image_tag(
+                    config,
+                    f"profiles.dask_worker.{k}.kubespawner_override.image",
+                    current_image,
+                    __rounded_finish_version__,
                 )
-                config["profiles"]["dask_worker"][k]["image"] = newimage
 
         # Run any version-specific tasks
         return self._version_specific_upgrade(
@@ -226,7 +282,9 @@ def _version_specific_upgrade(
         """
         if config.get("default_images", {}).get("conda_store", None) is None:
             newimage = "quansight/conda-store-server:v0.3.3"
-            print(f"Adding default_images: conda_store image as {newimage}")
+            rich.print(
+                f"Adding default_images: conda_store image as [green]{newimage}[/green]"
+            )
             config["default_images"]["conda_store"] = newimage
         return config
 
@@ -259,8 +317,8 @@ def _version_specific_upgrade(
                     f"{customauth_warning}\n\nRun `nebari upgrade --attempt-fixes` to switch to basic Keycloak authentication instead."
                 )
             else:
-                print(f"\nWARNING: {customauth_warning}")
-                print(
+                rich.print(f"\nWARNING: {customauth_warning}")
+                rich.print(
                     "\nSwitching to basic Keycloak authentication instead since you specified --attempt-fixes."
                 )
                 config["security"]["authentication"] = {"type": "password"}
@@ -297,8 +355,8 @@ def _version_specific_upgrade(
         with realm_import_filename.open("wt") as f:
             json.dump(realm, f, indent=2)
 
-        print(
-            f"\nSaving user/group import file {realm_import_filename}.\n\n"
+        rich.print(
+            f"\nSaving user/group import file [purple]{realm_import_filename}[/purple].\n\n"
             "ACTION REQUIRED: You must import this file into the Keycloak admin webpage after you redeploy Nebari.\n"
             "Visit the URL path /auth/ and login as 'root'. Under Manage, click Import and select this file.\n\n"
             "Non-admin users will default to analyst group membership after the upgrade (no dask access), "
@@ -315,7 +373,7 @@ def _version_specific_upgrade(
 
         if "terraform_modules" in config:
             del config["terraform_modules"]
-            print(
+            rich.print(
                 "Removing terraform_modules field from config as it is no longer used.\n"
             )
 
@@ -333,8 +391,8 @@ def _version_specific_upgrade(
         )
         security.setdefault("keycloak", {})["initial_root_password"] = default_password
 
-        print(
-            f"Generated default random password={default_password} for Keycloak root user (Please change at /auth/ URL path).\n"
+        rich.print(
+            f"Generated default random password=[green]{default_password}[/green] for Keycloak root user (Please change at /auth/ URL path).\n"
         )
 
         # project was never needed in Azure - it remained as PLACEHOLDER in earlier nebari inits!
@@ -369,7 +427,7 @@ def _version_specific_upgrade(
         """
         Upgrade jupyterlab profiles.
         """
-        print("\nUpgrading jupyterlab profiles in order to specify access type:\n")
+        rich.print("\nUpgrading jupyterlab profiles in order to specify access type:\n")
 
         profiles_jupyterlab = config.get("profiles", {}).get("jupyterlab", [])
         for profile in profiles_jupyterlab:
@@ -380,8 +438,8 @@ def _version_specific_upgrade(
             else:
                 profile["access"] = "all"
 
-            print(
-                f"Setting access type of JupyterLab profile {name} to {profile['access']}"
+            rich.print(
+                f"Setting access type of JupyterLab profile [green]{name}[/green] to [green]{profile['access']}[/green]"
             )
         return config
 
@@ -408,18 +466,41 @@ def _version_specific_upgrade(
         )
 
         continue_ = Prompt.ask(
-            "Have you deleted the Argo Workflows CRDs and service accounts? [y/N]",
+            "Have you deleted the Argo Workflows CRDs and service accounts? [y/N] ",
             default="N",
         )
         if not continue_ == "y":
-            print(
-                f"You must delete the Argo Workflows CRDs and service accounts before upgrading to {self.version} (or later)."
+            rich.print(
+                f"You must delete the Argo Workflows CRDs and service accounts before upgrading to [green]{self.version}[/green] (or later)."
             )
             exit()
 
         return config
 
 
+class Upgrade_2023_5_2(UpgradeStep):
+    version = "2023.5.2"
+
+    def _version_specific_upgrade(
+        self, config, start_version, config_filename: Path, *args, **kwargs
+    ):
+        argo = config.get("argo_workflows", {})
+        if argo.get("enabled"):
+            response = Prompt.ask(
+                f"\nDo you want to enable the [green][link={NEBARI_WORKFLOW_CONTROLLER_DOCS}]Nebari Workflow Controller[/link][/green], required for [green][link={ARGO_JUPYTER_SCHEDULER_REPO}]Argo-Jupyter-Scheduler[/link][green]? [Y/n] ",
+                default="Y",
+            )
+            if response.lower() in ["y", "yes", ""]:
+                argo["nebari_workflow_controller"] = {"enabled": True}
+
+        rich.print("\n ⚠️ Deprecation Warnings ⚠️")
+        rich.print(
+            f"-> [green]{self.version}[/green] is the last Nebari version that supports CDS Dashboards"
+        )
+
+        return config
+
+
 __rounded_version__ = ".".join([str(c) for c in rounded_ver_parse(__version__)])
 
 # Manually-added upgrade steps must go above this line