Add Users and Groups permissions filters to DaskWorker profiles

src/_nebari/stages/kubernetes_services/__init__.py

@@ -90,14 +90,10 @@ class Config:
         extra = "allow"
 
 
-class JupyterLabProfile(schema.Base):
+class ProfileAccess(schema.Base):
     access: AccessEnum = AccessEnum.all
-    display_name: str
-    description: str
-    default: bool = False
     users: typing.Optional[typing.List[str]]
     groups: typing.Optional[typing.List[str]]
-    kubespawner_override: typing.Optional[KubeSpawner]
 
     @pydantic.root_validator
     def only_yaml_can_have_groups_and_users(cls, values):
@@ -112,7 +108,14 @@ def only_yaml_can_have_groups_and_users(cls, values):
         return values
 
 
-class DaskWorkerProfile(schema.Base):
+class JupyterLabProfile(ProfileAccess):
+    display_name: str
+    description: str
+    default: bool = False
+    kubespawner_override: typing.Optional[KubeSpawner]
+
+
+class DaskWorkerProfile(ProfileAccess):
     worker_cores_limit: int
     worker_cores: int
     worker_memory_limit: str

src/_nebari/stages/kubernetes_services/template/modules/kubernetes/services/dask-gateway/files/gateway_config.py

@@ -81,6 +81,8 @@ async def authenticate(self, request):
 
         user.admin = "dask_gateway_admin" in data["roles"]
         user.groups = [Path(group).name for group in data["groups"]]
+        user.keycloak_profile_names = data.get("dask_profiles", [])
+
         return user
 
 
@@ -227,21 +229,61 @@ def base_username_mount(username, uid=1000, gid=100):
     }
 
 
+def _sanitize_permissions(profile):
+    keys_to_remove = ["groups", "users", "access"]
+
+    for key in keys_to_remove:
+        profile.pop(key, None)
+
+    return profile
+
+
 def worker_profile(options, user):
     namespace, name = options.conda_environment.split("/")
+
     return functools.reduce(
         deep_merge,
         [
             base_node_group(options),
             base_conda_store_mounts(namespace, name),
             base_username_mount(user.name),
-            config["profiles"][options.profile],
+            _sanitize_permissions(config["profiles"][options.profile]),
             {"environment": {**options.environment_vars}},
         ],
         {},
     )
 
 
+def _filter_profiles(user, profiles):
+    """
+    Filter access to profiles based on user's groups and username
+    """
+
+    def has_group_access(profile):
+        return not profile.get("groups") or set(user.groups).intersection(
+            profile["groups"]
+        )
+
+    def has_user_access(profile):
+        return not profile.get("users") or user.name in profile["users"]
+
+    user_profiles = list(profiles.keys())
+
+    for name in list(user_profiles):
+        profile = profiles[name]
+        access_type = profile.get("access", "all")
+
+        if access_type == "yaml":
+            if not (has_group_access(profile) and has_user_access(profile)):
+                user_profiles.remove(name)
+        elif access_type == "keycloak":
+            # Keycloak mapper should provide the 'daskworker_profiles' attribute from groups/user
+            if name not in user.keycloak_profilenames:
+                user_profiles.remove(name)
+
+    return user_profiles
+
+
 def user_options(user):
     default_namespace = config["default-conda-store-namespace"]
     allowed_namespaces = set(
@@ -253,6 +295,8 @@ def user_options(user):
             continue
         conda_environments.append(f"{namespace}/{namespace}-{name}")
 
+    user_profiles = _filter_profiles(user, config["profiles"])
+
     args = []
     if conda_environments:
         args += [
@@ -267,8 +311,8 @@ def user_options(user):
         args += [
             Select(
                 "profile",
-                list(config["profiles"].keys()),
-                default=list(config["profiles"].keys())[0],
+                user_profiles,
+                default=user_profiles[0],
                 label="Cluster Profile",
             )
         ]

src/_nebari/stages/kubernetes_services/template/modules/kubernetes/services/jupyterhub/main.tf

@@ -283,6 +283,7 @@ module "jupyterhub-openid-client" {
     var.jupyterhub-logout-redirect-url
   ]
   jupyterlab_profiles_mapper = true
+  daskworker_profiles_mapper = true
 }
 
 

src/_nebari/stages/kubernetes_services/template/modules/kubernetes/services/keycloak-client/main.tf

@@ -62,6 +62,23 @@ resource "keycloak_openid_user_attribute_protocol_mapper" "jupyterlab_profiles"
   aggregate_attributes = true
 }
 
+resource "keycloak_openid_user_attribute_protocol_mapper" "daskworker_profiles" {
+  count = var.daskworker_profiles_mapper ? 1 : 0
+
+  realm_id   = var.realm_id
+  client_id  = keycloak_openid_client.main.id
+  name       = "daskworker_profiles_mapper"
+  claim_name = "daskworker_profiles"
+
+  add_to_id_token     = true
+  add_to_access_token = true
+  add_to_userinfo     = true
+
+  user_attribute       = "daskworker_profiles"
+  multivalued          = true
+  aggregate_attributes = true
+}
+
 resource "keycloak_role" "main" {
   for_each = toset(flatten(values(var.role_mapping)))
 

src/_nebari/stages/kubernetes_services/template/modules/kubernetes/services/keycloak-client/variables.tf

@@ -33,3 +33,9 @@ variable "jupyterlab_profiles_mapper" {
   type        = bool
   default     = false
 }
+
+variable "daskworker_profiles_mapper" {
+  description = "Create a mapper for daskworker_profiles group/user attributes"
+  type        = bool
+  default     = false
+}