WIP

Author: PawelPlesniak

src/drunc/process_manager/ssh_process_manager.py

@@ -448,6 +448,10 @@ def _ps_impl(self, query: ProcessQuery) -> ProcessInstanceList:
                 boot_request_dict=self.boot_request,
                 order_by="random",
             )
+            if query.session:
+                self.log.warning(
+                    f"{self.name} found {len(process_uuids)} processes matching {query} for ps"
+                )
 
             # Iterate through all processes matching the query
             for proc_uuid in process_uuids:
@@ -495,13 +499,24 @@ def _ps_impl(self, query: ProcessQuery) -> ProcessInstanceList:
                 else:
                     pi.remote_pid = remote_pid_result.reason
                 ret += [pi]
-
-            return ProcessInstanceList(
+            ret_fmt = ProcessInstanceList(
                 name=self.name,
                 token=None,
                 values=ret,
                 flag=ResponseFlag.EXECUTED_SUCCESSFULLY,
             )
+            if query.session:
+                self.log.critical(
+                    f"{self.name} returning {len(ret)} processes from ps query {query}"
+                )
+                self.log.critical(ret_fmt)
+                self.log.critical(f"TEST: {ret_fmt=}")
+            else:
+                self.log.warning(
+                    f"{self.name} returning {len(ret)} processes from ps query {query}"
+                )
+                self.log.warning(ret)
+            return ret_fmt
 
     def _boot_impl(self, boot_request: BootRequest) -> ProcessInstanceList:
         self.log.debug(f"{self.name} running boot command")

src/drunc/unified_shell/commands.py

@@ -7,6 +7,7 @@
 from drunc.controller.interface.shell_utils import controller_setup
 from drunc.exceptions import DruncSetupException
 from drunc.process_manager.interface.context import ProcessManagerContext
+from drunc.process_manager.utils import tabulate_process_instance_list
 from drunc.unified_shell.context import UnifiedShellMode
 from drunc.utils.shell_utils import InterruptedCommand
 from drunc.utils.utils import get_logger
@@ -45,12 +46,12 @@ def boot(
         override_logs_boot = obj.override_logs
     else:
         override_logs_boot = override_logs
-    if len(processes.values) > 0:
-        log.error(
-            f"Cannot boot: session {session_name} already has {len(processes.values)} processes running. "
-            "Please terminate the existing session first."
-        )
-        return
+    # if len(processes.values) > 0:
+    #     log.error(
+    #         f"Cannot boot: session {session_name} already has {len(processes.values)} processes running. "
+    #         "Please terminate the existing session first."
+    #     )
+    #     return
 
     try:
         results = obj.get_driver("process_manager").boot(
@@ -124,6 +125,43 @@ def boot(
             sys.exit(1)
 
 
+@click.command("terminate")
+@click.pass_obj
+@click.pass_context
+def terminate(ctx, obj):
+    """
+    Execute the process manager terminate command, but only do this for the current
+    session
+    """
+
+    log = get_logger("unified_shell.terminate")
+    session_query = ProcessQuery(session=ctx.obj.session_name)
+    log.info(f"Terminating session [green]{ctx.obj.session_name}[/]")
+    obj.get_driver("process_manager").kill(session_query)
+
+
+@click.command("ps")
+@click.pass_obj
+@click.pass_context
+def ps(ctx, obj):
+    """
+    Execute the process manager terminate command, but only do this for the current
+    session
+    """
+
+    log = get_logger("unified_shell.ps")
+    session_query = ProcessQuery(session=ctx.obj.session_name)
+    log.info(f"Listing session [green]{ctx.obj.session_name}[/]")
+    results = obj.get_driver("process_manager").ps(session_query)
+    obj.print(
+        tabulate_process_instance_list(
+            results, title=f"Processes running in session {ctx.obj.session_name}"
+        ),
+        overflow="fold",
+        soft_wrap=True,
+    )
+
+
 @click.command("start-shell")
 @click.pass_obj
 @click.pass_context

src/drunc/unified_shell/shell.py

@@ -52,12 +52,12 @@
     flush,
     kill,
     logs,
-    ps,
+    # ps,
     restart,
-    terminate,
+    # terminate,
 )
 from drunc.process_manager.interface.process_manager import run_pm
-from drunc.unified_shell.commands import boot, start_shell
+from drunc.unified_shell.commands import boot, ps, start_shell, terminate
 from drunc.unified_shell.context import UnifiedShellMode
 from drunc.unified_shell.shell_utils import generate_fsm_sequence_command
 from drunc.utils.configuration import ConfTypes, OKSKey
@@ -156,14 +156,14 @@ def unified_shell(
     ctx.obj.log.debug("Setting up the [green]unified_shell[/green] logger")
 
     # Parse the process manager argument to determine if it's a config or an address
-    unified_shell_log.critical(
+    ctx.obj.log.critical(
         f"Parsing the process manager argument: {process_manager}"
     )
     process_manager_url: ParseResult = urlparse(process_manager)
     internal_pm: bool = True
     if process_manager_url.scheme == "grpc":  # i.e. if it's an address
         internal_pm = False
-    unified_shell_log.critical(
+    ctx.obj.log.critical(
         f"Process manager argument parsed, internal_pm set to {internal_pm}"
     )
 
@@ -178,7 +178,7 @@ def unified_shell(
     #     )
     #     sys.exit(1)
 
-    unified_shell_log.critical("TEST")
+    ctx.obj.log.critical("TEST")
     # Setup configuration related context variables
     ctx.obj.configuration_file = f"oksconflibs:{configuration_file}"
     ctx.obj.configuration_id = configuration_id
@@ -267,7 +267,7 @@ def unified_shell(
         )
 
     else:  # Connect to an existing process manager at the provided address
-        unified_shell_log.critical(
+        ctx.obj.log.critical(
             "Connecting to an existing process manager at the provided address"
         )
         process_manager_address = process_manager.replace(
@@ -286,7 +286,7 @@ def unified_shell(
 
     # Run a simple command (describe) to check the connection with the process manager
     desc: Description | None = None
-    unified_shell_log.critical("Getting driver")
+    ctx.obj.log.critical("Getting driver")
     try:
         desc = ctx.obj.get_driver().describe()
     except Exception as e:
@@ -313,6 +313,7 @@ def unified_shell(
             ctx.obj.pm_process.join()
 
         sys.exit(1)
+    ctx.obj.log.critical("Process manager described successfully")
 
     # Broadcasting configuration if requested
     if desc.HasField("broadcast"):
@@ -323,18 +324,20 @@ def unified_shell(
 
     # Add the unified shell Click commands to the CLI
     ctx.obj.log.debug("Adding [green]unified_shell[/green] commands")
-    ctx.command.add_command(boot, "boot")
-    ctx.obj.dynamic_commands.add("boot")
+    unified_shell_commands = [boot, ps, terminate]
+    for cmd in unified_shell_commands:
+        ctx.command.add_command(cmd, format_name_for_cli(cmd.name))
+        ctx.obj.dynamic_commands.add(format_name_for_cli(cmd.name))
 
     # Add the process manager Click commands to the CLI
     ctx.obj.log.debug("Adding [green]process_manager[/green] commands")
     process_manager_commands: list[click.Command] = [
         kill,
-        terminate,
+        # terminate,
         flush,
         logs,
         restart,
-        ps,
+        # ps,
     ]
     for cmd in process_manager_commands:
         ctx.command.add_command(cmd, format_name_for_cli(cmd.name))
@@ -475,7 +478,8 @@ def cleanup():
 
         # Terminate any residual processes
         if ctx.obj.get_driver("process_manager"):
-            ctx.obj.get_driver("process_manager").terminate()
+            session_processes = ProcessQuery(session=ctx.obj.session_name)
+            ctx.obj.get_driver("process_manager").kill(session_processes)
 
         # Check if any processes are still running
         if (