Add resource attributes to the LogRecord (#26)

Author: jgiannuzzi

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "partial_span_processor"
-version = "0.0.5"
+version = "0.0.6"
 authors = [
     { name = "Mladjan Gadzic", email = "[email protected]" }
 ]

src/partial_span_processor/__init__.py

@@ -12,25 +12,26 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from __future__ import annotations
+
 import base64
 import threading
 import time
 from queue import Queue
-from typing import Optional
+from typing import TYPE_CHECKING
 
-from opentelemetry import context as context_api
 from opentelemetry._logs.severity import SeverityNumber
 from opentelemetry.exporter.otlp.proto.common.trace_encoder import encode_spans
 from opentelemetry.proto.trace.v1 import trace_pb2
 from opentelemetry.sdk._logs import LogData, LogRecord
-from opentelemetry.sdk._logs.export import LogExporter
-from opentelemetry.sdk.trace import (
-  SpanProcessor,
-  Span,
-  ReadableSpan
-)
+from opentelemetry.sdk.trace import ReadableSpan, Span, SpanProcessor
 from opentelemetry.trace import TraceFlags
 
+if TYPE_CHECKING:
+  from opentelemetry import context as context_api
+  from opentelemetry.sdk._logs.export import LogExporter
+  from opentelemetry.sdk.resources import Resource
+
 WORKER_THREAD_NAME = "OtelPartialSpanProcessor"
 
 
@@ -39,12 +40,15 @@ class PartialSpanProcessor(SpanProcessor):
   def __init__(
       self,
       log_exporter: LogExporter,
-      heartbeat_interval_millis: int
-  ):
+      heartbeat_interval_millis: int,
+      resource: Resource | None = None,
+  ) -> None:
     if heartbeat_interval_millis <= 0:
-      raise ValueError("heartbeat_interval_ms must be greater than 0")
+      msg = "heartbeat_interval_ms must be greater than 0"
+      raise ValueError(msg)
     self.log_exporter = log_exporter
     self.heartbeat_interval_millis = heartbeat_interval_millis
+    self.resource = resource
 
     self.active_spans = {}
     self.ended_spans = Queue()
@@ -53,11 +57,11 @@ def __init__(
     self.done = False
     self.condition = threading.Condition(threading.Lock())
     self.worker_thread = threading.Thread(
-      name=WORKER_THREAD_NAME, target=self.worker, daemon=True
+      name=WORKER_THREAD_NAME, target=self.worker, daemon=True,
     )
     self.worker_thread.start()
 
-  def worker(self):
+  def worker(self) -> None:
     while not self.done:
       with self.condition:
         self.condition.wait(self.heartbeat_interval_millis / 1000)
@@ -73,17 +77,17 @@ def worker(self):
 
       self.heartbeat()
 
-  def heartbeat(self):
+  def heartbeat(self) -> None:
     with self.lock:
-      for span_key, span in list(self.active_spans.items()):
+      for span in list(self.active_spans.values()):
         attributes = self.get_heartbeat_attributes()
-        log_data = get_log_data(span, attributes)
+        log_data = self.get_log_data(span, attributes)
         self.log_exporter.export([log_data])
 
-  def on_start(self, span: "Span",
-      parent_context: Optional[context_api.Context] = None) -> None:
+  def on_start(self, span: Span,
+      parent_context: context_api.Context | None = None) -> None:
     attributes = self.get_heartbeat_attributes()
-    log_data = get_log_data(span, attributes)
+    log_data = self.get_log_data(span, attributes)
     self.log_exporter.export([log_data])
 
     span_key = (span.context.trace_id, span.context.span_id)
@@ -92,7 +96,7 @@ def on_start(self, span: "Span",
 
   def on_end(self, span: ReadableSpan) -> None:
     attributes = get_stop_attributes()
-    log_data = get_log_data(span, attributes)
+    log_data = self.get_log_data(span, attributes)
     self.log_exporter.export([log_data])
 
     span_key = (span.context.trace_id, span.context.span_id)
@@ -105,39 +109,38 @@ def shutdown(self) -> None:
       self.condition.notify_all()
     self.worker_thread.join()
 
-  def get_heartbeat_attributes(self):
+  def get_heartbeat_attributes(self) -> dict[str, str]:
     return {
       "partial.event": "heartbeat",
       "partial.frequency": str(self.heartbeat_interval_millis) + "ms",
     }
 
+  def get_log_data(self, span: Span, attributes: dict[str, str]) -> LogData:
+    span_context = Span.get_span_context(span)
+
+    enc_spans = encode_spans([span]).resource_spans
+    traces_data = trace_pb2.TracesData()
+    traces_data.resource_spans.extend(enc_spans)
+    serialized_traces_data = traces_data.SerializeToString()
+
+    log_record = LogRecord(
+      timestamp=time.time_ns(),
+      observed_timestamp=time.time_ns(),
+      trace_id=span_context.trace_id,
+      span_id=span_context.span_id,
+      trace_flags=TraceFlags().get_default(),
+      severity_text="INFO",
+      severity_number=SeverityNumber.INFO,
+      body=base64.b64encode(serialized_traces_data).decode("utf-8"),
+      resource=self.resource,
+      attributes=attributes,
+    )
+    return LogData(
+      log_record=log_record, instrumentation_scope=span.instrumentation_scope,
+    )
 
-def get_stop_attributes():
+
+def get_stop_attributes() -> dict[str, str]:
   return {
     "partial.event": "stop",
   }
-
-
-def get_log_data(span, attributes):
-  span_context = Span.get_span_context(span)
-
-  enc_spans = encode_spans([span]).resource_spans
-  traces_data = trace_pb2.TracesData()
-  traces_data.resource_spans.extend(enc_spans)
-  serialized_traces_data = traces_data.SerializeToString()
-
-  log_record = LogRecord(
-    timestamp=time.time_ns(),
-    observed_timestamp=time.time_ns(),
-    trace_id=span_context.trace_id,
-    span_id=span_context.span_id,
-    trace_flags=TraceFlags().get_default(),
-    severity_text="INFO",
-    severity_number=SeverityNumber.INFO,
-    body=base64.b64encode(serialized_traces_data).decode('utf-8'),
-    attributes=attributes,
-  )
-  log_data = LogData(
-    log_record=log_record, instrumentation_scope=span.instrumentation_scope
-  )
-  return log_data

tests/partial_span_processor/in_memory_log_exporter.py

@@ -12,12 +12,16 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from __future__ import annotations
+
 import threading
 import typing
 
-from opentelemetry.sdk._logs import LogData
 from opentelemetry.sdk._logs.export import LogExporter, LogExportResult
 
+if typing.TYPE_CHECKING:
+  from opentelemetry.sdk._logs import LogData
+
 
 class InMemoryLogExporter(LogExporter):
   """Implementation of :class:`.LogExporter` that stores logs in memory.
@@ -27,7 +31,7 @@ class InMemoryLogExporter(LogExporter):
   :func:`.get_finished_logs` method.
   """
 
-  def __init__(self):
+  def __init__(self) -> None:
     self._logs = []
     self._lock = threading.Lock()
     self._stopped = False
@@ -36,7 +40,7 @@ def clear(self) -> None:
     with self._lock:
       self._logs.clear()
 
-  def get_finished_logs(self) -> typing.Tuple[LogData, ...]:
+  def get_finished_logs(self) -> tuple[LogData, ...]:
     with self._lock:
       return tuple(self._logs)
 

tests/partial_span_processor/test_partial_span_processor.py

@@ -17,27 +17,27 @@
 
 from opentelemetry.sdk.resources import Resource
 from opentelemetry.sdk.trace import TracerProvider
-from opentelemetry.trace import SpanContext, TraceFlags
+from opentelemetry.trace import Span, SpanContext, TraceFlags
 
 from src.partial_span_processor import PartialSpanProcessor
-from tests.partial_span_processor.in_memory_log_exporter import \
-  InMemoryLogExporter
+from tests.partial_span_processor.in_memory_log_exporter import InMemoryLogExporter
 
 
 class TestPartialSpanProcessor(unittest.TestCase):
-  def setUp(self):
+  def setUp(self) -> None:
     # Set up an in-memory log exporter and processor
     self.log_exporter = InMemoryLogExporter()
     self.processor = PartialSpanProcessor(
       log_exporter=self.log_exporter,
       heartbeat_interval_millis=1000,  # 1 second
+      resource=Resource(attributes={"service.name": "test"}),
     )
 
-  def tearDown(self):
+  def tearDown(self) -> None:
     # Shut down the processor
     self.processor.shutdown()
 
-  def create_mock_span(self, trace_id=1, span_id=1):
+  def create_mock_span(self, trace_id: int = 1, span_id: int = 1) -> Span:
     # Create a mock tracer
     tracer_provider = TracerProvider(resource=Resource.create({}))
     tracer = tracer_provider.get_tracer("test_tracer")
@@ -54,36 +54,37 @@ def create_mock_span(self, trace_id=1, span_id=1):
       span._context = span_context  # Modify the span's context for testing
       return span
 
-  def test_on_start(self):
+  def test_on_start(self) -> None:
     # Test the on_start method
     span = self.create_mock_span()
     self.processor.on_start(span)
 
     # Verify the span is added to active_spans
     span_key = (span.context.trace_id, span.context.span_id)
-    self.assertIn(span_key, self.processor.active_spans)
+    assert span_key in self.processor.active_spans
 
     # Verify a log is emitted
     logs = self.log_exporter.get_finished_logs()
-    self.assertEqual(len(logs), 1)
-    self.assertEqual(logs[0].log_record.attributes["partial.event"],
-                     "heartbeat")
+    assert len(logs) == 1
+    assert logs[0].log_record.attributes["partial.event"] == "heartbeat"
+    assert logs[0].log_record.resource.attributes["service.name"] == "test"
 
-  def test_on_end(self):
+  def test_on_end(self) -> None:
     # Test the on_end method
     span = self.create_mock_span()
     self.processor.on_start(span)
     self.processor.on_end(span)
 
     # Verify the span is added to ended_spans
-    self.assertFalse(self.processor.ended_spans.empty())
+    assert not self.processor.ended_spans.empty()
 
     # Verify a log is emitted
     logs = self.log_exporter.get_finished_logs()
-    self.assertEqual(len(logs), 2)
-    self.assertEqual(logs[1].log_record.attributes["partial.event"], "stop")
+    assert len(logs) == 2
+    assert logs[1].log_record.attributes["partial.event"] == "stop"
+    assert logs[0].log_record.resource.attributes["service.name"] == "test"
 
-  def test_heartbeat(self):
+  def test_heartbeat(self) -> None:
     # Test the heartbeat method
     span = self.create_mock_span()
     self.processor.on_start(span)
@@ -93,18 +94,18 @@ def test_heartbeat(self):
     logs = self.log_exporter.get_finished_logs()
 
     # Verify heartbeat logs are emitted
-    self.assertGreaterEqual(len(logs), 2)
-    self.assertEqual(logs[1].log_record.attributes["partial.event"],
-                     "heartbeat")
+    assert len(logs) >= 2
+    assert logs[1].log_record.attributes["partial.event"] == "heartbeat"
+    assert logs[0].log_record.resource.attributes["service.name"] == "test"
 
-  def test_shutdown(self):
+  def test_shutdown(self) -> None:
     # Test the shutdown method
     self.processor.shutdown()
 
     # Verify the worker thread is stopped
-    self.assertTrue(self.processor.done)
+    assert self.processor.done
 
-  def test_worker_thread(self):
+  def test_worker_thread(self) -> None:
     # Test the worker thread processes ended spans
     span = self.create_mock_span()
     self.processor.on_start(span)
@@ -115,7 +116,7 @@ def test_worker_thread(self):
 
     # Verify the span is removed from active_spans
     span_key = (span.context.trace_id, span.context.span_id)
-    self.assertNotIn(span_key, self.processor.active_spans)
+    assert span_key not in self.processor.active_spans
 
 
 if __name__ == "__main__":