Move encoding of spans to opentelemetry-exporter-otlp-proto-common

Author: jack-burridge-cfh

exporter/opentelemetry-exporter-otlp-proto-http/src/opentelemetry/exporter/otlp/proto/http/trace_exporter/encoder/__init__.py renamed to exporter/opentelemetry-exporter-otlp-proto-common/src/opentelemetry/exporter/otlp/proto/common/trace_encoder/__init__.py

@@ -13,19 +13,18 @@
 # limitations under the License.
 
 import logging
-from collections import abc
-from typing import Any, List, Optional, Sequence
-
+from collections import defaultdict
+from typing import List, Optional, Sequence
+
+from opentelemetry.exporter.otlp.proto.common import (
+    _encode_trace_id,
+    _encode_span_id,
+    _encode_key_value,
+    _encode_instrumentation_scope,
+)
 from opentelemetry.proto.collector.trace.v1.trace_service_pb2 import (
     ExportTraceServiceRequest as PB2ExportTraceServiceRequest,
 )
-from opentelemetry.proto.common.v1.common_pb2 import AnyValue as PB2AnyValue
-from opentelemetry.proto.common.v1.common_pb2 import (
-    ArrayValue as PB2ArrayValue,
-)
-from opentelemetry.proto.common.v1.common_pb2 import (
-    InstrumentationScope as PB2InstrumentationScope,
-)
 from opentelemetry.proto.common.v1.common_pb2 import KeyValue as PB2KeyValue
 from opentelemetry.proto.resource.v1.resource_pb2 import (
     Resource as PB2Resource,
@@ -38,10 +37,8 @@
 )
 from opentelemetry.proto.trace.v1.trace_pb2 import Span as PB2SPan
 from opentelemetry.proto.trace.v1.trace_pb2 import Status as PB2Status
-from opentelemetry.sdk.trace import Event
-from opentelemetry.sdk.util.instrumentation import InstrumentationScope
+from opentelemetry.sdk.trace import Event, ReadableSpan
 from opentelemetry.sdk.trace import Resource
-from opentelemetry.sdk.trace import Span as SDKSpan
 from opentelemetry.trace import Link
 from opentelemetry.trace import SpanKind
 from opentelemetry.trace.span import SpanContext, TraceState, Status
@@ -59,20 +56,16 @@
 _logger = logging.getLogger(__name__)
 
 
-class _ProtobufEncoder:
-    @classmethod
-    def serialize(cls, sdk_spans: Sequence[SDKSpan]) -> str:
-        return cls.encode(sdk_spans).SerializeToString()
-
-    @staticmethod
-    def encode(sdk_spans: Sequence[SDKSpan]) -> PB2ExportTraceServiceRequest:
-        return PB2ExportTraceServiceRequest(
-            resource_spans=_encode_resource_spans(sdk_spans)
-        )
+def encode_spans(
+    sdk_spans: Sequence[ReadableSpan],
+) -> PB2ExportTraceServiceRequest:
+    return PB2ExportTraceServiceRequest(
+        resource_spans=_encode_resource_spans(sdk_spans)
+    )
 
 
 def _encode_resource_spans(
-    sdk_spans: Sequence[SDKSpan],
+    sdk_spans: Sequence[ReadableSpan],
 ) -> List[PB2ResourceSpans]:
     # We need to inspect the spans and group + structure them as:
     #
@@ -85,25 +78,14 @@ def _encode_resource_spans(
     #
     # Second loop encodes the data into Protobuf format.
     #
-    sdk_resource_spans = {}
+    sdk_resource_spans = defaultdict(lambda: defaultdict(list))
 
     for sdk_span in sdk_spans:
         sdk_resource = sdk_span.resource
         sdk_instrumentation = sdk_span.instrumentation_scope or None
         pb2_span = _encode_span(sdk_span)
 
-        if sdk_resource not in sdk_resource_spans.keys():
-            sdk_resource_spans[sdk_resource] = {
-                sdk_instrumentation: [pb2_span]
-            }
-        elif (
-            sdk_instrumentation not in sdk_resource_spans[sdk_resource].keys()
-        ):
-            sdk_resource_spans[sdk_resource][sdk_instrumentation] = [pb2_span]
-        else:
-            sdk_resource_spans[sdk_resource][sdk_instrumentation].append(
-                pb2_span
-            )
+        sdk_resource_spans[sdk_resource][sdk_instrumentation].append(pb2_span)
 
     pb2_resource_spans = []
 
@@ -126,7 +108,7 @@ def _encode_resource_spans(
     return pb2_resource_spans
 
 
-def _encode_span(sdk_span: SDKSpan) -> PB2SPan:
+def _encode_span(sdk_span: ReadableSpan) -> PB2SPan:
     span_context = sdk_span.get_span_context()
     return PB2SPan(
         trace_id=_encode_trace_id(span_context.trace_id),
@@ -141,6 +123,9 @@ def _encode_span(sdk_span: SDKSpan) -> PB2SPan:
         events=_encode_events(sdk_span.events),
         links=_encode_links(sdk_span.links),
         status=_encode_status(sdk_span.status),
+        dropped_attributes_count=sdk_span.dropped_attributes,
+        dropped_events_count=sdk_span.dropped_events,
+        dropped_links_count=sdk_span.dropped_links,
     )
 
 
@@ -154,6 +139,7 @@ def _encode_events(
             encoded_event = PB2SPan.Event(
                 name=event.name,
                 time_unix_nano=event.timestamp,
+                dropped_attributes_count=event.attributes.dropped,
             )
             for key, value in event.attributes.items():
                 try:
@@ -167,14 +153,15 @@ def _encode_events(
     return pb2_events
 
 
-def _encode_links(links: List[Link]) -> List[PB2SPan.Link]:
+def _encode_links(links: Sequence[Link]) -> Sequence[PB2SPan.Link]:
     pb2_links = None
     if links:
         pb2_links = []
         for link in links:
             encoded_link = PB2SPan.Link(
                 trace_id=_encode_trace_id(link.context.trace_id),
                 span_id=_encode_span_id(link.context.span_id),
+                dropped_attributes_count=link.attributes.dropped,
             )
             for key, value in link.attributes.items():
                 try:
@@ -208,11 +195,9 @@ def _encode_trace_state(trace_state: TraceState) -> Optional[str]:
 
 
 def _encode_parent_id(context: Optional[SpanContext]) -> Optional[bytes]:
-    if isinstance(context, SpanContext):
-        encoded_parent_id = _encode_span_id(context.span_id)
-    else:
-        encoded_parent_id = None
-    return encoded_parent_id
+    if context:
+        return _encode_span_id(context.span_id)
+    return None
 
 
 def _encode_attributes(
@@ -239,50 +224,3 @@ def _encode_resource(resource: Resource) -> PB2Resource:
         except Exception as error:  # pylint: disable=broad-except
             _logger.exception(error)
     return pb2_resource
-
-
-def _encode_instrumentation_scope(
-    instrumentation_scope: InstrumentationScope,
-) -> PB2InstrumentationScope:
-    if instrumentation_scope is None:
-        pb2_instrumentation_scope = PB2InstrumentationScope()
-    else:
-        pb2_instrumentation_scope = PB2InstrumentationScope(
-            name=instrumentation_scope.name,
-            version=instrumentation_scope.version,
-        )
-    return pb2_instrumentation_scope
-
-
-def _encode_value(value: Any) -> PB2AnyValue:
-    if isinstance(value, bool):
-        any_value = PB2AnyValue(bool_value=value)
-    elif isinstance(value, str):
-        any_value = PB2AnyValue(string_value=value)
-    elif isinstance(value, int):
-        any_value = PB2AnyValue(int_value=value)
-    elif isinstance(value, float):
-        any_value = PB2AnyValue(double_value=value)
-    elif isinstance(value, abc.Sequence):
-        any_value = PB2AnyValue(
-            array_value=PB2ArrayValue(values=[_encode_value(v) for v in value])
-        )
-    # tracing specs currently does not support Mapping type attributes.
-    # elif isinstance(value, abc.Mapping):
-    #     pass
-    else:
-        raise Exception(f"Invalid type {type(value)} of value {value}")
-    return any_value
-
-
-def _encode_key_value(key: str, value: Any) -> PB2KeyValue:
-    any_value = _encode_value(value)
-    return PB2KeyValue(key=key, value=any_value)
-
-
-def _encode_span_id(span_id: int) -> bytes:
-    return span_id.to_bytes(length=8, byteorder="big", signed=False)
-
-
-def _encode_trace_id(trace_id: int) -> bytes:
-    return trace_id.to_bytes(length=16, byteorder="big", signed=False)

exporter/opentelemetry-exporter-otlp-proto-http/tests/test_protobuf_encoder.py renamed to exporter/opentelemetry-exporter-otlp-proto-common/tests/test_trace_encoder.py

@@ -17,12 +17,14 @@
 import unittest
 from typing import List, Tuple
 
-from opentelemetry.exporter.otlp.proto.http.trace_exporter.encoder import (
-    _SPAN_KIND_MAP,
+from opentelemetry.exporter.otlp.proto.common import (
     _encode_span_id,
-    _encode_status,
     _encode_trace_id,
-    _ProtobufEncoder,
+)
+from opentelemetry.exporter.otlp.proto.common.trace_encoder import (
+    _SPAN_KIND_MAP,
+    _encode_status,
+    encode_spans,
 )
 from opentelemetry.proto.collector.trace.v1.trace_service_pb2 import (
     ExportTraceServiceRequest as PB2ExportTraceServiceRequest,
@@ -55,19 +57,10 @@
 from opentelemetry.trace.status import StatusCode as SDKStatusCode
 
 
-class TestProtobufEncoder(unittest.TestCase):
-    def test_encode(self):
+class TestOTLPTraceEncoder(unittest.TestCase):
+    def test_encode_spans(self):
         otel_spans, expected_encoding = self.get_exhaustive_test_spans()
-        self.assertEqual(
-            _ProtobufEncoder().encode(otel_spans), expected_encoding
-        )
-
-    def test_serialize(self):
-        otel_spans, expected_encoding = self.get_exhaustive_test_spans()
-        self.assertEqual(
-            _ProtobufEncoder().serialize(otel_spans),
-            expected_encoding.SerializeToString(),
-        )
+        self.assertEqual(encode_spans(otel_spans), expected_encoding)
 
     @staticmethod
     def get_exhaustive_otel_span_list() -> List[SDKSpan]: