Merge pull request #91 from kbase/dev_jupyterhub

fix hub permission issue

Dockerfile

@@ -73,7 +73,7 @@ RUN chown -R spark_user:spark ${JUPYTERHUB_CONFIG_DIR}
 # Jupyter Hub user home directory
 ENV JUPYTERHUB_USER_HOME=/jupyterhub/users_home
 RUN mkdir -p $JUPYTERHUB_USER_HOME
-RUN chown -R spark_user:spark $JUPYTERHUB_USER_HOME
+RUN chown -R spark_user:spark /jupyterhub
 
 RUN npm install -g configurable-http-proxy
 

src/jupyterhub_config/custom_spawner.py

@@ -1,4 +1,5 @@
 import fcntl
+import grp
 import os
 import pwd
 import subprocess
@@ -31,7 +32,7 @@ def start(self):
         # Ensure the system user exists
         self._ensure_system_user(username, group='jupyterhub')
 
-        # Ensure the user directory exists and has correct permissions
+        # Ensure the user directory exists
         self._ensure_user_directory(user_dir, username)
 
         # Ensure the user's Jupyter directory exists
@@ -47,6 +48,9 @@ def start(self):
         # Configure the notebook directory based on whether the user is an admin
         self._configure_notebook_dir(username, user_dir)
 
+        # Ensure the user's workspace has the correct permissions
+        self._ensure_workspace_permission(user_dir, username)
+
         # Set the command to start the notebook
         env_vars = [f'{key}={value}' for key, value in self.environment.items()]
 
@@ -101,28 +105,11 @@ def _ensure_system_user(self, username: str, group: str = None):
 
     def _ensure_user_directory(self, user_dir: Path, username: str):
         """
-        Ensure the user's home directory exists and is correctly owned and permissioned.
+        Ensure the user's home directory exists.
         """
         if not user_dir.exists():
-
-            self.log.info(f'Getting user info for {username}')
-            try:
-                user_info = pwd.getpwnam(username)
-            except KeyError:
-                raise ValueError(f'System user {username} does not exist')
-            # Get the Jupyter user's UID and GID
-            uid = user_info.pw_uid
-            gid = user_info.pw_gid
-
             self.log.info(f'Creating user directory for {username}')
             user_dir.mkdir(parents=True, exist_ok=True)  # guard against race conditions
-
-            # Change the directory's ownership to the user
-            os.chown(user_dir, uid, gid)
-
-            # Set directory permissions to 750: Owner (rwx), Group (r-x), Others (---)
-            os.chmod(user_dir, 0o750)
-
         else:
             self.log.info(f'Reusing user directory for {username}')
 
@@ -153,7 +140,7 @@ def _ensure_virtual_environment(self, user_env_dir: Path):
         created with the system site-packages included.
         """
         if not user_env_dir.exists():
-            user_env_dir.mkdir(parents=True)
+            user_env_dir.mkdir(parents=True, exist_ok=True)
             self.log.info(f'Creating virtual environment for {self.user.name}')
             try:
                 # Create a virtual environment with system site-packages access
@@ -211,3 +198,25 @@ def _configure_notebook_dir(self, username: str, user_dir: Path):
         else:
             self.log.info(f'Non-admin user detected: {username}. Setting up user-specific workspace.')
             self.notebook_dir = str(user_dir)
+
+    def _ensure_workspace_permission(self, user_dir: Path, username: str):
+        """
+        Ensure the user's workspace has the correct permissions.
+        """
+        try:
+            user_info = pwd.getpwnam(username)
+        except KeyError:
+            raise ValueError(f'System user {username} does not exist')
+        gid = user_info.pw_gid
+        group_name = grp.getgrgid(gid).gr_name
+
+        self.log.info(f'Configuring workspace permissions for {username}')
+        # Change the directory's ownership to the user
+        subprocess.run(['sudo', 'chown', '-R', f'{username}:{group_name}', user_dir], check=True)
+
+        self.log.info(f'Add spark_user to the group of {group_name}')
+        subprocess.run(['sudo', 'usermod', '-aG', group_name, 'spark_user'], check=True)
+
+        # TODO: Set directory permissions to 750 or 700 or switch to use docker spawner
+        # Set directory permissions to 777: Owner (rwx), Group (rwx), Others (rwx)
+        subprocess.run(['sudo', 'chmod', '-R', '777', user_dir], check=True)

test/src/jupyterhub_config/custom_spawner_test.py

@@ -32,8 +32,14 @@ def spawner():
 @patch.object(VirtualEnvSpawner, '_ensure_virtual_environment')
 @patch.object(VirtualEnvSpawner, '_configure_environment')
 @patch.object(VirtualEnvSpawner, '_configure_notebook_dir')
-def test_start(mock_configure_notebook_dir, mock_configure_environment, mock_ensure_virtual_environment,
-               mock_ensure_user_jupyter_directory, mock_ensure_user_directory, mock_ensure_system_user,
+@patch.object(VirtualEnvSpawner, '_ensure_workspace_permission')
+def test_start(mock_ensure_workspace_permission,
+               mock_configure_notebook_dir,
+               mock_configure_environment,
+               mock_ensure_virtual_environment,
+               mock_ensure_user_jupyter_directory,
+               mock_ensure_user_directory,
+               mock_ensure_system_user,
                spawner):
 
     # set spawner.environment (_configure_environment is mocked, so `self.environment` won't be set by the method)
@@ -165,17 +171,9 @@ def test_ensure_system_user_error(mock_run):
     mock_run.assert_has_calls(expected_calls, any_order=False)
 
 
-@patch('pwd.getpwnam')
-@patch('os.chown')
-def test_ensure_user_directory_with_logging(mock_chown, mock_getpwnam, caplog):
+def test_ensure_user_directory_with_logging(caplog):
     username = 'testuser'
 
-    # Mock pwd.getpwnam to return a mock user info
-    mock_user_info = MagicMock()
-    mock_user_info.pw_uid = 1000
-    mock_user_info.pw_gid = 1000
-    mock_getpwnam.return_value = mock_user_info
-
     with tempfile.TemporaryDirectory() as temp_dir:
         user_dir = Path(temp_dir) / username
 
@@ -187,44 +185,10 @@ def test_ensure_user_directory_with_logging(mock_chown, mock_getpwnam, caplog):
         assert user_dir.exists()
         assert user_dir.is_dir()
 
-        # Assert that chown was called with correct parameters
-        mock_chown.assert_called_once_with(user_dir, 1000, 1000)
-
-        # Check directory permissions
-        st = os.stat(user_dir)
-        # Permissions should be 0o750 (rwxr-x---)
-        assert (st.st_mode & 0o777) == 0o750
-
         # Check log messages
-        assert f'Getting user info for {username}' in caplog.text
         assert f'Creating user directory for {username}' in caplog.text
 
-
-@patch('pwd.getpwnam')
-def test_ensure_user_directory_user_not_found(mock_getpwnam, caplog):
-    username = 'nonexistentuser'
-
-    # Mock pwd.getpwnam to raise KeyError (simulating that the user does not exist)
-    mock_getpwnam.side_effect = KeyError
-
-    with tempfile.TemporaryDirectory() as temp_dir:
-        user_dir = Path(temp_dir) / username
-
-        with caplog.at_level(logging.INFO):
-            with pytest.raises(ValueError, match=f'System user {username} does not exist'):
-                spawner = VirtualEnvSpawner()
-                spawner._ensure_user_directory(user_dir, username)
-
-        # Check that the directory was not created
-        assert not user_dir.exists()
-
-        # Check log messages
-        assert f'Getting user info for {username}' in caplog.text
-
-
-@patch('os.chown')
-@patch('os.chmod')
-def test_ensure_user_directory_reuse_existing(mock_chown, mock_chmod, caplog):
+def test_ensure_user_directory_reuse_existing(caplog):
     username = 'testuser'
 
     with tempfile.TemporaryDirectory() as temp_dir:
@@ -237,10 +201,7 @@ def test_ensure_user_directory_reuse_existing(mock_chown, mock_chmod, caplog):
             spawner = VirtualEnvSpawner()
             spawner._ensure_user_directory(user_dir, username)
 
-        # Check that mkdir, chown, and chmod were not called since directory exists
         assert user_dir.exists()
-        mock_chown.assert_not_called()
-        mock_chmod.assert_not_called()
 
         # Check log message
         assert f'Reusing user directory for {username}' in caplog.text