diff --git a/nodescraper/cli/cli.py b/nodescraper/cli/cli.py index 129ef136..30dc8792 100644 --- a/nodescraper/cli/cli.py +++ b/nodescraper/cli/cli.py @@ -32,6 +32,7 @@ import platform import sys import uuid +from collections.abc import Callable, Sequence from typing import Optional import nodescraper @@ -65,6 +66,7 @@ from nodescraper.constants import DEFAULT_LOGGER from nodescraper.enums import ExecutionStatus, SystemInteractionLevel, SystemLocation from nodescraper.models import SystemInfo +from nodescraper.models.pluginresult import PluginResult from nodescraper.pluginexecutor import PluginExecutor from nodescraper.pluginregistry import PluginRegistry @@ -461,6 +463,7 @@ def main( arg_input: Optional[list[str]] = None, *, host_cli_args: Optional[argparse.Namespace] = None, + plugin_run_result_hooks: Optional[Sequence[Callable[[PluginResult], None]]] = None, ): """Main entry point for the CLI @@ -468,6 +471,8 @@ def main( arg_input (Optional[list[str]], optional): list of args to parse. Defaults to None. host_cli_args: Optional namespace from an embedding host (e.g. detect-errors) for code that calls get_plugin_run_invocation during the plugin queue. + plugin_run_result_hooks: Optional callbacks invoked with each plugin's :class:`PluginResult` + after ``run()`` completes (used by embedded hosts such as error-scraper). """ if arg_input is None: arg_input = sys.argv[1:] @@ -643,6 +648,7 @@ def main( sname=sname, host_cli_args=host_cli_args, session_id=str(uuid.uuid4()), + plugin_run_result_hooks=plugin_run_result_hooks, ) log_system_info(log_path, system_info, logger) diff --git a/nodescraper/cli/embed.py b/nodescraper/cli/embed.py index 60d94515..b1e91c37 100644 --- a/nodescraper/cli/embed.py +++ b/nodescraper/cli/embed.py @@ -27,9 +27,11 @@ from __future__ import annotations import argparse +from collections.abc import Callable, Sequence from typing import Optional from nodescraper.cli.cli import get_cli_top_level_subcommands +from nodescraper.models.pluginresult import PluginResult CLI_TOP_LEVEL_SUBCOMMANDS = get_cli_top_level_subcommands() @@ -45,29 +47,38 @@ def run_cli_return_code( argv: list[str], *, host_cli_args: Optional[argparse.Namespace] = None, + plugin_run_result_hooks: Optional[Sequence[Callable[[PluginResult], None]]] = None, ) -> int: """Run nodescraper in-process; same behavior as :func:`run_main_return_code`. Args: argv: Tokens after the program name. host_cli_args: Optional host namespace forwarded to :func:`nodescraper.cli.cli.main`. + plugin_run_result_hooks: Optional callbacks invoked with each + :class:`~nodescraper.models.pluginresult.PluginResult` after a plugin finishes (embed hosts). Returns: Integer exit code (``SystemExit`` is mapped, not raised). """ - return run_main_return_code(argv, host_cli_args=host_cli_args) + return run_main_return_code( + argv, + host_cli_args=host_cli_args, + plugin_run_result_hooks=plugin_run_result_hooks, + ) def run_main_return_code( arg_input: list[str], *, host_cli_args: Optional[argparse.Namespace] = None, + plugin_run_result_hooks: Optional[Sequence[Callable[[PluginResult], None]]] = None, ) -> int: """Run :func:`nodescraper.cli.cli.main` and map ``SystemExit`` to an exit code. Args: arg_input: Tokens after the program name. host_cli_args: Optional host namespace for embedded runs. + plugin_run_result_hooks: Optional per-plugin result callbacks for embedded runs. Returns: Integer exit code. @@ -75,7 +86,11 @@ def run_main_return_code( from nodescraper.cli.cli import main try: - main(arg_input, host_cli_args=host_cli_args) + main( + arg_input, + host_cli_args=host_cli_args, + plugin_run_result_hooks=plugin_run_result_hooks, + ) except SystemExit as exc: code = exc.code if code is None: diff --git a/nodescraper/cli/invocation.py b/nodescraper/cli/invocation.py index ee59e4a6..9edc7214 100644 --- a/nodescraper/cli/invocation.py +++ b/nodescraper/cli/invocation.py @@ -28,6 +28,7 @@ import argparse import logging +from collections.abc import Callable, Sequence from contextlib import contextmanager from contextvars import ContextVar from dataclasses import dataclass @@ -72,6 +73,7 @@ class PluginRunInvocation: sname: str host_cli_args: Optional[argparse.Namespace] = None session_id: Optional[str] = None + plugin_run_result_hooks: tuple[Callable[[PluginResult], None], ...] = () def run_plugin_queue_with_invocation( @@ -86,8 +88,12 @@ def run_plugin_queue_with_invocation( sname: str, host_cli_args: Optional[argparse.Namespace] = None, session_id: Optional[str] = None, + plugin_run_result_hooks: Optional[Sequence[Callable[[PluginResult], None]]] = None, ) -> list[PluginResult]: """Constructs the plugin executor, binds invocation context, and runs the plugin queue.""" + hooks_tuple: tuple[Callable[[PluginResult], None], ...] = ( + tuple(plugin_run_result_hooks) if plugin_run_result_hooks else () + ) inv = PluginRunInvocation( plugin_reg=plugin_reg, parsed_args=parsed_args, @@ -99,6 +105,7 @@ def run_plugin_queue_with_invocation( sname=sname, host_cli_args=host_cli_args, session_id=session_id, + plugin_run_result_hooks=hooks_tuple, ) plugin_executor = PluginExecutor( logger=logger, @@ -108,6 +115,7 @@ def run_plugin_queue_with_invocation( log_path=log_path, plugin_registry=plugin_reg, session_id=session_id, + plugin_run_result_hooks=hooks_tuple, ) with plugin_run_invocation_scope(inv): return plugin_executor.run_queue() diff --git a/nodescraper/pluginexecutor.py b/nodescraper/pluginexecutor.py index 4f3febed..bb22b8a9 100644 --- a/nodescraper/pluginexecutor.py +++ b/nodescraper/pluginexecutor.py @@ -30,6 +30,7 @@ import logging import uuid from collections import deque +from collections.abc import Callable, Sequence from typing import Optional, Type, Union from pydantic import BaseModel @@ -38,6 +39,7 @@ from nodescraper.connection.oob_ssh import OobSshConnectionManager from nodescraper.constants import DEFAULT_LOGGER from nodescraper.interfaces import ConnectionManager, DataPlugin, PluginInterface +from nodescraper.interfaces.taskresulthook import TaskResultHook from nodescraper.models import PluginConfig, SystemInfo from nodescraper.models.pluginresult import PluginResult from nodescraper.pluginregistry import PluginRegistry @@ -57,6 +59,7 @@ def __init__( plugin_registry: Optional[PluginRegistry] = None, log_path: Optional[str] = None, session_id: Optional[str] = None, + plugin_run_result_hooks: Optional[Sequence[Callable[[PluginResult], None]]] = None, ): if logger is None: @@ -89,7 +92,11 @@ def __init__( self.log_path = log_path - self.connection_result_hooks = [] + self.plugin_run_result_hooks: list[Callable[[PluginResult], None]] = ( + list(plugin_run_result_hooks) if plugin_run_result_hooks else [] + ) + + self.connection_result_hooks: list[TaskResultHook] = [] if log_path: self.connection_result_hooks.append(FileSystemLogHook(log_base_path=log_path)) @@ -263,7 +270,10 @@ def run_queue(self) -> list[PluginResult]: continue self.logger.info("-" * 50) - plugin_results.append(plugin_inst.run(**run_payload)) + plugin_result = plugin_inst.run(**run_payload) + plugin_results.append(plugin_result) + for hook in self.plugin_run_result_hooks: + hook(plugin_result) except Exception as e: self.logger.exception( "Unexpected exception when running plugin %s: %s", plugin_name, e diff --git a/test/unit/cli/test_cli_embed_api.py b/test/unit/cli/test_cli_embed_api.py index db44f6cf..54b95043 100644 --- a/test/unit/cli/test_cli_embed_api.py +++ b/test/unit/cli/test_cli_embed_api.py @@ -53,7 +53,12 @@ def test_run_cli_return_code_and_run_main_return_code_delegate( ) -> None: calls: list[list[str]] = [] - def fake_main(arg_input: list[str], *, host_cli_args=None) -> None: + def fake_main( + arg_input: list[str], + *, + host_cli_args=None, + plugin_run_result_hooks=None, + ) -> None: calls.append(list(arg_input)) raise SystemExit(7) diff --git a/test/unit/framework/test_plugin_executor.py b/test/unit/framework/test_plugin_executor.py index fe9a8954..494551ce 100644 --- a/test/unit/framework/test_plugin_executor.py +++ b/test/unit/framework/test_plugin_executor.py @@ -186,3 +186,18 @@ def test_connection_manager_from_plugin_when_not_in_registry(): assert len(results) == 1 assert results[0].source == "testB" assert results[0].status == ExecutionStatus.OK + + +def test_plugin_run_result_hooks_called_after_each_plugin(plugin_registry): + seen: list[str] = [] + + def hook(res: PluginResult) -> None: + seen.append(res.source) + + executor = PluginExecutor( + plugin_configs=[PluginConfig(plugins={"TestPluginB": {}})], + plugin_registry=plugin_registry, + plugin_run_result_hooks=[hook], + ) + executor.run_queue() + assert seen == ["testB"]