diff --git a/README.md b/README.md index e96be43..541cf78 100644 --- a/README.md +++ b/README.md @@ -35,8 +35,50 @@ logger: ## Features - JSON event streaming to message brokers + - RabbitMQ stream queue support via STOMP queue headers - SSL/TLS support - - Configurable formatters (default + JLab SWF schema) + - Configurable formatters (default + JLab SWF schema + ComprehensiveEventFormatter) - Event filtering (include/exclude) - Heartbeat management - - Environment variable support for secrets \ No newline at end of file + - Optional application-level consumer heartbeat events + - Environment variable support for secrets + +## Consumer Heartbeat Events + +In addition to STOMP transport heartbeats, you can emit periodic heartbeat messages +that are delivered to consumers on the configured destination. + +Set `consumer_heartbeat_interval` in seconds: + +```yaml +logger: + stomp: + consumer_heartbeat_interval: 30 +``` + +Use `0` (default) to disable consumer heartbeat events. + +## RabbitMQ Streams + +RabbitMQ streams are supported as a queue type when the broker is accessed via the +RabbitMQ STOMP adapter. + +Example profile configuration: + +```yaml +logger: + stomp: + host: "rabbitmq.example.com" + port: 61613 + user: "${STOMP_USER}" + password: "${STOMP_PASSWORD}" + queue: "/queue/snakemake.stream" + use_stream: true + stream_filter_by_workflow: true +``` + +Notes: + +- RabbitMQ stream support in this plugin is publish-only. +- RabbitMQ queue type is immutable. If a destination already exists as a classic queue, + it must be deleted and recreated as a stream before `use_stream: true` will work. \ No newline at end of file diff --git a/docs/further.md b/docs/further.md index f8e5502..b3f27be 100644 --- a/docs/further.md +++ b/docs/further.md @@ -108,6 +108,40 @@ snakemake --logger stomp \ The `SWF_RUN_NUMBER` environment variable can be set to populate the `correlation.run_number` field for tracking related workflow runs. +### ComprehensiveEventFormatter + +Human-readable event view with rich passthrough of Snakemake event attributes. +This formatter keeps top-level context fields and adds an `event_details` section +containing nearly all event-specific data available on the log record. + +```json +{ + "timestamp": "2024-01-15T14:30:00Z", + "hostname": "compute-node-01", + "workflow_id": "550e8400-e29b-41d4-a716-446655440000", + "event_type": "job_finished", + "level": "INFO", + "message": "Rule completed", + "summary": "Job Finished | rule=align_reads | job=5 | wildcards={'sample': 'A'} | Rule completed", + "event_details": { + "event": "job_finished", + "jobid": 5, + "name": "align_reads", + "wildcards": {"sample": "A"}, + "threads": 4, + "resources": {"mem_mb": 16000}, + "benchmark": "bench/align_reads.txt" + } +} +``` + +**To use ComprehensiveEventFormatter:** + +```bash +snakemake --logger stomp \ + --logger-stomp-formatter-class "snakemake_logger_plugin_stomp.formatters.ComprehensiveEventFormatter" +``` + ## Custom Formatters Create your own formatter by implementing the `BaseFormatter` interface: @@ -198,6 +232,53 @@ queue: "/topic/snakemake.broadcast" Use **queues** when you want load balancing across multiple consumers. Use **topics** when multiple systems need to receive all events (e.g., monitoring + archival + alerting). +## RabbitMQ Streams + +RabbitMQ streams are a queue type exposed by the RabbitMQ STOMP adapter. This plugin +supports stream publishing and stream declaration for queue destinations. + +### Creating a Stream + +Use a `/queue/` destination together with `use_stream: true` to have RabbitMQ +declare the destination as a stream on first publish: + +```yaml +logger: + stomp: + queue: "/queue/snakemake.stream" + use_stream: true +``` + +### Publishing to an Existing Stream + +Use `/amq/queue/` when the stream already exists and should not be redeclared: + +```yaml +logger: + stomp: + queue: "/amq/queue/snakemake.stream" + use_stream: true +``` + +### Stream Filter Headers + +RabbitMQ streams support outbound filter values on published messages. Use +`stream_filter_by_workflow: true` to set `x-stream-filter-value` to the current +Snakemake `workflow_id` on each published event once the workflow ID is available: + +```yaml +logger: + stomp: + queue: "/queue/snakemake.stream" + use_stream: true + stream_filter_by_workflow: true +``` + +### Constraints + +- Stream support is publish-only in this plugin. +- RabbitMQ queue type is immutable, so an existing classic queue cannot be converted in place. + ## Complete Production Example A full production configuration combining SSL, event filtering, custom formatter, and heartbeat tuning: diff --git a/docs/intro.md b/docs/intro.md index 1ccf26d..3c140f2 100644 --- a/docs/intro.md +++ b/docs/intro.md @@ -5,8 +5,9 @@ Stream Snakemake workflow events to STOMP message brokers (ActiveMQ, RabbitMQ, A ## Key Features - **Real-time event streaming** to STOMP-compatible message brokers +- **RabbitMQ stream queue support** for append-only log style destinations - **SSL/TLS encryption** for secure production deployments -- **Flexible formatters** - default flat JSON or JLab Scientific Workflow schema +- **Flexible formatters** - default flat JSON, JLab Scientific Workflow schema, or comprehensive Snakemake logs - **Event filtering** - include/exclude specific event types to reduce noise - **Custom formatters** - plugin your own formatter classes - **Environment variable support** for credentials and sensitive configuration @@ -48,6 +49,25 @@ Then run: snakemake --profile profiles/stomp ``` +### RabbitMQ Stream Destinations + +When using RabbitMQ's STOMP adapter, the plugin can publish to stream queues. + +```yaml +logger: + stomp: + host: "rabbitmq.example.com" + port: 61613 + user: "${STOMP_USER}" + password: "${STOMP_PASSWORD}" + queue: "/queue/snakemake.stream" + use_stream: true + stream_filter_by_workflow: true +``` + +Use `/queue/` to create a new stream on first publish. Use `/amq/queue/` +to publish to an existing stream without redeclaring it. + ### Environment Variables Secure credentials using environment variables: diff --git a/pyproject.toml b/pyproject.toml index 9964c47..1bbe80d 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -22,6 +22,7 @@ dependencies = [ "snakemake-interface-common>=1.22.0,<2", "snakemake-interface-logger-plugins>=2.0.0,<3", "stomp-py>=8.2.0", + "paramiko>=3.4.0", ] [[project.authors]] diff --git a/src/snakemake_logger_plugin_stomp/__init__.py b/src/snakemake_logger_plugin_stomp/__init__.py index 713e6d4..dfd5c02 100644 --- a/src/snakemake_logger_plugin_stomp/__init__.py +++ b/src/snakemake_logger_plugin_stomp/__init__.py @@ -9,10 +9,15 @@ import importlib import json import logging as py_logging +import os +from pathlib import Path import socket import ssl +import threading +import time import uuid from dataclasses import dataclass, field +from datetime import UTC, datetime from typing import Optional import stomp @@ -21,6 +26,7 @@ from snakemake_interface_logger_plugins.settings import LogHandlerSettingsBase from snakemake_logger_plugin_stomp.formatters import DefaultJSONFormatter +from snakemake_logger_plugin_stomp.ssh_tunnel import SSHTunnelManager @dataclass @@ -125,6 +131,104 @@ class LogHandlerSettings(LogHandlerSettingsBase): "env_var": False, }, ) + use_stream: bool = field( + default=False, + metadata={ + "help": "Declare the configured queue destination as a RabbitMQ stream", + "env_var": False, + }, + ) + stream_filter_by_workflow: bool = field( + default=False, + metadata={ + "help": "Set RabbitMQ x-stream-filter-value header to workflow_id on outbound messages", + "env_var": False, + }, + ) + consumer_heartbeat_interval: float = field( + default=0, + metadata={ + "help": "Application-level consumer heartbeat interval in seconds (0 to disable)", + "env_var": False, + }, + ) + use_ssh_tunnel: bool = field( + default=False, + metadata={ + "help": "Enable SSH tunnel transport between logger and STOMP broker", + "env_var": False, + }, + ) + ssh_host: Optional[str] = field( + default=None, + metadata={ + "help": "SSH jump host used for tunnel", + "env_var": False, + }, + ) + ssh_port: int = field( + default=22, + metadata={ + "help": "SSH server port", + "env_var": False, + }, + ) + ssh_username: Optional[str] = field( + default=None, + metadata={ + "help": "SSH username", + "env_var": True, + }, + ) + ssh_private_key: Optional[str] = field( + default=None, + metadata={ + "help": "Path to SSH private key for tunnel authentication", + "env_var": False, + }, + ) + ssh_key_passphrase: Optional[str] = field( + default=None, + metadata={ + "help": "Passphrase for SSH private key", + "env_var": True, + }, + ) + ssh_connect_timeout: int = field( + default=10, + metadata={ + "help": "SSH connect timeout in seconds", + "env_var": False, + }, + ) + ssh_local_bind_port: int = field( + default=0, + metadata={ + "help": "Local port for SSH tunnel (0 to auto-select)", + "env_var": False, + }, + ) + reconnect_attempts: int = field( + default=3, + metadata={ + "help": "Number of connection recovery attempts after send failures", + "env_var": False, + }, + ) + reconnect_initial_backoff_seconds: float = field( + default=1.0, + metadata={ + "help": "Initial reconnection backoff delay in seconds", + "env_var": False, + }, + ) + reconnect_max_backoff_seconds: float = field( + default=8.0, + metadata={ + "help": "Maximum reconnection backoff delay in seconds", + "env_var": False, + }, + ) class StompConnectionListener(stomp.ConnectionListener): @@ -184,13 +288,28 @@ def __post_init__(self) -> None: # STOMP connection handle self.connection = None + self._ssh_tunnel_manager = None + self._connection_host = self.settings.host + self._connection_port = self.settings.port + self._stream_declared = False + self._consumer_heartbeat_stop_event = threading.Event() + self._consumer_heartbeat_thread = None # Workflow tracking metadata self.workflow_metadata = { - "workflow_id": None, # Generated on WORKFLOW_STARTED event + "workflow_id": None, # Generated on WORKFLOW_STARTED event or consumer heartbeat "hostname": socket.gethostname(), + "working_directory": str(Path.cwd()), + "user": os.getenv("USER") or os.getenv("USERNAME") or "unknown", + "heartbeat_interval_seconds": self.settings.consumer_heartbeat_interval, + "workflow_initiated": datetime.now(UTC).isoformat(), } + self._validate_stream_settings() + self._validate_consumer_heartbeat_settings() + self._validate_reconnect_settings() + self._validate_ssh_settings() + # Load and initialize the configured formatter self.formatter_instance = self._init_formatter() @@ -207,6 +326,7 @@ def __post_init__(self) -> None: # Establish connection to STOMP broker self._connect_to_broker() + self._start_consumer_heartbeat_loop() def _init_formatter(self): """Load and instantiate the configured formatter class. @@ -233,15 +353,42 @@ def _init_formatter(self): ) return DefaultJSONFormatter() - def _connect_to_broker(self): + def _resolve_broker_endpoint(self) -> tuple[str, int]: + """Resolve broker endpoint, optionally by creating an SSH tunnel.""" + if not self.settings.use_ssh_tunnel: + return self.settings.host, self.settings.port + + if self._ssh_tunnel_manager is None: + self._ssh_tunnel_manager = SSHTunnelManager( + ssh_host=self.settings.ssh_host, + ssh_port=self.settings.ssh_port, + ssh_username=self.settings.ssh_username, + ssh_private_key=self.settings.ssh_private_key, + ssh_key_passphrase=self.settings.ssh_key_passphrase, + remote_host=self.settings.host, + remote_port=self.settings.port, + local_bind_port=self.settings.ssh_local_bind_port, + connect_timeout=self.settings.ssh_connect_timeout, + logger=self._internal_logger, + ) + + endpoint = self._ssh_tunnel_manager.connect() + self._internal_logger.info( + "[STOMP] SSH tunnel established: " + f"127.0.0.1:{endpoint[1]} -> {self.settings.host}:{self.settings.port}" + ) + return endpoint + + def _connect_to_broker(self, raise_on_error: bool = False) -> bool: """Establish connection to STOMP message broker. Configures SSL if enabled and sets up connection listener for error handling. Connection errors are logged but don't halt execution. """ - hosts = [(self.settings.host, self.settings.port)] - try: + self._connection_host, self._connection_port = self._resolve_broker_endpoint() + hosts = [(self._connection_host, self._connection_port)] + # Create STOMP connection with heartbeat configuration self.connection = stomp.Connection( host_and_ports=hosts, @@ -273,17 +420,178 @@ def _connect_to_broker(self): ) self._internal_logger.info( - f"[STOMP] Connected to {self.settings.host}:{self.settings.port}" + f"[STOMP] Connected to {self._connection_host}:{self._connection_port}" ) + return True except (ConnectionError, OSError, Exception) as e: self._internal_logger.error( f"[STOMP] Failed to connect to {self.settings.host}:{self.settings.port}: {e}" ) - if self.settings.fail_on_connection_error: + if self._ssh_tunnel_manager is not None: + self._ssh_tunnel_manager.close() + self._ssh_tunnel_manager = None + self.connection = None + + if self.settings.fail_on_connection_error or raise_on_error: raise RuntimeError( f"STOMP logging required but broker unavailable: {e}" ) from e + return False + + def _validate_stream_settings(self) -> None: + """Validate RabbitMQ stream-specific settings.""" + if not self.settings.use_stream: + return + + def _validate_consumer_heartbeat_settings(self) -> None: + """Validate consumer heartbeat settings.""" + if self.settings.consumer_heartbeat_interval < 0: + raise ValueError("consumer_heartbeat_interval must be >= 0") + + def _validate_reconnect_settings(self) -> None: + """Validate reconnection backoff settings.""" + if self.settings.reconnect_attempts < 0: + raise ValueError("reconnect_attempts must be >= 0") + if self.settings.reconnect_initial_backoff_seconds <= 0: + raise ValueError("reconnect_initial_backoff_seconds must be > 0") + if self.settings.reconnect_max_backoff_seconds <= 0: + raise ValueError("reconnect_max_backoff_seconds must be > 0") + if ( + self.settings.reconnect_initial_backoff_seconds + > self.settings.reconnect_max_backoff_seconds + ): + raise ValueError( + "reconnect_initial_backoff_seconds must be <= reconnect_max_backoff_seconds" + ) + + def _validate_ssh_settings(self) -> None: + """Validate SSH tunnel settings when enabled.""" + if not self.settings.use_ssh_tunnel: + return + + if not self.settings.ssh_host: + raise ValueError("ssh_host is required when use_ssh_tunnel is true") + if not self.settings.ssh_username: + raise ValueError("ssh_username is required when use_ssh_tunnel is true") + if not self.settings.ssh_private_key: + raise ValueError("ssh_private_key is required when use_ssh_tunnel is true") + if self.settings.ssh_connect_timeout <= 0: + raise ValueError("ssh_connect_timeout must be > 0") + if self.settings.ssh_local_bind_port < 0: + raise ValueError("ssh_local_bind_port must be >= 0") + if self.settings.ssh_port <= 0: + raise ValueError("ssh_port must be > 0") + + key_path = Path(self.settings.ssh_private_key) + if not key_path.exists() or not key_path.is_file(): + raise ValueError( + "ssh_private_key must point to an existing private key file" + ) + + def _recover_connection(self) -> bool: + """Attempt to recover the broker connection with bounded exponential backoff.""" + attempts = self.settings.reconnect_attempts + if attempts == 0: + return False + + backoff = self.settings.reconnect_initial_backoff_seconds + max_backoff = self.settings.reconnect_max_backoff_seconds + + for attempt in range(1, attempts + 1): + self._internal_logger.warning( + f"[STOMP] Attempting connection recovery ({attempt}/{attempts})" + ) + if self._connect_to_broker(raise_on_error=False): + self._internal_logger.info("[STOMP] Connection recovery succeeded") + return True + + if attempt < attempts: + time.sleep(backoff) + backoff = min(backoff * 2, max_backoff) + + self._internal_logger.error("[STOMP] Connection recovery failed") + return False + + def _start_consumer_heartbeat_loop(self) -> None: + """Start background loop that emits consumer heartbeat events.""" + if self.settings.consumer_heartbeat_interval <= 0: + return + + if not self.workflow_metadata["workflow_id"]: + self.workflow_metadata["workflow_id"] = str(uuid.uuid4()) + self._internal_logger.info( + f"[STOMP] Generated workflow ID: {self.workflow_metadata['workflow_id']}" + ) + + self._consumer_heartbeat_thread = threading.Thread( + target=self._consumer_heartbeat_loop, + name="stomp-consumer-heartbeat", + daemon=True, + ) + self._consumer_heartbeat_thread.start() + self._internal_logger.info( + "[STOMP] Consumer heartbeat enabled every " + f"{self.settings.consumer_heartbeat_interval}s" + ) + + def _consumer_heartbeat_loop(self) -> None: + """Emit consumer heartbeat immediately, then at the configured interval.""" + interval = self.settings.consumer_heartbeat_interval + + if self._consumer_heartbeat_stop_event.is_set(): + return + + self._emit_consumer_heartbeat() + + while not self._consumer_heartbeat_stop_event.wait(interval): + self._emit_consumer_heartbeat() + + def _emit_consumer_heartbeat(self) -> None: + """Emit one synthetic heartbeat event intended for downstream consumers.""" + try: + event_data = self.formatter_instance.format( + None, self.workflow_metadata + ) + except Exception as e: + self._internal_logger.error( + f"[STOMP] Formatter error for consumer heartbeat event: {e}" + ) + return + + self._send_to_broker(event_data) + + def _should_declare_stream_on_send(self) -> bool: + """Whether the next publish should include stream declaration headers.""" + return self.settings.use_stream and not self._stream_declared + + def _build_stream_declare_headers(self) -> dict[str, str]: + """Build RabbitMQ queue declaration headers for stream destinations.""" + if not self._should_declare_stream_on_send(): + return {} + + return {"x-queue-type": "stream"} + + def _build_stream_publish_headers(self) -> dict[str, str]: + """Build per-message RabbitMQ stream headers.""" + if not self.settings.use_stream or not self.settings.stream_filter_by_workflow: + return {} + + workflow_id = self.workflow_metadata.get("workflow_id") + if not workflow_id: + return {} + + return {"x-stream-filter-value": workflow_id} + + def _build_send_headers(self) -> dict[str, str]: + """Build STOMP SEND headers for the next event publish.""" + headers = { + "persistent": "true", + "content-type": "application/json", + } + headers.update(self._build_stream_declare_headers()) + headers.update(self._build_stream_publish_headers()) + return headers def _send_to_broker(self, event_data: dict): """Send formatted event data to STOMP broker. @@ -293,22 +601,27 @@ def _send_to_broker(self, event_data: dict): """ # Skip if not connected if not self.connection or not self.connection.is_connected(): - self._internal_logger.debug( - "[STOMP] Skipping send - not connected to broker" + self._internal_logger.warning( + "[STOMP] Not connected to broker; attempting recovery before send" ) - return + if not self._recover_connection(): + self._internal_logger.debug( + "[STOMP] Skipping send - unable to recover broker connection" + ) + return try: # Send message with JSON body and appropriate headers + headers = self._build_send_headers() self.connection.send( body=json.dumps(event_data, default=str), destination=self.settings.queue, - headers={ - "persistent": "true", # Request broker persistence - "content-type": "application/json", - }, + headers=headers, ) + if self._should_declare_stream_on_send(): + self._stream_declared = True + self._internal_logger.debug( f"[STOMP] Sent event: {event_data.get('event_type', 'UNKNOWN')}" ) @@ -387,6 +700,10 @@ def close(self): Called when Snakemake workflow completes or logger is being shut down. """ + self._consumer_heartbeat_stop_event.set() + if self._consumer_heartbeat_thread and self._consumer_heartbeat_thread.is_alive(): + self._consumer_heartbeat_thread.join(timeout=2) + if self.connection and self.connection.is_connected(): try: self.connection.disconnect() @@ -394,6 +711,16 @@ def close(self): except Exception as e: self._internal_logger.error(f"[STOMP] Error during disconnect: {e}") + if self._ssh_tunnel_manager is not None: + try: + self._ssh_tunnel_manager.close() + except Exception as e: + self._internal_logger.error( + f"[STOMP] Error during SSH tunnel cleanup: {e}" + ) + finally: + self._ssh_tunnel_manager = None + @property def writes_to_stream(self) -> bool: """Indicate this handler does not write to stdout/stderr. diff --git a/src/snakemake_logger_plugin_stomp/formatters.py b/src/snakemake_logger_plugin_stomp/formatters.py index 11971a7..4344b82 100644 --- a/src/snakemake_logger_plugin_stomp/formatters.py +++ b/src/snakemake_logger_plugin_stomp/formatters.py @@ -13,6 +13,32 @@ from typing import Any +# Standard logging.LogRecord attributes that are not Snakemake event payload. +_LOG_RECORD_INTERNAL_ATTRS = { + "args", + "asctime", + "created", + "exc_info", + "exc_text", + "filename", + "funcName", + "levelname", + "levelno", + "lineno", + "module", + "msecs", + "msg", + "name", + "pathname", + "process", + "processName", + "relativeCreated", + "stack_info", + "thread", + "threadName", +} + + class BaseFormatter(ABC): """Abstract base class for STOMP message formatters. @@ -159,3 +185,59 @@ def format(self, record, workflow_metadata): "correlation": {"run_number": os.getenv("SWF_RUN_NUMBER")}, "payload": payload, } + + +class ComprehensiveEventFormatter(BaseFormatter): + """Formatter for logging comprehensive Snakemake event details.""" + + def _normalize_value(self, value: Any) -> Any: + """Convert common Snakemake/LogRecord values into JSON-friendly forms.""" + if value is None or isinstance(value, (str, int, float, bool)): + return value + + if isinstance(value, dict): + return { + str(key): self._normalize_value(item) for key, item in value.items() + } + + if isinstance(value, (list, tuple, set)): + return [self._normalize_value(item) for item in value] + + # Snakemake wildcards and similar objects often expose items(). + if hasattr(value, "items"): + try: + return { + str(key): self._normalize_value(item) + for key, item in value.items() + } + except Exception: # pragma: no cover - defensive fallback + return str(value) + + return str(value) + + def _extract_event_details(self, record: Any) -> dict: + """Extract event-centric fields while excluding logging internals.""" + details = {} + + for attr, value in vars(record).items(): + if attr.startswith("_") or attr in _LOG_RECORD_INTERNAL_ATTRS: + continue + if callable(value): + continue + details[attr] = self._normalize_value(value) + + return details + + def format(self, record, workflow_metadata): + """Format record with rich Snakemake event details.""" + return { + "timestamp": datetime.now(UTC).isoformat(), + "hostname": workflow_metadata.get("hostname"), + "workflow_id": workflow_metadata.get("workflow_id"), + "level": getattr(record, "levelname", "INFO"), + "user": workflow_metadata.get("user"), + "workflow_initiated": workflow_metadata.get("workflow_initiated"), + "workflow_directory": workflow_metadata.get("working_directory"), + "heartbeat_interval_seconds": workflow_metadata.get("heartbeat_interval_seconds"), + "event": self._extract_event_details(record) if record is not None else None, + } diff --git a/src/snakemake_logger_plugin_stomp/ssh_tunnel.py b/src/snakemake_logger_plugin_stomp/ssh_tunnel.py new file mode 100644 index 0000000..20fa736 --- /dev/null +++ b/src/snakemake_logger_plugin_stomp/ssh_tunnel.py @@ -0,0 +1,156 @@ +"""SSH tunnel utilities for STOMP broker connectivity.""" + +import logging +import select +import socketserver +import threading +from typing import Optional + +import paramiko + + +class _ForwardHandler(socketserver.BaseRequestHandler): + """Forward one local TCP connection to the remote broker through SSH.""" + + ssh_transport: paramiko.Transport + remote_host: str + remote_port: int + + def handle(self) -> None: + try: + channel = self.ssh_transport.open_channel( + "direct-tcpip", + (self.remote_host, self.remote_port), + self.request.getpeername(), + ) + except Exception: + return + + if channel is None: + return + + try: + while True: + readable, _, _ = select.select([self.request, channel], [], []) + if self.request in readable: + data = self.request.recv(4096) + if not data: + break + channel.sendall(data) + + if channel in readable: + data = channel.recv(4096) + if not data: + break + self.request.sendall(data) + finally: + channel.close() + self.request.close() + + +class _ForwardServer(socketserver.ThreadingTCPServer): + """Threaded local server that forwards connections over SSH.""" + + daemon_threads = True + allow_reuse_address = True + + +class SSHTunnelManager: + """Manage SSH tunnel lifecycle for STOMP connectivity.""" + + def __init__( + self, + ssh_host: str, + ssh_port: int, + ssh_username: str, + ssh_private_key: str, + ssh_key_passphrase: Optional[str], + remote_host: str, + remote_port: int, + local_bind_port: int, + connect_timeout: int, + logger: logging.Logger, + ) -> None: + self._ssh_host = ssh_host + self._ssh_port = ssh_port + self._ssh_username = ssh_username + self._ssh_private_key = ssh_private_key + self._ssh_key_passphrase = ssh_key_passphrase + self._remote_host = remote_host + self._remote_port = remote_port + self._local_bind_port = local_bind_port + self._connect_timeout = connect_timeout + self._logger = logger + + self._ssh_client: Optional[paramiko.SSHClient] = None + self._forward_server: Optional[_ForwardServer] = None + self._forward_thread: Optional[threading.Thread] = None + self._endpoint: Optional[tuple[str, int]] = None + + def connect(self) -> tuple[str, int]: + """Create a tunnel and return local endpoint for STOMP connection.""" + if self._endpoint is not None: + return self._endpoint + + ssh_client = paramiko.SSHClient() + ssh_client.load_system_host_keys() + ssh_client.set_missing_host_key_policy(paramiko.RejectPolicy()) + ssh_client.connect( + hostname=self._ssh_host, + port=self._ssh_port, + username=self._ssh_username, + key_filename=self._ssh_private_key, + passphrase=self._ssh_key_passphrase, + timeout=self._connect_timeout, + allow_agent=False, + look_for_keys=False, + ) + + transport = ssh_client.get_transport() + if transport is None: + ssh_client.close() + raise RuntimeError("SSH transport was not established") + + transport.set_keepalive(30) + + class Handler(_ForwardHandler): + ssh_transport = transport + remote_host = self._remote_host + remote_port = self._remote_port + + forward_server = _ForwardServer(("127.0.0.1", self._local_bind_port), Handler) + forward_thread = threading.Thread( + target=forward_server.serve_forever, + name="stomp-ssh-forwarder", + daemon=True, + ) + forward_thread.start() + + self._ssh_client = ssh_client + self._forward_server = forward_server + self._forward_thread = forward_thread + self._endpoint = forward_server.server_address + + self._logger.info( + "[STOMP] Opened SSH tunnel " + f"127.0.0.1:{self._endpoint[1]} -> {self._remote_host}:{self._remote_port} " + f"via {self._ssh_host}:{self._ssh_port}" + ) + return self._endpoint + + def close(self) -> None: + """Close local forward server and SSH session.""" + if self._forward_server is not None: + self._forward_server.shutdown() + self._forward_server.server_close() + + if self._forward_thread is not None and self._forward_thread.is_alive(): + self._forward_thread.join(timeout=2) + + if self._ssh_client is not None: + self._ssh_client.close() + + self._forward_server = None + self._forward_thread = None + self._ssh_client = None + self._endpoint = None diff --git a/tests/test_plugin.py b/tests/test_plugin.py index 191a406..a4eeb4b 100644 --- a/tests/test_plugin.py +++ b/tests/test_plugin.py @@ -23,6 +23,7 @@ def __init__(self) -> None: LogHandler, LogHandlerSettings, ) +from snakemake_logger_plugin_stomp.formatters import ComprehensiveEventFormatter class DummyConnection: @@ -100,6 +101,40 @@ def test_settings_defaults(): assert settings.queue == "/queue/snakemake.events" assert settings.user is None assert settings.password is None + assert settings.stream_filter_by_workflow is False + assert settings.consumer_heartbeat_interval == 0 + assert settings.use_ssh_tunnel is False + assert settings.ssh_port == 22 + assert settings.reconnect_attempts == 3 + + +def test_consumer_heartbeat_interval_must_be_non_negative(): + common = MockOutputSettings() + settings = LogHandlerSettings(consumer_heartbeat_interval=-1) + + with pytest.raises(ValueError, match="consumer_heartbeat_interval must be >= 0"): + LogHandler(common_settings=common, settings=settings) + + +def test_ssh_requires_settings_when_enabled(): + common = MockOutputSettings() + settings = LogHandlerSettings(use_ssh_tunnel=True) + + with pytest.raises(ValueError, match="ssh_host is required"): + LogHandler(common_settings=common, settings=settings) + + +def test_ssh_requires_existing_private_key_file(tmp_path): + common = MockOutputSettings() + settings = LogHandlerSettings( + use_ssh_tunnel=True, + ssh_host="jump.example.com", + ssh_username="alice", + ssh_private_key=str(tmp_path / "missing.key"), + ) + + with pytest.raises(ValueError, match="ssh_private_key must point to an existing"): + LogHandler(common_settings=common, settings=settings) def test_formatter_loads_and_falls_back(monkeypatch): @@ -138,6 +173,8 @@ def test_emit_creates_workflow_id_and_sends(): handler.emit(rec) assert handler.workflow_metadata["workflow_id"] is not None + assert handler.workflow_metadata["workflow_initiated"] is not None + assert handler.workflow_metadata["working_directory"] assert len(handler.connection.sent) == 1 sent = handler.connection.sent[0] @@ -220,6 +257,68 @@ def capture_connection(*args, **kwargs): assert conn_calls[0]["heartbeats"] == (5000, 15000) + +def test_ssh_tunnel_endpoint_used_for_stomp_connection(monkeypatch, tmp_path): + common = MockOutputSettings() + conn_calls = [] + tunnel_calls = [] + + class DummyTunnelManager: + def __init__(self, **kwargs): + tunnel_calls.append(("init", kwargs)) + + def connect(self): + tunnel_calls.append(("connect", None)) + return ("127.0.0.1", 19001) + + def close(self): + tunnel_calls.append(("close", None)) + + def capture_connection(*args, **kwargs): + conn_calls.append(kwargs) + return DummyConnection(*args, **kwargs) + + key_file = tmp_path / "id_ed25519" + key_file.write_text("dummy-private-key") + + monkeypatch.setattr("snakemake_logger_plugin_stomp.SSHTunnelManager", DummyTunnelManager) + monkeypatch.setattr("stomp.Connection", capture_connection) + + settings = LogHandlerSettings( + use_ssh_tunnel=True, + ssh_host="jump.example.com", + ssh_username="alice", + ssh_private_key=str(key_file), + ) + handler = LogHandler(common_settings=common, settings=settings) + + assert conn_calls[0]["host_and_ports"] == [("127.0.0.1", 19001)] + assert tunnel_calls[1][0] == "connect" + + handler.close() + assert tunnel_calls[-1][0] == "close" + + +def test_send_recovers_connection_when_disconnected(monkeypatch): + common = MockOutputSettings() + handler = LogHandler(common_settings=common, settings=LogHandlerSettings()) + handler.connection = DummyConnection() + handler.connection._connected = False + + recover_calls = [] + + def fake_recover(): + recover_calls.append(True) + handler.connection._connected = True + return True + + monkeypatch.setattr(handler, "_recover_connection", fake_recover) + + handler._send_to_broker({"event_type": "test"}) + + assert len(recover_calls) == 1 + assert len(handler.connection.sent) == 1 + def test_send_when_not_connected(): """Verify _send_to_broker skips gracefully when disconnected.""" common = MockOutputSettings() @@ -230,6 +329,76 @@ def test_send_when_not_connected(): # Should not raise, should log debug message handler._send_to_broker({"test": "data"}) + +def test_emit_consumer_heartbeat_sends_event_for_consumers(): + common = MockOutputSettings() + settings = LogHandlerSettings(consumer_heartbeat_interval=30) + handler = LogHandler(common_settings=common, settings=settings) + handler.connection = DummyConnection() + handler.workflow_metadata["workflow_id"] = "wf-123" + + handler._emit_consumer_heartbeat() + + assert len(handler.connection.sent) == 1 + payload = json.loads(handler.connection.sent[0]["body"]) + assert payload["event_type"] == "consumer_heartbeat" + assert payload["workflow_id"] == "wf-123" + + +def test_stream_first_send_declares_queue_once(): + common = MockOutputSettings() + settings = LogHandlerSettings( + queue="/queue/snakemake.stream", + use_stream=True, + ) + handler = LogHandler(common_settings=common, settings=settings) + handler.connection = DummyConnection() + + handler.emit(DummyRecord(event="workflow_started", msg="start")) + handler.emit(DummyRecord(event="job_finished", msg="done")) + + first_headers = handler.connection.sent[0]["headers"] + second_headers = handler.connection.sent[1]["headers"] + + assert first_headers["x-queue-type"] == "stream" + assert "x-queue-type" not in second_headers + + +def test_stream_amq_queue_destination_includes_declaration_and_workflow_filter_headers(): + common = MockOutputSettings() + settings = LogHandlerSettings( + queue="/amq/queue/snakemake.stream", + use_stream=True, + stream_filter_by_workflow=True, + ) + handler = LogHandler(common_settings=common, settings=settings) + handler.connection = DummyConnection() + + handler.emit(DummyRecord(event="workflow_started", msg="start")) + + headers = handler.connection.sent[0]["headers"] + assert headers["x-queue-type"] == "stream" + assert headers["x-stream-filter-value"] == handler.workflow_metadata["workflow_id"] + + +def test_stream_filter_by_workflow_skips_filter_header_until_workflow_id_exists(): + common = MockOutputSettings() + settings = LogHandlerSettings( + queue="/queue/snakemake.stream", + use_stream=True, + stream_filter_by_workflow=True, + ) + handler = LogHandler(common_settings=common, settings=settings) + handler.connection = DummyConnection() + + handler.emit(DummyRecord(event="job_started", msg="start")) + + headers = handler.connection.sent[0]["headers"] + assert "x-stream-filter-value" not in headers + + + + def test_listener_on_error_called(monkeypatch): """Verify StompConnectionListener logs broker errors.""" common = MockOutputSettings() @@ -262,11 +431,78 @@ def test_formatter_output_structure(): result = formatter.format(record, metadata) assert result["event_type"] == "job_finished" - assert result["message"] == "Job done" assert result["workflow_id"] == "wf-abc" assert "timestamp" in result # If your formatter adds timestamps +def test_comprehensive_event_formatter_preserves_high_detail_fields(): + formatter = ComprehensiveEventFormatter() + + class Wildcards: + def items(self): + return {"sample": "A", "lane": "L001"}.items() + + record = DummyRecord( + event="job_finished", + msg="Rule completed", + jobid=17, + name="align_reads", + wildcards=Wildcards(), + threads=8, + resources={"mem_mb": 16000}, + benchmark="bench/align_reads.txt", + ) + metadata = {"workflow_id": "wf-123", "hostname": "node-01"} + + result = formatter.format(record, metadata) + + assert result["event_type"] == "job_finished" + assert result["workflow_id"] == "wf-123" + assert result["event"]["jobid"] == 17 + assert result["event"]["wildcards"] == {"sample": "A", "lane": "L001"} + assert result["event"]["benchmark"] == "bench/align_reads.txt" + + +def test_comprehensive_event_formatter_handles_sparse_record(): + formatter = ComprehensiveEventFormatter() + record = DummyRecord(event="workflow_finished", msg="done") + metadata = {"workflow_id": "wf-xyz", "hostname": "node-02"} + + result = formatter.format(record, metadata) + + assert result["event_type"] == "workflow_finished" + assert result["event"]["event"] == "workflow_finished" + + +def test_emit_uses_comprehensive_event_formatter_when_configured(): + common = MockOutputSettings() + settings = LogHandlerSettings( + formatter_class=( + "snakemake_logger_plugin_stomp.formatters.ComprehensiveEventFormatter" + ) + ) + handler = LogHandler(common_settings=common, settings=settings) + handler.connection = DummyConnection() + + rec = DummyRecord( + event="job_started", + msg="Launching rule", + jobid=99, + name="qc", + resources={"mem_mb": 4000}, + extra_context={"attempt": 1}, + ) + handler.emit(rec) + + sent = handler.connection.sent[0] + payload = json.loads(sent["body"]) + + assert payload["event_type"] == "job_started" + assert payload["event"]["jobid"] == 99 + assert payload["event"]["resources"]["mem_mb"] == 4000 + assert payload["event"]["extra_context"]["attempt"] == 1 + + def test_close_handles_disconnect_errors(monkeypatch): """Verify close() logs but doesn't raise on disconnect errors.""" common = MockOutputSettings()