Jamesturner246/sm load session

src/drunc/session_manager/interface/commands.py

@@ -23,3 +23,27 @@ def list_all_configs(ctx) -> None:
     """List all available configurations."""
     response = ctx.obj.get_driver("session_manager").list_all_configs()
     click.echo(response)
+
+
+@click.command("load_session")
+@click.option(
+    "--session-file",
+    "-f",
+    type=str,
+    required=True,
+    help="The file containing the session to load.",
+)
+@click.option(
+    "--session-id",
+    "-i",
+    type=str,
+    required=True,
+    help="The ID of the session to load.",
+)
+@click.pass_context
+def load_session(ctx, session_file: str, session_id: str) -> None:
+    """Load a session based on the provided configuration file path and ID."""
+    response = ctx.obj.get_driver("session_manager").load_session(
+        session_file, session_id
+    )
+    click.echo(response)

src/drunc/session_manager/interface/shell.py

@@ -8,6 +8,7 @@
     describe,
     list_all_configs,
     list_all_sessions,
+    load_session,
 )
 from drunc.utils.utils import (
     CONTEXT_SETTINGS,
@@ -42,6 +43,7 @@ def cleanup() -> None:
     ctx.command.add_command(describe, "describe")
     ctx.command.add_command(list_all_sessions, "list_all_sessions")
     ctx.command.add_command(list_all_configs, "list_all_configs")
+    ctx.command.add_command(load_session, "load_session")
     ctx.call_on_close(cleanup)
 
     log.info("Starting session manager shell")

src/drunc/session_manager/session_manager.py

@@ -3,6 +3,7 @@
 import abc
 from os import getenv
 from pathlib import Path
+from threading import Lock
 
 from conffwk import Configuration
 from druncschema.description_pb2 import CommandDescription, Description
@@ -15,6 +16,8 @@
     AllActiveSessions,
     AllConfigKeys,
     ConfigKey,
+    LoadSessionRequest,
+    LoadSessionResponse,
 )
 from druncschema.session_manager_pb2_grpc import SessionManagerServicer
 from grpc import ServicerContext
@@ -40,12 +43,14 @@ def __init__(self, name: str, configuration: SessionManagerConfHandler):
         """
         super().__init__()
 
-        self.log = get_logger("session_manager", rich_handler=True)
-        self.log.debug(pid_info_str())
-        self.log.debug("Initialized SessionManager")
-
         self.name = name
         self.configuration = configuration
+        self._active_sessions: dict[str, ActiveSession] = {}
+        self._active_sessions_lock = Lock()
+
+        self.log = get_logger("session_manager", rich_handler=True)
+        self.log.debug(pid_info_str())
+        self.log.debug("Initialised session manager")
 
     def describe(self, request: Request, context: ServicerContext) -> Description:
         """Respond with a description of this session manager service.
@@ -78,6 +83,12 @@ def describe(self, request: Request, context: ServicerContext) -> Description:
                 help="List all available configurations.",
                 return_type="session_manager_pb2.AllConfigKeys",
             ),
+            CommandDescription(
+                name="load_session",
+                data_type=["None"],
+                help="Load a session based on the provided configuration key.",
+                return_type="session_manager_pb2.LoadSessionResponse",
+            ),
         ]
 
         return Description(
@@ -102,21 +113,13 @@ def list_all_sessions(
         """
         self.log.debug(f"{self.name} running list_all_sessions")
 
-        dummy_config = ConfigKey(
-            file="dummy_config_file",
-            session_id="dummy_config_session_id",
-        )
-
-        dummy_session = ActiveSession(
-            name="dummy_session",
-            user="dummy_user",
-            config_key=dummy_config,
-        )
+        with self._active_sessions_lock:
+            active_sessions = list(self._active_sessions.values())
 
         return AllActiveSessions(
             name=self.name,
             token=None,
-            active_sessions=[dummy_session],
+            active_sessions=active_sessions,
             flag=ResponseFlag.EXECUTED_SUCCESSFULLY,
         )
 
@@ -188,3 +191,40 @@ def list_all_configs(
             config_keys=configs,
             flag=ResponseFlag.EXECUTED_SUCCESSFULLY,
         )
+
+    def load_session(
+        self, request: LoadSessionRequest, context: ServicerContext
+    ) -> LoadSessionResponse:
+        """Load a session based on the provided configuration key.
+
+        Args:
+            request: The incoming request containing the configuration key.
+            context: The gRPC context (not used).
+
+        Returns:
+            LoadSessionResponse: A response containing loaded session metadata.
+        """
+        self.log.debug(f"{self.name} running load_session")
+
+        file = request.config_key.file
+        session_id = request.config_key.session_id
+        self.log.info(f"Loading session: '{session_id}' from file: '{file}'")
+
+        session = ActiveSession(
+            name="session_name", user="session_user", config_key=request.config_key
+        )
+
+        with self._active_sessions_lock:
+            if session_id in self._active_sessions:
+                raise DruncSetupException(
+                    message="Unable to load session",
+                    details=f"Session '{session_id}' already exists.",
+                )
+
+            # TODO: Implement the actual session loading logic here.
+            # TODO: Set user name and session name (different from session ID).
+            # TODO: Store connection URL in `ActiveSession`.
+
+            self._active_sessions[session_id] = session
+
+            return LoadSessionResponse(session=session)

src/drunc/session_manager/session_manager_driver.py

@@ -3,7 +3,13 @@
 import grpc
 from druncschema.description_pb2 import Description
 from druncschema.request_response_pb2 import Request
-from druncschema.session_manager_pb2 import AllActiveSessions, AllConfigKeys
+from druncschema.session_manager_pb2 import (
+    AllActiveSessions,
+    AllConfigKeys,
+    ConfigKey,
+    LoadSessionRequest,
+    LoadSessionResponse,
+)
 from druncschema.session_manager_pb2_grpc import SessionManagerStub
 from druncschema.token_pb2 import Token
 
@@ -50,10 +56,9 @@ def describe(self, timeout: int | float = 60) -> Description:
         Returns:
             A response containing the description of the service.
         """
+        self.log.info(f"Sending describe request to session manager at {self.address}")
+
         request = Request(token=copy_token(self.token))
-        self.log.info(
-            f"Sending describe request to session manager at {self.address} with timeout {timeout}s"
-        )
 
         response: Description = self.stub.describe(request, timeout=timeout)
 
@@ -68,11 +73,12 @@ def list_all_sessions(self, timeout: int | float = 60) -> AllActiveSessions:
         Returns:
             A response containing a list of all active sessions.
         """
-        request = Request(token=copy_token(self.token))
         self.log.info(
-            f"Sending list_all_sessions request to session manager at {self.address} with timeout {timeout}s"
+            f"Sending list_all_sessions request to session manager at {self.address}"
         )
 
+        request = Request(token=copy_token(self.token))
+
         response: AllActiveSessions = self.stub.list_all_sessions(
             request, timeout=timeout
         )
@@ -88,11 +94,40 @@ def list_all_configs(self, timeout: int | float = 60) -> AllConfigKeys:
         Returns:
             A response containing all available configuration keys.
         """
-        request = Request(token=copy_token(self.token))
         self.log.info(
-            f"Sending list_all_configs request to session manager at {self.address} with timeout {timeout}s"
+            f"Sending list_all_configs request to session manager at {self.address}"
         )
 
+        request = Request(token=copy_token(self.token))
+
         response: AllConfigKeys = self.stub.list_all_configs(request, timeout=timeout)
 
         return response
+
+    def load_session(
+        self, file: str, session_id: str, timeout: int | float = 60
+    ) -> LoadSessionResponse:
+        """Load a session based on the provided configuration key.
+
+        Args:
+            file: The file containing the session to load.
+            session_id: The ID of the session to load.
+            timeout: The timeout for the gRPC call in seconds.
+
+        Returns:
+            LoadSessionResponse: A response containing loaded session metadata.
+        """
+        self.log.info(
+            f"Sending load_session request to session manager at {self.address}"
+        )
+
+        request = LoadSessionRequest(
+            config_key=ConfigKey(
+                file=file,
+                session_id=session_id,
+            )
+        )
+
+        response: LoadSessionResponse = self.stub.load_session(request, timeout=timeout)
+
+        return response

tests/session_manager/conftest.py

@@ -289,6 +289,12 @@ def command_description_list():
             help="List all available configurations.",
             return_type="session_manager_pb2.AllConfigKeys",
         ),
+        CommandDescription(
+            name="load_session",
+            data_type=["None"],
+            help="Load a session based on the provided configuration key.",
+            return_type="session_manager_pb2.LoadSessionResponse",
+        ),
     ]
 
 

tests/session_manager/test_session_manager_servicer.py

@@ -9,6 +9,8 @@
     AllActiveSessions,
     AllConfigKeys,
     ConfigKey,
+    LoadSessionRequest,
+    LoadSessionResponse,
 )
 
 from drunc.exceptions import DruncSetupException
@@ -18,7 +20,7 @@ def test_describe(
     session_manager, mock_request, mock_context, command_description_list, mock_logger
 ):
     response = session_manager.describe(mock_request, mock_context)
-    mock_logger.debug.assert_any_call("Initialized SessionManager")
+    mock_logger.debug.assert_any_call("Initialised session manager")
 
     assert isinstance(response, Description)
     assert response.name == "dummy_name"
@@ -35,6 +37,7 @@ def test_list_all_sessions(session_manager, mock_request, mock_context, mock_log
         name="dummy_session", user="dummy_user", config_key=mock_config
     )
 
+    session_manager._active_sessions = {mock_config.session_id: mock_session}
     response = session_manager.list_all_sessions(mock_request, mock_context)
     mock_logger.debug.assert_any_call(f"{response.name} running list_all_sessions")
 
@@ -158,3 +161,48 @@ def test_list_all_configs_dals_missing(
             mock_logger.error.assert_any_call(
                 "Failed to get DALs from mock_file_1.data.xml: DALs missing or invalid"
             )
+
+
+def test_load_session(session_manager, mock_context, mock_logger):
+    """
+    Test loading a session with a given configuration key.
+    """
+    session_file = "dummy_config_file"
+    session_id = "dummy_config_session_id"
+    session_name = "session_name"
+    session_user = "session_user"
+
+    mock_config = ConfigKey(file=session_file, session_id=session_id)
+    mock_request = LoadSessionRequest(config_key=mock_config)
+
+    response = session_manager.load_session(mock_request, mock_context)
+    session = session_manager._active_sessions[session_id]
+    mock_logger.debug.assert_any_call(f"{session_manager.name} running load_session")
+
+    assert session.name == session_name
+    assert session.user == session_user
+    assert session.config_key == mock_config
+
+    assert isinstance(response, LoadSessionResponse)
+    assert response.session.name == session_name
+    assert response.session.user == session_user
+    assert response.session.config_key == mock_config
+
+
+def test_load_session_duplicate_id(session_manager, mock_context):
+    """
+    Test loading a session with a duplicate session ID.
+    """
+    session_file = "dummy_config_file"
+    session_id = "dummy_config_session_id"
+
+    mock_config = ConfigKey(file=session_file, session_id=session_id)
+    mock_request = LoadSessionRequest(config_key=mock_config)
+
+    session_manager.load_session(mock_request, mock_context)
+    with pytest.raises(DruncSetupException) as excinfo:
+        session_manager.load_session(mock_request, mock_context)
+
+    assert "Unable to load session" in str(excinfo.value)
+    assert len(session_manager._active_sessions) == 1
+    assert session_manager._active_sessions[session_id].config_key == mock_config