From 101163cf1df251681b86bed0dbe09eb5fc921f96 Mon Sep 17 00:00:00 2001
From: Tianhao-Gu <tgu@anl.gov>
Date: Tue, 1 Oct 2024 14:38:01 -0500
Subject: [PATCH] add shared dir to user's favorite list

---
 Dockerfile                                    |  5 +++++
 .../custom_docker_spawner.py                  | 20 +++++++++++--------
 src/jupyterhub_config/hub_singleuser.py       |  2 +-
 3 files changed, 18 insertions(+), 9 deletions(-)

diff --git a/Dockerfile b/Dockerfile
index c9ded25..3009f22 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -98,6 +98,11 @@ ENV CDM_SHARED_DIR=/cdm_shared_workspace
 RUN mkdir -p ${CDM_SHARED_DIR} && chmod -R 777 ${CDM_SHARED_DIR}
 RUN chown -R spark_user:spark $CDM_SHARED_DIR
 
+# TODO: Config through a config file or DB as the number of groups increases.
+ENV KBASE_GROUP_SHARED_DIR=$CDM_SHARED_DIR/kbase_group_shared
+RUN mkdir -p ${KBASE_GROUP_SHARED_DIR} && chmod -R 777 ${KBASE_GROUP_SHARED_DIR}
+RUN chown -R spark_user:spark $KBASE_GROUP_SHARED_DIR
+
 # Set a directory for hosting Hive metastore files - defined in config/hive-site-template.xml
 ENV HIVE_METASTORE_DIR=$CDM_SHARED_DIR/hive_metastore
 RUN mkdir -p ${HIVE_METASTORE_DIR}
diff --git a/src/jupyterhub_config/custom_docker_spawner.py b/src/jupyterhub_config/custom_docker_spawner.py
index ad936c3..ef57dc1 100644
--- a/src/jupyterhub_config/custom_docker_spawner.py
+++ b/src/jupyterhub_config/custom_docker_spawner.py
@@ -34,9 +34,8 @@ def start(self):
         # Ensure the user's volume is correctly mounted in the container
         self._ensure_user_volume()
 
-        # Add the user's home directory to JupyterLab favorites
-        # TODO: include shared group directories in favorites
-        self._add_favorite_dir(user_dir)
+        # Add the user's home and group shared directory to Jupyterhub favorites
+        self._add_favorite_dir(user_dir, favorites={Path(os.environ['KBASE_GROUP_SHARED_DIR'])})
 
         return super().start()
 
@@ -138,7 +137,9 @@ def _configure_notebook_dir(self, username: str, user_dir: Path):
             self.notebook_dir = str('/')
         else:
             self.log.info(f'Non-admin user detected: {username}. Setting up user-specific workspace.')
-            self.notebook_dir = str(user_dir)
+            # TODO: It appears that notebook_dir must be the parent of the favorites directory - investigate if it's possible to set notebook_dir to user_dir
+            # self.notebook_dir = str(user_dir)
+            self.notebook_dir = str('/')
 
     def _is_rw_minio_user(self):
         """
@@ -160,6 +161,7 @@ def _ensure_user_volume(self):
 
         cdm_shared_dir = Path(os.environ['CDM_SHARED_DIR'])  # Legacy data volume from JupyterLab
         hive_metastore_dir = Path(os.environ['HIVE_METASTORE_DIR'])  # within cdm_shared_dir
+        kbase_shared_dir = Path(os.environ['KBASE_GROUP_SHARED_DIR'])  # within cdm_shared_dir
 
         if self.user.admin:
             self.log.info(f'Admin user detected: {self.user.name}. Setting up admin mount points.')
@@ -174,7 +176,8 @@ def _ensure_user_volume(self):
             self.volumes.update({
                 f'{mount_base_dir}/{hive_metastore_dir}': {'bind': f'{hive_metastore_dir}', 'mode': access_mode},
                 # User specific home directory
-                f'{mount_base_dir}/{user_home_dir}/{self.user.name}': f'{user_home_dir}/{self.user.name}'
+                f'{mount_base_dir}/{user_home_dir}/{self.user.name}': f'{user_home_dir}/{self.user.name}',
+                f'{mount_base_dir}/{kbase_shared_dir}': f'{kbase_shared_dir}',
             })
 
     def _add_favorite_dir(self, user_dir: Path, favorites: set[Path] = None):
@@ -211,8 +214,9 @@ def _add_favorite_dir(self, user_dir: Path, favorites: set[Path] = None):
                 if not fav.is_dir():
                     raise ValueError(f"Favorite {fav} is not a directory or does not exist")
 
-                root_str = str(fav)
-                path_str = ""
+                # same approach used by NERSC JupyterHub
+                root_str = "/"
+                path_str = str(fav.relative_to(root_str))
 
                 if (root_str, path_str) not in existing_fav_set:
                     exist_favorites["favorites"].append({
@@ -220,7 +224,7 @@ def _add_favorite_dir(self, user_dir: Path, favorites: set[Path] = None):
                         "path": path_str,
                         "contentType": "directory",
                         "iconLabel": "ui-components:folder",
-                        "name": "$HOME" if root_str == str(user_dir) else fav.name,
+                        "name": "$HOME" if str(fav) == str(user_dir) else fav.name,
                     })
 
             with open(jupyterlab_favorites_path, 'w') as f:
diff --git a/src/jupyterhub_config/hub_singleuser.py b/src/jupyterhub_config/hub_singleuser.py
index 5d97b1c..2f978c5 100644
--- a/src/jupyterhub_config/hub_singleuser.py
+++ b/src/jupyterhub_config/hub_singleuser.py
@@ -15,7 +15,7 @@ class SingleUserEnvManager:
 
     def __init__(self,
                  username,
-                 groupname: str = 'jupyterhub', ):
+                 groupname: str = 'kbase', ):
         """
         Initializes the environment manager for the specified user.
         :param username: The name of the user for whom the environment is set up.