Session manager shell

pyproject.toml

@@ -47,6 +47,7 @@ drunc-fsm-tests = "drunc.tests.fsm:main"
 drunc-process-manager = "drunc.apps.pm:main"
 drunc-process-manager-shell = "drunc.apps.pm_shell:main"
 drunc-session-manager = "drunc.apps.session_manager:main"
+drunc-session-manager-shell = "drunc.apps.session_manager_shell:main"
 drunc-setup-ssh-config = "drunc.apps.ssh_configurator:main"
 drunc-ssh-doctor = "drunc.apps.ssh_doctor:main"
 drunc-unified-shell = "drunc.apps.unified_shell:main"

src/drunc/apps/session_manager_shell.py

@@ -0,0 +1,19 @@
+from drunc.session_manager.interface.context import SessionManagerContext
+from drunc.session_manager.interface.shell import session_manager_shell
+from drunc.utils.utils import get_logger, get_root_logger
+
+
+def main():
+    context = SessionManagerContext()
+    try:
+        session_manager_shell(obj=context)
+    except Exception as e:
+        get_root_logger("INFO")
+        log = get_logger("session_manager", rich_handler=True)
+        log.error("[red bold]:fire::fire: Exception thrown :fire::fire:")
+        log.exception(e)
+        exit(1)
+
+
+if __name__ == "__main__":
+    main()

src/drunc/session_manager/interface/commands.py

@@ -0,0 +1,25 @@
+import click
+
+
+@click.command("describe")
+@click.pass_context
+def describe(ctx) -> None:
+    """List the methods exposed by this endpoint."""
+    response = ctx.obj.get_driver("session_manager").describe()
+    click.echo(response)
+
+
+@click.command("list_all_sessions")
+@click.pass_context
+def list_all_sessions(ctx) -> None:
+    """List all active sessions."""
+    response = ctx.obj.get_driver("session_manager").list_all_sessions()
+    click.echo(response)
+
+
+@click.command("list_all_configs")
+@click.pass_context
+def list_all_configs(ctx) -> None:
+    """List all available configurations."""
+    response = ctx.obj.get_driver("session_manager").list_all_configs()
+    click.echo(response)

src/drunc/session_manager/interface/context.py

@@ -0,0 +1,41 @@
+from druncschema.token_pb2 import Token
+
+from drunc.session_manager.session_manager_driver import SessionManagerDriver
+from drunc.utils.shell_utils import (
+    ShellContext,
+    create_dummy_token_from_uname,
+)
+from drunc.utils.utils import resolve_localhost_to_hostname
+
+
+class SessionManagerContext(ShellContext):
+    shell_id = "session_manager_shell"
+
+    def __init__(self, *args, **kwargs):
+        self.status_receiver = None
+        super().__init__(*args, **kwargs)
+
+    def reset(self, address: str = "", **kwargs):
+        self.address = resolve_localhost_to_hostname(address)
+        super()._reset(
+            name="session_manager_context",
+            token_args={},
+            driver_args={},
+        )
+
+    def create_drivers(self, **kwargs) -> dict[str, object]:
+        if not self.address:
+            return {}
+        return {
+            "session_manager": SessionManagerDriver(
+                self.address,
+                self._token,
+            )
+        }
+
+    def create_token(self, **kwargs) -> Token:
+        return create_dummy_token_from_uname()
+
+    def terminate(self):
+        if self.status_receiver:
+            self.status_receiver.stop()

src/drunc/session_manager/interface/session_manager.py

@@ -5,6 +5,7 @@
 
 import click
 import grpc
+from daqpytools.logging import logging_log_levels
 from druncschema.session_manager_pb2_grpc import add_SessionManagerServicer_to_server
 
 from drunc.grpc_settings import (
@@ -38,16 +39,22 @@ def serve(session_manager: SessionManager, address: str) -> None:
 
 
 @click.command()
-def session_manager_cli() -> None:
+@click.option(
+    "-l",
+    "--log-level",
+    type=click.Choice(logging_log_levels, case_sensitive=False),
+    default="INFO",
+    help="Set the log level (default is 'INFO').",
+)
+def session_manager_cli(log_level: str) -> None:
     """CLI interface for the Drunc session manager.
 
     This command starts the session manager service, which allows clients to manage
     and interact with drunc sessions.
     """
-    app_name = "session_manager"
-    log_level = "DEBUG"
-
     get_root_logger(log_level)
+
+    app_name = "session_manager"
     logger = get_logger(app_name, rich_handler=True)
 
     # Load the configuration for the session manager.

src/drunc/session_manager/interface/shell.py

@@ -0,0 +1,47 @@
+import os
+
+import click
+import click_shell
+from daqpytools.logging import logging_log_levels
+
+from drunc.session_manager.interface.commands import (
+    describe,
+    list_all_configs,
+    list_all_sessions,
+)
+from drunc.utils.utils import (
+    CONTEXT_SETTINGS,
+    get_logger,
+    get_root_logger,
+)
+
+
+@click_shell.shell(
+    prompt="drunc-session-manager > ",
+    chain=True,
+    context_settings=CONTEXT_SETTINGS,
+    hist_file=os.path.expanduser("~/.drunc-sm-shell.history"),
+)
+@click.option(
+    "-l",
+    "--log-level",
+    type=click.Choice(logging_log_levels.keys(), case_sensitive=False),
+    default="INFO",
+    help="Set the log level (default is 'INFO').",
+)
+@click.argument("session-manager-address", type=str)
+@click.pass_context
+def session_manager_shell(ctx, session_manager_address: str, log_level: str) -> None:
+    get_root_logger(log_level)
+    log = get_logger("session_manager.shell", rich_handler=True)
+
+    def cleanup() -> None:
+        log.info("Exiting session manager shell")
+
+    ctx.obj.reset(address=session_manager_address)
+    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.call_on_close(cleanup)
+
+    log.info("Starting session manager shell")