From 2a4caa9c4c039aa45977ec240ecdd89bbef16b90 Mon Sep 17 00:00:00 2001 From: Garry O'Donnell Date: Thu, 18 Apr 2024 13:58:01 +0000 Subject: [PATCH] Return inner span processor method results --- src/graph_energy_scan/telemetry.py | 13 ++++++++----- 1 file changed, 8 insertions(+), 5 deletions(-) diff --git a/src/graph_energy_scan/telemetry.py b/src/graph_energy_scan/telemetry.py index 5e192b4..5bc572c 100644 --- a/src/graph_energy_scan/telemetry.py +++ b/src/graph_energy_scan/telemetry.py @@ -18,17 +18,20 @@ def __init__(self, inner: SpanProcessor, filter_types: list[str]) -> None: super().__init__() def on_start(self, span: Span, parent_context: trace.Context | None = None) -> None: - self.inner.on_start(span, parent_context) + return self.inner.on_start(span, parent_context) def on_end(self, span: ReadableSpan) -> None: - if span.attributes.get("type") not in self.filter_types: - self.inner.on_end(span) + if ( + span.attributes is not None + and span.attributes.get("type") not in self.filter_types + ): + return self.inner.on_end(span) def shutdown(self) -> None: - self.inner.shutdown() + return self.inner.shutdown() def force_flush(self, timeout_millis: int = 30000) -> bool: - self.inner.force_flush(timeout_millis) + return self.inner.force_flush(timeout_millis) def setup_telemetry(otel_collector_url: str):