From 7bc2c4a609e3f80656acba5337ff2a2848bf47c9 Mon Sep 17 00:00:00 2001 From: Ammon Knaupp Date: Fri, 27 Mar 2026 18:35:56 -0400 Subject: [PATCH 1/9] Add support for RabbitMQ stream queues --- README.md | 28 ++++++- docs/further.md | 46 ++++++++++ docs/intro.md | 20 +++++ src/snakemake_logger_plugin_stomp/__init__.py | 62 +++++++++++++- .../formatters.py | 83 +++++++++++++++++++ tests/test_plugin.py | 40 +++++++++ 6 files changed, 274 insertions(+), 5 deletions(-) diff --git a/README.md b/README.md index e96be43..fe4cc4e 100644 --- a/README.md +++ b/README.md @@ -35,8 +35,34 @@ 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) - Event filtering (include/exclude) - Heartbeat management - - Environment variable support for secrets \ No newline at end of file + - Environment variable support for secrets + +## 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_value: "snakemake-events" +``` + +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..2181f22 100644 --- a/docs/further.md +++ b/docs/further.md @@ -198,6 +198,52 @@ 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_value` +to set a static filter value applied to each published event: + +```yaml +logger: + stomp: + queue: "/queue/snakemake.stream" + use_stream: true + stream_filter_value: "snakemake-workflow-logs" +``` + +### 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..ec5c9ac 100644 --- a/docs/intro.md +++ b/docs/intro.md @@ -5,6 +5,7 @@ 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 - **Event filtering** - include/exclude specific event types to reduce noise @@ -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_value: "snakemake-events" +``` + +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/src/snakemake_logger_plugin_stomp/__init__.py b/src/snakemake_logger_plugin_stomp/__init__.py index 713e6d4..4da4832 100644 --- a/src/snakemake_logger_plugin_stomp/__init__.py +++ b/src/snakemake_logger_plugin_stomp/__init__.py @@ -125,6 +125,20 @@ 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_value: Optional[str] = field( + default=None, + metadata={ + "help": "Static RabbitMQ x-stream-filter-value for outbound messages", + "env_var": False, + }, + ) class StompConnectionListener(stomp.ConnectionListener): @@ -184,6 +198,7 @@ def __post_init__(self) -> None: # STOMP connection handle self.connection = None + self._stream_declared = False # Workflow tracking metadata self.workflow_metadata = { @@ -191,6 +206,8 @@ def __post_init__(self) -> None: "hostname": socket.gethostname(), } + self._validate_stream_settings() + # Load and initialize the configured formatter self.formatter_instance = self._init_formatter() @@ -285,6 +302,42 @@ def _connect_to_broker(self): f"STOMP logging required but broker unavailable: {e}" ) from e + def _validate_stream_settings(self) -> None: + """Validate RabbitMQ stream-specific settings.""" + if not self.settings.use_stream: + return + + if self.settings.stream_filter_value is not None and not str(self.settings.stream_filter_value).strip(): + raise ValueError("stream_filter_value cannot be empty") + + 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, event_data: dict) -> dict[str, str]: + """Build per-message RabbitMQ stream headers.""" + if not self.settings.use_stream or self.settings.stream_filter_value is None: + return {} + + return {"x-stream-filter-value": self.settings.stream_filter_value} + + def _build_send_headers(self, event_data: dict) -> 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(event_data)) + return headers + def _send_to_broker(self, event_data: dict): """Send formatted event data to STOMP broker. @@ -300,15 +353,16 @@ def _send_to_broker(self, event_data: dict): try: # Send message with JSON body and appropriate headers + headers = self._build_send_headers(event_data) 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')}" ) diff --git a/src/snakemake_logger_plugin_stomp/formatters.py b/src/snakemake_logger_plugin_stomp/formatters.py index 11971a7..2c2e531 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,60 @@ def format(self, record, workflow_metadata): "correlation": {"run_number": os.getenv("SWF_RUN_NUMBER")}, "payload": payload, } + + +class SnakemakeEventFormatter(BaseFormatter): + """Human-focused formatter with rich Snakemake event details. + + This formatter keeps a readable top-level summary while preserving as much + event data from the underlying log record as possible. + """ + + 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"), + "event_type": str(getattr(record, "event", "UNKNOWN")), + "level": getattr(record, "levelname", "INFO"), + "event": self._extract_event_details(record), + } diff --git a/tests/test_plugin.py b/tests/test_plugin.py index 191a406..46ec9a5 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 SnakemakeEventFormatter class DummyConnection: @@ -230,6 +231,45 @@ def test_send_when_not_connected(): # Should not raise, should log debug message handler._send_to_broker({"test": "data"}) + +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_filter_headers(): + common = MockOutputSettings() + settings = LogHandlerSettings( + queue="/amq/queue/snakemake.stream", + use_stream=True, + stream_filter_value="job-events", + ) + 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 headers["x-queue-type"] == "stream" + assert headers["x-stream-filter-value"] == "job-events" + + + + def test_listener_on_error_called(monkeypatch): """Verify StompConnectionListener logs broker errors.""" common = MockOutputSettings() From b0dd8b622273d5eef44f66ca474025d3a71147b2 Mon Sep 17 00:00:00 2001 From: Ammon Knaupp Date: Fri, 27 Mar 2026 18:49:53 -0400 Subject: [PATCH 2/9] Add ComprehensiveEventFormatter --- README.md | 2 +- docs/further.md | 34 +++++++++ docs/intro.md | 2 +- .../formatters.py | 8 +-- tests/test_plugin.py | 71 ++++++++++++++++++- 5 files changed, 107 insertions(+), 10 deletions(-) diff --git a/README.md b/README.md index fe4cc4e..6e7bd54 100644 --- a/README.md +++ b/README.md @@ -37,7 +37,7 @@ logger: - 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 diff --git a/docs/further.md b/docs/further.md index 2181f22..03cbd06 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: diff --git a/docs/intro.md b/docs/intro.md index ec5c9ac..d1d7d74 100644 --- a/docs/intro.md +++ b/docs/intro.md @@ -7,7 +7,7 @@ Stream Snakemake workflow events to STOMP message brokers (ActiveMQ, RabbitMQ, A - **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 diff --git a/src/snakemake_logger_plugin_stomp/formatters.py b/src/snakemake_logger_plugin_stomp/formatters.py index 2c2e531..493921d 100644 --- a/src/snakemake_logger_plugin_stomp/formatters.py +++ b/src/snakemake_logger_plugin_stomp/formatters.py @@ -187,12 +187,8 @@ def format(self, record, workflow_metadata): } -class SnakemakeEventFormatter(BaseFormatter): - """Human-focused formatter with rich Snakemake event details. - - This formatter keeps a readable top-level summary while preserving as much - event data from the underlying log record as possible. - """ +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.""" diff --git a/tests/test_plugin.py b/tests/test_plugin.py index 46ec9a5..1159ee2 100644 --- a/tests/test_plugin.py +++ b/tests/test_plugin.py @@ -23,7 +23,7 @@ def __init__(self) -> None: LogHandler, LogHandlerSettings, ) -from snakemake_logger_plugin_stomp.formatters import SnakemakeEventFormatter +from snakemake_logger_plugin_stomp.formatters import ComprehensiveEventFormatter class DummyConnection: @@ -302,11 +302,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() From f661c8391da7bc58a2c4d05e9164f2769674d2e3 Mon Sep 17 00:00:00 2001 From: Ammon Knaupp Date: Mon, 30 Mar 2026 14:41:23 -0400 Subject: [PATCH 3/9] feat: add application-level consumer heartbeat events and validation --- README.md | 16 ++++ src/snakemake_logger_plugin_stomp/__init__.py | 92 +++++++++++++++++++ tests/test_plugin.py | 29 ++++++ 3 files changed, 137 insertions(+) diff --git a/README.md b/README.md index 6e7bd54..0345553 100644 --- a/README.md +++ b/README.md @@ -40,8 +40,24 @@ logger: - Configurable formatters (default + JLab SWF schema + ComprehensiveEventFormatter) - Event filtering (include/exclude) - Heartbeat management + - 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 diff --git a/src/snakemake_logger_plugin_stomp/__init__.py b/src/snakemake_logger_plugin_stomp/__init__.py index 4da4832..bb8cdef 100644 --- a/src/snakemake_logger_plugin_stomp/__init__.py +++ b/src/snakemake_logger_plugin_stomp/__init__.py @@ -11,8 +11,10 @@ import logging as py_logging import socket import ssl +import threading import uuid from dataclasses import dataclass, field +from datetime import UTC, datetime from typing import Optional import stomp @@ -139,6 +141,13 @@ class LogHandlerSettings(LogHandlerSettingsBase): "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, + }, + ) class StompConnectionListener(stomp.ConnectionListener): @@ -168,6 +177,18 @@ def on_disconnected(self): self.logger.warning("[STOMP] Disconnected from broker") +class _ConsumerHeartbeatRecord: + """Lightweight synthetic record for consumer heartbeat events.""" + + event = "consumer_heartbeat" + levelname = "INFO" + msg = "Consumer heartbeat" + heartbeat = True + + def getMessage(self): + return self.msg + + class LogHandler(LogHandlerBase): """Snakemake logger handler for STOMP message broker integration. @@ -199,14 +220,18 @@ def __post_init__(self) -> None: # STOMP connection handle self.connection = None 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_start_timestamp": None, "hostname": socket.gethostname(), } self._validate_stream_settings() + self._validate_consumer_heartbeat_settings() # Load and initialize the configured formatter self.formatter_instance = self._init_formatter() @@ -224,6 +249,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. @@ -310,6 +336,65 @@ def _validate_stream_settings(self) -> None: if self.settings.stream_filter_value is not None and not str(self.settings.stream_filter_value).strip(): raise ValueError("stream_filter_value cannot be empty") + 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 _start_consumer_heartbeat_loop(self) -> None: + """Start background loop that emits consumer heartbeat events.""" + if self.settings.consumer_heartbeat_interval <= 0: + return + + 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 events at a configured interval.""" + interval = self.settings.consumer_heartbeat_interval + + 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( + _ConsumerHeartbeatRecord(), self.workflow_metadata + ) + except Exception as e: + self._internal_logger.error( + f"[STOMP] Formatter error for consumer heartbeat event: {e}" + ) + return + + # Ensure heartbeat semantics are explicit even with custom formatters. + if isinstance(event_data, dict): + event_data.setdefault("event_type", "consumer_heartbeat") + event_data.setdefault("heartbeat", True) + event_data.setdefault( + "workflow_id", + self.workflow_metadata.get("workflow_id"), + ) + event_data.setdefault( + "workflow_start_timestamp", + self.workflow_metadata.get("workflow_start_timestamp"), + ) + event_data.setdefault( + "heartbeat_interval_seconds", + self.settings.consumer_heartbeat_interval, + ) + + 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 @@ -420,6 +505,9 @@ def emit(self, record): and not self.workflow_metadata["workflow_id"] ): self.workflow_metadata["workflow_id"] = str(uuid.uuid4()) + self.workflow_metadata["workflow_start_timestamp"] = datetime.now( + UTC + ).isoformat() self._internal_logger.info( f"[STOMP] Generated workflow ID: {self.workflow_metadata['workflow_id']}" ) @@ -441,6 +529,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() diff --git a/tests/test_plugin.py b/tests/test_plugin.py index 1159ee2..4f8099f 100644 --- a/tests/test_plugin.py +++ b/tests/test_plugin.py @@ -101,6 +101,15 @@ def test_settings_defaults(): assert settings.queue == "/queue/snakemake.events" assert settings.user is None assert settings.password is None + assert settings.consumer_heartbeat_interval == 0 + + +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_formatter_loads_and_falls_back(monkeypatch): @@ -139,6 +148,7 @@ 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_start_timestamp"] is not None assert len(handler.connection.sent) == 1 sent = handler.connection.sent[0] @@ -232,6 +242,25 @@ def test_send_when_not_connected(): 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.workflow_metadata["workflow_start_timestamp"] = "2026-03-30T12:00:00+00:00" + + 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["heartbeat"] is True + assert payload["workflow_id"] == "wf-123" + assert payload["workflow_start_timestamp"] == "2026-03-30T12:00:00+00:00" + assert payload["heartbeat_interval_seconds"] == 30 + + def test_stream_first_send_declares_queue_once(): common = MockOutputSettings() settings = LogHandlerSettings( From 6a38c15e7f2d89308d508b5abcf3e9aa6bbe49d7 Mon Sep 17 00:00:00 2001 From: Ammon Knaupp Date: Mon, 30 Mar 2026 14:41:50 -0400 Subject: [PATCH 4/9] fix: stream filter should use workflow ID for RabbitMQ messages --- README.md | 2 +- docs/further.md | 7 +++--- docs/intro.md | 2 +- src/snakemake_logger_plugin_stomp/__init__.py | 25 ++++++++++--------- tests/test_plugin.py | 25 ++++++++++++++++--- 5 files changed, 40 insertions(+), 21 deletions(-) diff --git a/README.md b/README.md index 0345553..541cf78 100644 --- a/README.md +++ b/README.md @@ -74,7 +74,7 @@ logger: password: "${STOMP_PASSWORD}" queue: "/queue/snakemake.stream" use_stream: true - stream_filter_value: "snakemake-events" + stream_filter_by_workflow: true ``` Notes: diff --git a/docs/further.md b/docs/further.md index 03cbd06..b3f27be 100644 --- a/docs/further.md +++ b/docs/further.md @@ -262,15 +262,16 @@ logger: ### Stream Filter Headers -RabbitMQ streams support outbound filter values on published messages. Use `stream_filter_value` -to set a static filter value applied to each published event: +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_value: "snakemake-workflow-logs" + stream_filter_by_workflow: true ``` ### Constraints diff --git a/docs/intro.md b/docs/intro.md index d1d7d74..3c140f2 100644 --- a/docs/intro.md +++ b/docs/intro.md @@ -62,7 +62,7 @@ logger: password: "${STOMP_PASSWORD}" queue: "/queue/snakemake.stream" use_stream: true - stream_filter_value: "snakemake-events" + stream_filter_by_workflow: true ``` Use `/queue/` to create a new stream on first publish. Use `/amq/queue/` diff --git a/src/snakemake_logger_plugin_stomp/__init__.py b/src/snakemake_logger_plugin_stomp/__init__.py index bb8cdef..40fef2d 100644 --- a/src/snakemake_logger_plugin_stomp/__init__.py +++ b/src/snakemake_logger_plugin_stomp/__init__.py @@ -134,10 +134,10 @@ class LogHandlerSettings(LogHandlerSettingsBase): "env_var": False, }, ) - stream_filter_value: Optional[str] = field( - default=None, + stream_filter_by_workflow: bool = field( + default=False, metadata={ - "help": "Static RabbitMQ x-stream-filter-value for outbound messages", + "help": "Set RabbitMQ x-stream-filter-value header to workflow_id on outbound messages", "env_var": False, }, ) @@ -333,9 +333,6 @@ def _validate_stream_settings(self) -> None: if not self.settings.use_stream: return - if self.settings.stream_filter_value is not None and not str(self.settings.stream_filter_value).strip(): - raise ValueError("stream_filter_value cannot be empty") - def _validate_consumer_heartbeat_settings(self) -> None: """Validate consumer heartbeat settings.""" if self.settings.consumer_heartbeat_interval < 0: @@ -406,21 +403,25 @@ def _build_stream_declare_headers(self) -> dict[str, str]: return {"x-queue-type": "stream"} - def _build_stream_publish_headers(self, event_data: dict) -> dict[str, str]: + def _build_stream_publish_headers(self) -> dict[str, str]: """Build per-message RabbitMQ stream headers.""" - if not self.settings.use_stream or self.settings.stream_filter_value is None: + 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": self.settings.stream_filter_value} + return {"x-stream-filter-value": workflow_id} - def _build_send_headers(self, event_data: dict) -> dict[str, str]: + 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(event_data)) + headers.update(self._build_stream_publish_headers()) return headers def _send_to_broker(self, event_data: dict): @@ -438,7 +439,7 @@ def _send_to_broker(self, event_data: dict): try: # Send message with JSON body and appropriate headers - headers = self._build_send_headers(event_data) + headers = self._build_send_headers() self.connection.send( body=json.dumps(event_data, default=str), destination=self.settings.queue, diff --git a/tests/test_plugin.py b/tests/test_plugin.py index 4f8099f..0a08e03 100644 --- a/tests/test_plugin.py +++ b/tests/test_plugin.py @@ -101,6 +101,7 @@ 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 @@ -280,21 +281,37 @@ def test_stream_first_send_declares_queue_once(): assert "x-queue-type" not in second_headers -def test_stream_amq_queue_destination_includes_declaration_and_filter_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_value="job-events", + stream_filter_by_workflow=True, ) handler = LogHandler(common_settings=common, settings=settings) handler.connection = DummyConnection() - handler.emit(DummyRecord(event="job_started", msg="start")) + 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"] == "job-events" + 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 From 4adfff4e279cfb70b15e90faeb417a7b5273391d Mon Sep 17 00:00:00 2001 From: Ammon Knaupp Date: Thu, 2 Apr 2026 17:40:01 -0400 Subject: [PATCH 5/9] feat: include working directory and user in heartbeat payload --- src/snakemake_logger_plugin_stomp/__init__.py | 12 ++++++++++++ tests/test_plugin.py | 2 ++ 2 files changed, 14 insertions(+) diff --git a/src/snakemake_logger_plugin_stomp/__init__.py b/src/snakemake_logger_plugin_stomp/__init__.py index 40fef2d..b392ad8 100644 --- a/src/snakemake_logger_plugin_stomp/__init__.py +++ b/src/snakemake_logger_plugin_stomp/__init__.py @@ -9,6 +9,8 @@ import importlib import json import logging as py_logging +import os +from pathlib import Path import socket import ssl import threading @@ -228,6 +230,8 @@ def __post_init__(self) -> None: "workflow_id": None, # Generated on WORKFLOW_STARTED event "workflow_start_timestamp": None, "hostname": socket.gethostname(), + "working_directory": str(Path.cwd()), + "user": os.getenv("USER") or os.getenv("USERNAME") or "unknown", } self._validate_stream_settings() @@ -389,6 +393,14 @@ def _emit_consumer_heartbeat(self) -> None: "heartbeat_interval_seconds", self.settings.consumer_heartbeat_interval, ) + event_data.setdefault( + "working_directory", + self.workflow_metadata.get("working_directory"), + ) + event_data.setdefault( + "producer", + self.workflow_metadata.get("producer"), + ) self._send_to_broker(event_data) diff --git a/tests/test_plugin.py b/tests/test_plugin.py index 0a08e03..85c590e 100644 --- a/tests/test_plugin.py +++ b/tests/test_plugin.py @@ -150,6 +150,7 @@ def test_emit_creates_workflow_id_and_sends(): assert handler.workflow_metadata["workflow_id"] is not None assert handler.workflow_metadata["workflow_start_timestamp"] is not None + assert handler.workflow_metadata["working_directory_name"] assert len(handler.connection.sent) == 1 sent = handler.connection.sent[0] @@ -260,6 +261,7 @@ def test_emit_consumer_heartbeat_sends_event_for_consumers(): assert payload["workflow_id"] == "wf-123" assert payload["workflow_start_timestamp"] == "2026-03-30T12:00:00+00:00" assert payload["heartbeat_interval_seconds"] == 30 + assert payload["working_directory"] == handler.workflow_metadata["working_directory"] def test_stream_first_send_declares_queue_once(): From 95f8df1ca1a7de0e7eae9bcf2b5c6ad68eb519ad Mon Sep 17 00:00:00 2001 From: Ammon Knaupp <62727673+adknaupp@users.noreply.github.com> Date: Thu, 2 Apr 2026 18:10:10 -0400 Subject: [PATCH 6/9] Missed change belonging to previous commit --- src/snakemake_logger_plugin_stomp/__init__.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/snakemake_logger_plugin_stomp/__init__.py b/src/snakemake_logger_plugin_stomp/__init__.py index b392ad8..25ef7d6 100644 --- a/src/snakemake_logger_plugin_stomp/__init__.py +++ b/src/snakemake_logger_plugin_stomp/__init__.py @@ -398,8 +398,8 @@ def _emit_consumer_heartbeat(self) -> None: self.workflow_metadata.get("working_directory"), ) event_data.setdefault( - "producer", - self.workflow_metadata.get("producer"), + "user", + self.workflow_metadata.get("user"), ) self._send_to_broker(event_data) From f60b3de42c5a2a13f0a3a5a7e2974bcc7307b6a6 Mon Sep 17 00:00:00 2001 From: Ammon Knaupp Date: Sun, 5 Apr 2026 00:18:53 -0400 Subject: [PATCH 7/9] fix: initialize workflow_id for consumer heartbeat and simplify ComprehensiveEventFormatter.format --- src/snakemake_logger_plugin_stomp/__init__.py | 60 +++++-------------- .../formatters.py | 4 +- tests/test_plugin.py | 9 +-- 3 files changed, 20 insertions(+), 53 deletions(-) diff --git a/src/snakemake_logger_plugin_stomp/__init__.py b/src/snakemake_logger_plugin_stomp/__init__.py index 25ef7d6..9a06a57 100644 --- a/src/snakemake_logger_plugin_stomp/__init__.py +++ b/src/snakemake_logger_plugin_stomp/__init__.py @@ -179,18 +179,6 @@ def on_disconnected(self): self.logger.warning("[STOMP] Disconnected from broker") -class _ConsumerHeartbeatRecord: - """Lightweight synthetic record for consumer heartbeat events.""" - - event = "consumer_heartbeat" - levelname = "INFO" - msg = "Consumer heartbeat" - heartbeat = True - - def getMessage(self): - return self.msg - - class LogHandler(LogHandlerBase): """Snakemake logger handler for STOMP message broker integration. @@ -227,11 +215,12 @@ def __post_init__(self) -> None: # Workflow tracking metadata self.workflow_metadata = { - "workflow_id": None, # Generated on WORKFLOW_STARTED event - "workflow_start_timestamp": None, + "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() @@ -346,6 +335,12 @@ 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, @@ -359,9 +354,14 @@ def _start_consumer_heartbeat_loop(self) -> None: ) def _consumer_heartbeat_loop(self) -> None: - """Emit consumer heartbeat events at a configured interval.""" + """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() @@ -369,7 +369,7 @@ def _emit_consumer_heartbeat(self) -> None: """Emit one synthetic heartbeat event intended for downstream consumers.""" try: event_data = self.formatter_instance.format( - _ConsumerHeartbeatRecord(), self.workflow_metadata + self.workflow_metadata, self.workflow_metadata ) except Exception as e: self._internal_logger.error( @@ -377,31 +377,6 @@ def _emit_consumer_heartbeat(self) -> None: ) return - # Ensure heartbeat semantics are explicit even with custom formatters. - if isinstance(event_data, dict): - event_data.setdefault("event_type", "consumer_heartbeat") - event_data.setdefault("heartbeat", True) - event_data.setdefault( - "workflow_id", - self.workflow_metadata.get("workflow_id"), - ) - event_data.setdefault( - "workflow_start_timestamp", - self.workflow_metadata.get("workflow_start_timestamp"), - ) - event_data.setdefault( - "heartbeat_interval_seconds", - self.settings.consumer_heartbeat_interval, - ) - event_data.setdefault( - "working_directory", - self.workflow_metadata.get("working_directory"), - ) - event_data.setdefault( - "user", - self.workflow_metadata.get("user"), - ) - self._send_to_broker(event_data) def _should_declare_stream_on_send(self) -> bool: @@ -518,9 +493,6 @@ def emit(self, record): and not self.workflow_metadata["workflow_id"] ): self.workflow_metadata["workflow_id"] = str(uuid.uuid4()) - self.workflow_metadata["workflow_start_timestamp"] = datetime.now( - UTC - ).isoformat() self._internal_logger.info( f"[STOMP] Generated workflow ID: {self.workflow_metadata['workflow_id']}" ) diff --git a/src/snakemake_logger_plugin_stomp/formatters.py b/src/snakemake_logger_plugin_stomp/formatters.py index 493921d..91d192f 100644 --- a/src/snakemake_logger_plugin_stomp/formatters.py +++ b/src/snakemake_logger_plugin_stomp/formatters.py @@ -231,10 +231,10 @@ def _extract_event_details(self, record: Any) -> dict: def format(self, record, workflow_metadata): """Format record with rich Snakemake event details.""" return { + "heartbeat": record is workflow_metadata, "timestamp": datetime.now(UTC).isoformat(), "hostname": workflow_metadata.get("hostname"), "workflow_id": workflow_metadata.get("workflow_id"), - "event_type": str(getattr(record, "event", "UNKNOWN")), "level": getattr(record, "levelname", "INFO"), - "event": self._extract_event_details(record), + "event": workflow_metadata if record is workflow_metadata else self._extract_event_details(record), } diff --git a/tests/test_plugin.py b/tests/test_plugin.py index 85c590e..02c19f2 100644 --- a/tests/test_plugin.py +++ b/tests/test_plugin.py @@ -149,8 +149,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_start_timestamp"] is not None - assert handler.workflow_metadata["working_directory_name"] + 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] @@ -250,18 +250,13 @@ def test_emit_consumer_heartbeat_sends_event_for_consumers(): handler = LogHandler(common_settings=common, settings=settings) handler.connection = DummyConnection() handler.workflow_metadata["workflow_id"] = "wf-123" - handler.workflow_metadata["workflow_start_timestamp"] = "2026-03-30T12:00:00+00:00" 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["heartbeat"] is True assert payload["workflow_id"] == "wf-123" - assert payload["workflow_start_timestamp"] == "2026-03-30T12:00:00+00:00" - assert payload["heartbeat_interval_seconds"] == 30 - assert payload["working_directory"] == handler.workflow_metadata["working_directory"] def test_stream_first_send_declares_queue_once(): From ee878a92f6d8de0a3a3334b6187ceaa39d70efdc Mon Sep 17 00:00:00 2001 From: Ammon Knaupp Date: Tue, 7 Apr 2026 13:35:00 -0400 Subject: [PATCH 8/9] Move workflow metadata to top level --- src/snakemake_logger_plugin_stomp/__init__.py | 2 +- src/snakemake_logger_plugin_stomp/formatters.py | 7 +++++-- 2 files changed, 6 insertions(+), 3 deletions(-) diff --git a/src/snakemake_logger_plugin_stomp/__init__.py b/src/snakemake_logger_plugin_stomp/__init__.py index 9a06a57..35b5020 100644 --- a/src/snakemake_logger_plugin_stomp/__init__.py +++ b/src/snakemake_logger_plugin_stomp/__init__.py @@ -369,7 +369,7 @@ def _emit_consumer_heartbeat(self) -> None: """Emit one synthetic heartbeat event intended for downstream consumers.""" try: event_data = self.formatter_instance.format( - self.workflow_metadata, self.workflow_metadata + None, self.workflow_metadata ) except Exception as e: self._internal_logger.error( diff --git a/src/snakemake_logger_plugin_stomp/formatters.py b/src/snakemake_logger_plugin_stomp/formatters.py index 91d192f..4344b82 100644 --- a/src/snakemake_logger_plugin_stomp/formatters.py +++ b/src/snakemake_logger_plugin_stomp/formatters.py @@ -231,10 +231,13 @@ def _extract_event_details(self, record: Any) -> dict: def format(self, record, workflow_metadata): """Format record with rich Snakemake event details.""" return { - "heartbeat": record is workflow_metadata, "timestamp": datetime.now(UTC).isoformat(), "hostname": workflow_metadata.get("hostname"), "workflow_id": workflow_metadata.get("workflow_id"), "level": getattr(record, "levelname", "INFO"), - "event": workflow_metadata if record is workflow_metadata else self._extract_event_details(record), + "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, } From 864acc3cb10d45464a2b2f087d021070b35245e5 Mon Sep 17 00:00:00 2001 From: Ammon Knaupp Date: Fri, 10 Apr 2026 11:17:44 -0400 Subject: [PATCH 9/9] feat: add SSH tunnel support for STOMP broker connectivity --- pyproject.toml | 1 + src/snakemake_logger_plugin_stomp/__init__.py | 212 +++++++++++++++++- .../ssh_tunnel.py | 156 +++++++++++++ tests/test_plugin.py | 86 +++++++ 4 files changed, 447 insertions(+), 8 deletions(-) create mode 100644 src/snakemake_logger_plugin_stomp/ssh_tunnel.py 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 35b5020..dfd5c02 100644 --- a/src/snakemake_logger_plugin_stomp/__init__.py +++ b/src/snakemake_logger_plugin_stomp/__init__.py @@ -14,6 +14,7 @@ import socket import ssl import threading +import time import uuid from dataclasses import dataclass, field from datetime import UTC, datetime @@ -25,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 @@ -150,6 +152,83 @@ class LogHandlerSettings(LogHandlerSettingsBase): "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): @@ -209,6 +288,9 @@ 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 @@ -225,6 +307,8 @@ def __post_init__(self) -> None: 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() @@ -269,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, @@ -309,17 +420,24 @@ 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.""" @@ -331,6 +449,70 @@ def _validate_consumer_heartbeat_settings(self) -> None: 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: @@ -419,10 +601,14 @@ 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 @@ -525,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/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 02c19f2..a4eeb4b 100644 --- a/tests/test_plugin.py +++ b/tests/test_plugin.py @@ -103,6 +103,9 @@ def test_settings_defaults(): 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(): @@ -113,6 +116,27 @@ def test_consumer_heartbeat_interval_must_be_non_negative(): 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): common = MockOutputSettings() @@ -233,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()