chore(tracer): migrate to native based SpanLink class

ddtrace/_trace/_span_link.py

@@ -24,12 +24,9 @@
     s1.link_span(s2.context, link_attributes)
 """
 
-import dataclasses
 from enum import Enum
-from typing import Optional
 
-from ddtrace.internal.native._native import SpanLinkData
-from ddtrace.internal.utils.formats import flatten_key_value
+from ddtrace.internal.native._native import SpanLink
 
 
 class SpanLinkKind(Enum):
@@ -40,89 +37,4 @@ class SpanLinkKind(Enum):
     SPAN_POINTER = "span-pointer"  # Should not be used on normal SpanLinks.
 
 
-def _id_not_zero(self, attribute_name, value):
-    if not value > 0:
-        raise ValueError(f"{attribute_name} must be > 0. Value is {value}")
-
-
-@dataclasses.dataclass
-class SpanLink(SpanLinkData):
-    """
-    TraceId [required]: The span's 128-bit Trace ID
-    SpanId [required]: The span's 64-bit Span ID
-
-    Flags [optional]: The span's trace-flags field, as defined in the W3C standard. If only sampling
-    information is provided, the flags value must be 1 if the decision is keep, otherwise 0.
-
-    TraceState [optional]: The span's tracestate field, as defined in the W3C standard.
-
-    Attributes [optional]: Zero or more key-value pairs, where the key must be a non-empty string and the
-    value is either a string, bool, number or an array of primitive type values.
-    """
-
-    trace_id: int
-    span_id: int
-    tracestate: Optional[str] = None
-    flags: Optional[int] = None
-    attributes: dict = dataclasses.field(default_factory=dict)
-    _dropped_attributes: int = 0
-
-    def __post_init__(self):
-        _id_not_zero(self, "trace_id", self.trace_id)
-        _id_not_zero(self, "span_id", self.span_id)
-
-    @property
-    def name(self):
-        return self.attributes["link.name"]
-
-    @name.setter
-    def name(self, value):
-        self.attributes["link.name"] = value
-
-    @property
-    def kind(self) -> Optional[str]:
-        return self.attributes.get("link.kind")
-
-    @kind.setter
-    def kind(self, value: str) -> None:
-        self.attributes["link.kind"] = value
-
-    def to_dict(self):
-        d = {
-            "trace_id": f"{self.trace_id:032x}",
-            "span_id": f"{self.span_id:016x}",
-        }
-        if self.attributes:
-            d["attributes"] = {}
-            for k, v in self.attributes.items():
-                # flatten all values with the type list, tuple and set
-                for k1, v1 in flatten_key_value(k, v).items():
-                    # convert all values to string
-                    if isinstance(v1, str):
-                        d["attributes"][k1] = v1
-                    elif isinstance(v1, bool):
-                        # convert bool to lowercase string to be consistent with json encoding
-                        d["attributes"][k1] = str(v1).lower()
-                    else:
-                        d["attributes"][k1] = str(v1)
-
-        if self._dropped_attributes > 0:
-            d["dropped_attributes_count"] = self._dropped_attributes
-        if self.tracestate:
-            d["tracestate"] = self.tracestate
-        if self.flags is not None:
-            d["flags"] = self.flags
-
-        return d
-
-    def __getstate__(self):
-        return dataclasses.asdict(self)
-
-    def __setstate__(self, state):
-        self.__init__(**state)
-
-    def __repr__(self) -> str:
-        return (
-            f"SpanLink(trace_id={self.trace_id}, span_id={self.span_id}, attributes={self.attributes}, "
-            f"tracestate={self.tracestate}, flags={self.flags}, dropped_attributes={self._dropped_attributes})"
-        )
+__all__ = ["SpanLink", "SpanLinkKind"]

ddtrace/_trace/_span_pointer.py

@@ -3,10 +3,8 @@
 import random
 from typing import Any
 from typing import NamedTuple
-from typing import Optional
 
-from ddtrace._trace._span_link import SpanLink
-from ddtrace._trace._span_link import SpanLinkKind
+from ddtrace._trace._span_link import SpanLinkKind  # noqa: F401 - kept for backward compat
 from ddtrace._trace.telemetry import record_span_pointer_calculation_issue
 from ddtrace.internal.logger import get_logger
 
@@ -40,39 +38,6 @@ class _SpanPointerDescription(NamedTuple):
     extra_attributes: dict[str, Any]
 
 
-class _SpanPointer(SpanLink):
-    def __init__(
-        self,
-        pointer_kind: str,
-        pointer_direction: _SpanPointerDirection,
-        pointer_hash: str,
-        extra_attributes: Optional[dict[str, Any]] = None,
-    ):
-        super().__init__(
-            trace_id=_SPAN_POINTER_SPAN_LINK_TRACE_ID,
-            span_id=_SPAN_POINTER_SPAN_LINK_SPAN_ID,
-            attributes={
-                "ptr.kind": pointer_kind,
-                "ptr.dir": pointer_direction.value,
-                "ptr.hash": pointer_hash,
-                **(extra_attributes or {}),
-            },
-        )
-
-        self.kind = SpanLinkKind.SPAN_POINTER.value
-
-    def __post_init__(self):
-        # Do not want to do the trace_id and span_id checks that SpanLink does.
-        pass
-
-    def __repr__(self):
-        return (
-            f"SpanPointer(trace_id={self.trace_id}, span_id={self.span_id}, kind={self.kind}, "
-            f"direction={self.attributes.get('ptr.dir')}, hash={self.attributes.get('ptr.hash')}, "
-            f"attributes={self.attributes})"
-        )
-
-
 _STANDARD_HASHING_FUNCTION_FAILURE_PREFIX = "HashingFailure"
 
 

ddtrace/_trace/span.py

@@ -14,7 +14,6 @@
 from ddtrace._trace._limits import MAX_SPAN_META_VALUE_LEN
 from ddtrace._trace._span_link import SpanLink
 from ddtrace._trace._span_link import SpanLinkKind
-from ddtrace._trace._span_pointer import _SpanPointer
 from ddtrace._trace._span_pointer import _SpanPointerDirection
 from ddtrace._trace.context import Context
 from ddtrace._trace.types import _AttributeValueType
@@ -136,7 +135,7 @@ def __init__(
             else Context(trace_id=_trace_id, span_id=_span_id, is_remote=False)
         )
 
-        self._links: list[Union[SpanLink, _SpanPointer]] = []
+        self._links: list[SpanLink] = []
         if links:
             for new_link in links:
                 self._set_link_or_append_pointer(new_link)
@@ -669,15 +668,15 @@ def _add_span_pointer(
         # This is a Private API for now.
 
         self._set_link_or_append_pointer(
-            _SpanPointer(
+            SpanLink._SpanPointer(
                 pointer_kind=pointer_kind,
                 pointer_direction=pointer_direction,
                 pointer_hash=pointer_hash,
                 extra_attributes=extra_attributes,
             )
         )
 
-    def _set_link_or_append_pointer(self, link: Union[SpanLink, _SpanPointer]) -> None:
+    def _set_link_or_append_pointer(self, link: SpanLink) -> None:
         if link.kind == SpanLinkKind.SPAN_POINTER.value:
             self._links.append(link)
             return

ddtrace/internal/native/_native.pyi

@@ -581,17 +581,44 @@ class SpanEvent:
     def __iter__(self) -> Iterator[tuple[str, Any]]: ...
     def __reduce__(self) -> tuple: ...
 
-class SpanLinkData:
+class SpanLink:
+    SPAN_POINTER_KIND: str
+    trace_id: int
+    span_id: int
+    tracestate: Optional[str]
+    flags: Optional[int]
+    attributes: dict[str, Any]
+    _dropped_attributes: int
+
     def __init__(
         self,
         trace_id: int,
         span_id: int,
         tracestate: Optional[str] = None,
         flags: Optional[int] = None,
-        attributes: Optional[dict[str, str]] = None,
+        attributes: Optional[Mapping[str, Any]] = None,
         _dropped_attributes: int = 0,
-    ): ...
-
+        _skip_validation: bool = False,
+    ) -> None: ...
+    @property
+    def name(self) -> Any: ...
+    @property
+    def kind(self) -> Optional[Any]: ...
+    def to_dict(self) -> dict[str, Any]: ...
+    def __eq__(self, other: object) -> bool: ...
+    def __repr__(self) -> str: ...
+    def __reduce__(self) -> tuple: ...
+    @classmethod
+    def _SpanPointer(
+        cls,
+        pointer_kind: str,
+        pointer_direction: Any,
+        pointer_hash: str,
+        extra_attributes: Optional[dict[str, Any]] = None,
+    ) -> "SpanLink": ...
+
+def flatten_key_value(root_key: str, value: Any) -> dict[str, Any]: ...
+def is_sequence(obj: Any) -> bool: ...
 def seed() -> None: ...
 def rand64bits() -> int: ...
 def generate_128bit_trace_id() -> int: ...

ddtrace/internal/utils/formats.py

@@ -6,6 +6,8 @@
 from typing import Union  # noqa:F401
 
 from ddtrace.internal.constants import MAX_UINT_64BITS  # noqa:F401
+from ddtrace.internal.native._native import flatten_key_value  # noqa: F401
+from ddtrace.internal.native._native import is_sequence  # noqa: F401
 
 from ..compat import ensure_text
 
@@ -126,29 +128,6 @@ def stringify_cache_args(args: list[Any], value_max_len: int = VALUE_MAX_LEN, cm
     return " ".join(out)
 
 
-def is_sequence(obj: Any) -> bool:
-    try:
-        return isinstance(obj, (list, tuple, set, frozenset))
-    except TypeError:
-        # Checking the type of Generic Subclasses raises a TypeError
-        return False
-
-
-def flatten_key_value(root_key: str, value: Any) -> dict[str, Any]:
-    """Flattens attributes"""
-    if not is_sequence(value):
-        return {root_key: value}
-
-    flattened = dict()
-    for i, item in enumerate(value):
-        key = f"{root_key}.{i}"
-        if is_sequence(item):
-            flattened.update(flatten_key_value(key, item))
-        else:
-            flattened[key] = item
-    return flattened
-
-
 def format_trace_id(trace_id: int) -> str:
     """Translate a trace ID to a string format supported by the backend."""
     return "{:032x}".format(trace_id) if trace_id > MAX_UINT_64BITS else str(trace_id)

src/native/lib.rs

@@ -14,6 +14,7 @@ mod py_string;
 mod rand;
 mod span;
 mod tracer_flare;
+mod utils;
 
 use pyo3::prelude::*;
 
@@ -49,6 +50,8 @@ fn _native(m: &Bound<'_, PyModule>) -> PyResult<()> {
     data_pipeline::register_data_pipeline(m)?;
     span::register_native_span(m)?;
     rand::register_rand(m)?;
+    m.add_function(wrap_pyfunction!(utils::flatten_key_value, m)?)?;
+    m.add_function(wrap_pyfunction!(utils::is_sequence, m)?)?;
     m.add_wrapped(pyo3::wrap_pymodule!(config::config_module))?;
 
     // Add FFE submodule

src/native/span/mod.rs

@@ -7,10 +7,10 @@ pub mod utils;
 
 pub use span_data::SpanData;
 pub use span_event::SpanEvent;
-pub use span_link::SpanLinkData;
+pub use span_link::SpanLink;
 
 pub fn register_native_span(m: &pyo3::Bound<'_, pyo3::types::PyModule>) -> pyo3::PyResult<()> {
-    m.add_class::<SpanLinkData>()?;
+    m.add_class::<SpanLink>()?;
     m.add_class::<SpanEvent>()?;
     m.add_class::<SpanData>()?;
     Ok(())