From e94dc3505d4f1517b4156d37b8d580a148e1a37c Mon Sep 17 00:00:00 2001 From: Michal Skrivanek Date: Fri, 3 Oct 2025 10:16:44 +0200 Subject: [PATCH] add per-method description and timeout to shell commands introducing new format in addition to the simple methods: method: "command" you can now specify per-method customization with: methods: method: description: "This shows up in the CLI" method: "echo Hello" timeout: 5 --- packages/jumpstarter-driver-shell/README.md | 107 ++++++++-- .../jumpstarter_driver_shell/client.py | 24 ++- .../jumpstarter_driver_shell/driver.py | 51 ++++- .../jumpstarter_driver_shell/driver_test.py | 187 +++++++++++++----- 4 files changed, 286 insertions(+), 83 deletions(-) diff --git a/packages/jumpstarter-driver-shell/README.md b/packages/jumpstarter-driver-shell/README.md index d80f497bf..aa552cec5 100644 --- a/packages/jumpstarter-driver-shell/README.md +++ b/packages/jumpstarter-driver-shell/README.md @@ -11,7 +11,9 @@ $ pip3 install --extra-index-url {{index_url}} jumpstarter-driver-shell ## Configuration -Example configuration: +The shell driver supports two configuration formats for methods: + +### Format 1: Simple String e.g. for self-descriptive short commands ```yaml export: @@ -20,12 +22,40 @@ export: config: methods: ls: "ls" - method2: "echo 'Hello World 2'" - #multi line method - method3: | - echo 'Hello World $1' - echo 'Hello World $2' - env_var: "echo $1,$2,$ENV_VAR" + echo_hello: "echo 'Hello World'" +``` + +### Format 2: Unified Format with Descriptions + +```yaml +export: + shell: + type: jumpstarter_driver_shell.driver.Shell + config: + methods: + ls: + command: "ls -la" + description: "List directory contents with details" + deploy: + command: "ansible-playbook deploy.yml" + description: "Deploy application using Ansible" + # Multi-line commands work too + setup: + command: | + echo 'Setting up environment' + export PATH=$PATH:/usr/local/bin + ./setup.sh + description: "Set up the development environment" + # Description-only (uses default "echo Hello" command) + placeholder: + description: "Placeholder method for testing" + # Custom timeout for long-running operations + long_backup: + command: "tar -czf backup.tar.gz /data && rsync backup.tar.gz remote:/backups/" + description: "Create and sync backup (may take a while)" + timeout: 1800 # 30 minutes instead of default 5 minutes + # You can mix both formats + simple_echo: "echo 'simple'" # optional parameters cwd: "/tmp" log_level: "INFO" @@ -34,6 +64,25 @@ export: - "-c" ``` +### Configuration Parameters + +| Parameter | Description | Type | Required | Default | +|-----------|-------------|------|----------|---------| +| `methods` | Dictionary of methods. Values can be:
- String: just the command
- Dict: `{command: "...", description: "...", timeout: ...}` | `dict[str, str \| dict]` | Yes | - | +| `cwd` | Working directory for shell commands | `str` | No | `None` | +| `log_level` | Logging level | `str` | No | `"INFO"` | +| `shell` | Shell command to execute scripts | `list[str]` | No | `["bash", "-c"]` | +| `timeout` | Command timeout in seconds | `int` | No | `300` | + +**Method Configuration Options:** + +For the dict format, each method supports: +- `command`: The shell command to execute (optional, defaults to `"echo Hello"`) +- `description`: CLI help text (optional, defaults to `"Execute the {method_name} shell method"`) +- `timeout`: Command-specific timeout in seconds (optional, defaults to global `timeout` value) + +**Note:** You can mix both formats in the same configuration - use string format for simple commands and dict format when you want custom descriptions or timeouts. + ## API Reference Assuming the exporter driver is configured as in the example above, the client @@ -66,7 +115,11 @@ methods will be generated dynamically, and they will be available as follows: ## CLI Usage -The shell driver also provides a CLI when using `jmp shell`. All configured methods become available as CLI commands, except for methods starting with `_` which are considered private and hidden from the end user: +The shell driver also provides a CLI when using `jmp shell`. All configured methods become available as CLI commands, except for methods starting with `_` which are considered private and hidden from the end user. + +### CLI Help Output + +With unified format (custom descriptions): ```console $ jmp shell --exporter shell-exporter @@ -76,10 +129,40 @@ Usage: j shell [OPTIONS] COMMAND [ARGS]... Shell command executor Commands: - env_var Execute the env_var shell method - ls Execute the ls shell method - method2 Execute the method2 shell method - method3 Execute the method3 shell method + deploy Deploy application using Ansible + ls List directory contents with details + setup Set up the development environment +``` + +With simple string format (default descriptions): + +```console +$ j shell +Usage: j shell [OPTIONS] COMMAND [ARGS]... + + Shell command executor + +Commands: + deploy Execute the deploy shell method + ls Execute the ls shell method + setup Execute the setup shell method +``` + +**Mixed format example:** + +```yaml +methods: + deploy: + command: "ansible-playbook deploy.yml" + description: "Deploy using Ansible" + restart: "systemctl restart myapp" # Simple format +``` + +Results in: +```console +Commands: + deploy Deploy using Ansible + restart Execute the restart shell method ``` ### CLI Command Usage diff --git a/packages/jumpstarter-driver-shell/jumpstarter_driver_shell/client.py b/packages/jumpstarter-driver-shell/jumpstarter_driver_shell/client.py index ef8cee2f6..ea54137ed 100644 --- a/packages/jumpstarter-driver-shell/jumpstarter_driver_shell/client.py +++ b/packages/jumpstarter-driver-shell/jumpstarter_driver_shell/client.py @@ -58,13 +58,6 @@ def base(): def _add_method_command(self, group, method_name): """Add a Click command for a specific shell method""" - @group.command( - name=method_name, - context_settings={"ignore_unknown_options": True, "allow_interspersed_args": False}, - ) - @click.argument('args', nargs=-1, type=click.UNPROCESSED) - @click.option('--env', '-e', multiple=True, - help='Environment variables in KEY=VALUE format') def method_command(args, env): # Parse environment variables env_dict = {} @@ -81,5 +74,18 @@ def method_command(args, env): if returncode != 0: raise click.exceptions.Exit(returncode) - # Update the docstring dynamically - method_command.__doc__ = f"Execute the {method_name} shell method" + # Try to get custom description, fall back to default for older than 0.7 servers + try: + description = self.call("get_method_description", method_name) + except Exception: + description = f"Execute the {method_name} shell method" + + # Decorate and register the command + method_command.__doc__ = description + method_command = click.argument('args', nargs=-1, type=click.UNPROCESSED)(method_command) + method_command = click.option('--env', '-e', multiple=True, + help='Environment variables in KEY=VALUE format')(method_command) + method_command = group.command( + name=method_name, + context_settings={"ignore_unknown_options": True, "allow_interspersed_args": False}, + )(method_command) diff --git a/packages/jumpstarter-driver-shell/jumpstarter_driver_shell/driver.py b/packages/jumpstarter-driver-shell/jumpstarter_driver_shell/driver.py index f790a30f2..f504b1730 100644 --- a/packages/jumpstarter-driver-shell/jumpstarter_driver_shell/driver.py +++ b/packages/jumpstarter-driver-shell/jumpstarter_driver_shell/driver.py @@ -12,13 +12,36 @@ class Shell(Driver): """shell driver for Jumpstarter""" - # methods field is used to define the methods exported, and the shell script - # to be executed by each method - methods: dict[str, str] + # methods field defines the methods exported and their shell scripts + # Supports two formats: + # 1. Simple string: method_name: "command" + # 2. Dict with description: method_name: {command: "...", description: "...", timeout: ...} + methods: dict[str, str | dict[str, str | int]] shell: list[str] = field(default_factory=lambda: ["bash", "-c"]) timeout: int = 300 cwd: str | None = None + def _get_method_command(self, method: str) -> str: + """Extract the command string from a method configuration""" + method_config = self.methods[method] + if isinstance(method_config, str): + return method_config + return method_config.get("command", "echo Hello") + + def _get_method_description(self, method: str) -> str: + """Extract the description from a method configuration""" + method_config = self.methods[method] + if isinstance(method_config, str): + return f"Execute the {method} shell method" + return method_config.get("description", f"Execute the {method} shell method") + + def _get_method_timeout(self, method: str) -> int: + """Extract the timeout from a method configuration, fallback to global timeout""" + method_config = self.methods[method] + if isinstance(method_config, str): + return self.timeout + return method_config.get("timeout", self.timeout) + @classmethod def client(cls) -> str: return "jumpstarter_driver_shell.client.ShellClient" @@ -29,6 +52,11 @@ def get_methods(self) -> list[str]: self.logger.debug(f"get_methods called, returning methods: {methods}") return methods + @export + def get_method_description(self, method: str) -> str: + """Get the description for a specific method""" + return self._get_method_description(method) + @export async def call_method(self, method: str, env, *args) -> AsyncGenerator[tuple[str, str, int | None], None]: """ @@ -39,12 +67,13 @@ async def call_method(self, method: str, env, *args) -> AsyncGenerator[tuple[str self.logger.info(f"calling {method} with args: {args} and kwargs as env: {env}") if method not in self.methods: raise ValueError(f"Method '{method}' not found in available methods: {list(self.methods.keys())}") - script = self.methods[method] - self.logger.debug(f"running script: {script}") + script = self._get_method_command(method) + timeout = self._get_method_timeout(method) + self.logger.debug(f"running script: {script} with timeout: {timeout}") try: async for stdout_chunk, stderr_chunk, returncode in self._run_inline_shell_script( - method, script, *args, env_vars=env + method, script, *args, env_vars=env, timeout=timeout ): if stdout_chunk: self.logger.debug(f"{method} stdout:\n{stdout_chunk.rstrip()}") @@ -121,7 +150,7 @@ async def _read_process_output(self, process, read_all=False): return stdout_data, stderr_data async def _run_inline_shell_script( - self, method, script, *args, env_vars=None + self, method, script, *args, env_vars=None, timeout=None ) -> AsyncGenerator[tuple[str, str, int | None], None]: """ Run the given shell script with live streaming output. @@ -130,6 +159,7 @@ async def _run_inline_shell_script( :param script: The shell script contents as a string. :param args: Arguments to pass to the script (mapped to $1, $2, etc. in the script). :param env_vars: A dict of environment variables to make available to the script. + :param timeout: Customized command timeout in seconds. If None, uses global timeout. :yields: Tuples of (stdout_chunk, stderr_chunk, returncode). returncode is None until the process completes. @@ -151,9 +181,12 @@ async def _run_inline_shell_script( # Create a task to monitor the process timeout start_time = asyncio.get_event_loop().time() + if timeout is None: + timeout = self.timeout + # Read output in real-time while process.returncode is None: - if asyncio.get_event_loop().time() - start_time > self.timeout: + if asyncio.get_event_loop().time() - start_time > timeout: # Send SIGTERM to entire process group for graceful termination try: os.killpg(process.pid, signal.SIGTERM) @@ -168,7 +201,7 @@ async def _run_inline_shell_script( self.logger.warning(f"SIGTERM failed to terminate {process.pid}, sending SIGKILL") except (ProcessLookupError, OSError): pass - raise subprocess.TimeoutExpired(cmd, self.timeout) from None + raise subprocess.TimeoutExpired(cmd, timeout) from None try: stdout_data, stderr_data = await self._read_process_output(process, read_all=False) diff --git a/packages/jumpstarter-driver-shell/jumpstarter_driver_shell/driver_test.py b/packages/jumpstarter-driver-shell/jumpstarter_driver_shell/driver_test.py index 4e03b5f64..c3621b021 100644 --- a/packages/jumpstarter-driver-shell/jumpstarter_driver_shell/driver_test.py +++ b/packages/jumpstarter-driver-shell/jumpstarter_driver_shell/driver_test.py @@ -101,74 +101,155 @@ def test_cli_method_execution(client): """Test that CLI methods can be executed""" cli = client.cli() - # Test that we can get the echo command - echo_command = cli.commands.get('echo') - assert echo_command is not None + # Test that CLI commands exist and have the correct structure + echo_command = cli.commands['echo'] assert echo_command.name == 'echo' def test_cli_includes_all_methods(): - """Test that CLI includes all methods""" - from .driver import Shell - from jumpstarter.common.utils import serve - - shell_instance = Shell( - log_level="DEBUG", + """Test that CLI includes all configured methods with proper names""" + shell = Shell( methods={ - "method1": "echo method1", - "method2": "echo method2", - "method3": "echo method3", - }, + "echo": "echo", + "env": "echo $ENV_VAR", + "multi_line": "echo line1\necho line2", + "exit1": "exit 1", + "stderr": "echo 'error' 1>&2", + } ) - with serve(shell_instance) as test_client: - cli = test_client.cli() + with serve(shell) as client: + cli = client.cli() + + # Check that all methods are in the CLI + expected_methods = {"echo", "env", "multi_line", "exit1", "stderr"} available_commands = set(cli.commands.keys()) - # All methods should be available - expected_methods = {"method1", "method2", "method3"} assert available_commands == expected_methods, f"Expected {expected_methods}, got {available_commands}" def test_cli_exit_codes(): - """Test that CLI commands preserve shell command exit codes""" - import click + """Test that CLI methods correctly exit with shell command return codes""" + shell = Shell( + methods={ + "exit0": "exit 0", + "exit1": "exit 1", + "exit42": "exit 42", + } + ) - from .driver import Shell - from jumpstarter.common.utils import serve + with serve(shell) as client: + # Test successful command (exit 0) + returncode = client.exit0() + assert returncode == 0 - # Create a shell instance with methods that have different exit codes - shell_instance = Shell( - log_level="DEBUG", + # Test failed command (exit 1) + returncode = client.exit1() + assert returncode == 1 + + # Test custom exit code (exit 42) + returncode = client.exit42() + assert returncode == 42 + + +def test_cli_custom_descriptions_unified_format(): + """Test that CLI methods use custom descriptions with unified format""" + shell = Shell( methods={ - "success": "exit 0", - "fail_1": "exit 1", - "fail_42": "exit 42", - }, + "echo": { + "command": "echo", + "description": "Custom echo description" + }, + "test_method": { + "command": "echo 'test'", + "description": "Test method with custom description" + }, + } + ) + + with serve(shell) as client: + cli = client.cli() + + # Check that custom descriptions are used + echo_cmd = cli.commands['echo'] + assert echo_cmd.help == "Custom echo description", f"Expected custom description, got {echo_cmd.help}" + + test_cmd = cli.commands['test_method'] + assert test_cmd.help == "Test method with custom description" + + +def test_cli_default_descriptions(): + """Test that CLI methods use default descriptions when not configured""" + shell = Shell( + methods={ + "echo": "echo", + "test_method": "echo 'test'", + } + # No descriptions configured + ) + + with serve(shell) as client: + cli = client.cli() + + # Check that default descriptions are used + echo_cmd = cli.commands['echo'] + assert echo_cmd.help == "Execute the echo shell method" + + test_cmd = cli.commands['test_method'] + assert test_cmd.help == "Execute the test_method shell method" + + +def test_get_method_description_unified(): + """Test the get_method_description export method with unified format""" + shell = Shell( + methods={ + "method1": { + "command": "echo", + "description": "Custom description for method1" + }, + "method2": "ls", # String format, should use default description + } ) - with serve(shell_instance) as test_client: - cli = test_client.cli() - - # Test successful command (exit 0) - should not raise - success_cmd = cli.commands['success'] - try: - success_cmd.callback([], []) # Call with empty args and env - except click.exceptions.Exit: - raise AssertionError("Success command should not raise Exit exception") from None - - # Test command that exits with code 1 - should raise Exit(1) - fail1_cmd = cli.commands['fail_1'] - try: - fail1_cmd.callback([], []) - raise AssertionError("Command should have raised Exit exception") - except click.exceptions.Exit as e: - assert e.exit_code == 1 - - # Test command that exits with code 42 - should raise Exit(42) - fail42_cmd = cli.commands['fail_42'] - try: - fail42_cmd.callback([], []) - raise AssertionError("Command should have raised Exit exception") - except click.exceptions.Exit as e: - assert e.exit_code == 42 + with serve(shell) as client: + # Test with custom description (unified format) + assert client.call("get_method_description", "method1") == "Custom description for method1" + + # Test with default description (string format) + assert client.call("get_method_description", "method2") == "Execute the method2 shell method" + + +def test_mixed_format_methods(): + """Test that both string and dict formats work together""" + shell = Shell( + methods={ + "simple": "echo 'simple'", + "detailed": { + "command": "echo 'detailed'", + "description": "A detailed command with description" + }, + "default_cmd": { + # No command specified - should use default "echo Hello" + "description": "Method using default command" + } + } + ) + + with serve(shell) as client: + # Test string format works + returncode = client.simple() + assert returncode == 0 + + # Test dict format works + returncode = client.detailed() + assert returncode == 0 + + # Test default command works + returncode = client.default_cmd() + assert returncode == 0 + + # Test CLI descriptions + cli = client.cli() + assert cli.commands['simple'].help == "Execute the simple shell method" + assert cli.commands['detailed'].help == "A detailed command with description" + assert cli.commands['default_cmd'].help == "Method using default command"