Add compatibility prompt and notes for shared group mounting (#2739)

src/_nebari/keycloak.py

@@ -81,27 +81,16 @@ def list_users(keycloak_admin: keycloak.KeycloakAdmin):
         )
 
 
-def get_keycloak_admin_from_config(config: schema.Main):
-    keycloak_server_url = os.environ.get(
-        "KEYCLOAK_SERVER_URL", f"https://{config.domain}/auth/"
-    )
-
-    keycloak_username = os.environ.get("KEYCLOAK_ADMIN_USERNAME", "root")
-    keycloak_password = os.environ.get(
-        "KEYCLOAK_ADMIN_PASSWORD", config.security.keycloak.initial_root_password
-    )
-
-    should_verify_tls = config.certificate.type != CertificateEnum.selfsigned
-
+def get_keycloak_admin(server_url, username, password, verify=False):
     try:
         keycloak_admin = keycloak.KeycloakAdmin(
-            server_url=keycloak_server_url,
-            username=keycloak_username,
-            password=keycloak_password,
+            server_url=server_url,
+            username=username,
+            password=password,
             realm_name=os.environ.get("KEYCLOAK_REALM", "nebari"),
             user_realm_name="master",
             auto_refresh_token=("get", "put", "post", "delete"),
-            verify=should_verify_tls,
+            verify=verify,
         )
     except (
         keycloak.exceptions.KeycloakConnectionError,
@@ -112,6 +101,26 @@ def get_keycloak_admin_from_config(config: schema.Main):
     return keycloak_admin
 
 
+def get_keycloak_admin_from_config(config: schema.Main):
+    keycloak_server_url = os.environ.get(
+        "KEYCLOAK_SERVER_URL", f"https://{config.domain}/auth/"
+    )
+
+    keycloak_username = os.environ.get("KEYCLOAK_ADMIN_USERNAME", "root")
+    keycloak_password = os.environ.get(
+        "KEYCLOAK_ADMIN_PASSWORD", config.security.keycloak.initial_root_password
+    )
+
+    should_verify_tls = config.certificate.type != CertificateEnum.selfsigned
+
+    return get_keycloak_admin(
+        server_url=keycloak_server_url,
+        username=keycloak_username,
+        password=keycloak_password,
+        verify=should_verify_tls,
+    )
+
+
 def keycloak_rest_api_call(config: schema.Main = None, request: str = None):
     """Communicate directly with the Keycloak REST API by passing it a request"""
     keycloak_server_url = os.environ.get(

src/_nebari/upgrade.py

@@ -24,6 +24,7 @@
 from typing_extensions import override
 
 from _nebari.config import backup_configuration
+from _nebari.keycloak import get_keycloak_admin
 from _nebari.stages.infrastructure import (
     provider_enum_default_node_groups_map,
     provider_enum_name_map,
@@ -1235,6 +1236,95 @@ def _version_specific_upgrade(
             )
             rich.print("")
 
+        rich.print("\n ⚠️ Upgrade Warning ⚠️")
+
+        text = textwrap.dedent(
+            """
+            Please ensure no users are currently logged in prior to deploying this
+            update.
+
+            Nebari [green]2024.9.1[/green] introduces changes to how group
+            directories are mounted in JupyterLab pods.
+
+            Previously, every Keycloak group in the Nebari realm automatically created a
+            shared directory at ~/shared/<group-name>, accessible to all group members
+            in their JupyterLab pods.
+
+            Starting with Nebari [green]2024.9.1[/green], only groups assigned the
+            JupyterHub client role [magenta]allow-group-directory-creation[/magenta] will have their
+            directories mounted.
+
+            By default, the admin, analyst, and developer groups will have this
+            role assigned during the upgrade. For other groups, you'll now need to
+            assign this role manually in the Keycloak UI to have their directories
+            mounted.
+
+            For more details check our [green][link=https://www.nebari.dev/docs/references/release/]release notes[/link][/green].
+            """
+        )
+        rich.print(text)
+        keycloak_admin = None
+
+        # Prompt the user for role assignment (if yes, transforms the response into bool)
+        assign_roles = (
+            Prompt.ask(
+                "[bold]Would you like Nebari to assign the corresponding role to all of your current groups automatically?[/bold]",
+                choices=["y", "N"],
+                default="N",
+            ).lower()
+            == "y"
+        )
+
+        if assign_roles:
+            # In case this is done with a local deployment
+            import urllib3
+
+            urllib3.disable_warnings()
+
+            keycloak_admin = get_keycloak_admin(
+                server_url=f"https://{config['domain']}/auth/",
+                username="root",
+                password=config["security"]["keycloak"]["initial_root_password"],
+            )
+
+            # Proceed with updating group permissions
+            client_id = keycloak_admin.get_client_id("jupyterhub")
+            role_name = "allow-group-directory-creation-role"
+            role_id = keycloak_admin.get_client_role_id(
+                client_id=client_id, role_name=role_name
+            )
+            role_representation = keycloak_admin.get_role_by_id(role_id=role_id)
+
+            # Fetch all groups and groups with the role
+            all_groups = keycloak_admin.get_groups()
+            groups_with_role = keycloak_admin.get_client_role_groups(
+                client_id=client_id, role_name=role_name
+            )
+            groups_with_role_ids = {group["id"] for group in groups_with_role}
+
+            # Identify groups without the role
+            groups_without_role = [
+                group for group in all_groups if group["id"] not in groups_with_role_ids
+            ]
+
+            if groups_without_role:
+                group_names = ", ".join(group["name"] for group in groups_without_role)
+                rich.print(
+                    f"\n[bold]Updating the following groups with the required permissions:[/bold] {group_names}\n"
+                )
+                for group in groups_without_role:
+                    keycloak_admin.assign_group_client_roles(
+                        group_id=group["id"],
+                        client_id=client_id,
+                        roles=[role_representation],
+                    )
+                rich.print(
+                    "\n[green]Group permissions have been updated successfully.[/green]"
+                )
+            else:
+                rich.print(
+                    "\n[green]All groups already have the required permissions.[/green]"
+                )
         return config
 
 

tests/tests_unit/test_upgrade.py

@@ -67,6 +67,11 @@ def mock_input(prompt, **kwargs):
             == "Have you backed up your custom dashboards (if necessary), deleted the prometheus-node-exporter daemonset and updated the kube-prometheus-stack CRDs?"
         ):
             return "y"
+        elif (
+            prompt
+            == "[bold]Would you like Nebari to assign the corresponding role to all of your current groups automatically?[/bold]"
+        ):
+            return "N"
         # All other prompts will be answered with "y"
         else:
             return "y"