diff --git a/pyproject.toml b/pyproject.toml index 97f26b32f..a5d6b26e5 100644 --- a/pyproject.toml +++ b/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" diff --git a/src/drunc/apps/session_manager_shell.py b/src/drunc/apps/session_manager_shell.py new file mode 100644 index 000000000..32a3b302d --- /dev/null +++ b/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() diff --git a/src/drunc/session_manager/interface/commands.py b/src/drunc/session_manager/interface/commands.py new file mode 100644 index 000000000..027ef10cd --- /dev/null +++ b/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) diff --git a/src/drunc/session_manager/interface/context.py b/src/drunc/session_manager/interface/context.py new file mode 100644 index 000000000..6391c6246 --- /dev/null +++ b/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() diff --git a/src/drunc/session_manager/interface/session_manager.py b/src/drunc/session_manager/interface/session_manager.py index 0e7a79849..ce00103cd 100644 --- a/src/drunc/session_manager/interface/session_manager.py +++ b/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. diff --git a/src/drunc/session_manager/interface/shell.py b/src/drunc/session_manager/interface/shell.py new file mode 100644 index 000000000..77a66dd0a --- /dev/null +++ b/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")