Repository: elastic/apm-agent-python

Issue:
[META 643] Implement Span Links API
See meta issue for the description and details:
- Meta issue: https://github.com/elastic/apm/issues/643



Discussion:

Resolving patch:
diff --git a/elasticapm/base.py b/elasticapm/base.py
--- a/elasticapm/base.py
+++ b/elasticapm/base.py
@@ -44,7 +44,7 @@
 import warnings
 from copy import deepcopy
 from datetime import timedelta
-from typing import Optional, Tuple
+from typing import Optional, Sequence, Tuple
 
 import elasticapm
 from elasticapm.conf import Config, VersionedConfig, constants
@@ -52,6 +52,7 @@
 from elasticapm.metrics.base_metrics import MetricsRegistry
 from elasticapm.traces import Tracer, execution_context
 from elasticapm.utils import cgroup, cloud, compat, is_master_process, stacks, varmap
+from elasticapm.utils.disttracing import TraceParent
 from elasticapm.utils.encoding import enforce_label_format, keyword_field, shorten, transform
 from elasticapm.utils.logging import get_logger
 from elasticapm.utils.module_import import import_string
@@ -288,7 +289,14 @@ def queue(self, event_type, data, flush=False):
             flush = False
         self._transport.queue(event_type, data, flush)
 
-    def begin_transaction(self, transaction_type, trace_parent=None, start=None, auto_activate=True):
+    def begin_transaction(
+        self,
+        transaction_type: str,
+        trace_parent: Optional[TraceParent] = None,
+        start: Optional[float] = None,
+        auto_activate: bool = True,
+        links: Optional[Sequence[TraceParent]] = None,
+    ):
         """
         Register the start of a transaction on the client
 
@@ -300,7 +308,7 @@ def begin_transaction(self, transaction_type, trace_parent=None, start=None, aut
         """
         if self.config.is_recording:
             return self.tracer.begin_transaction(
-                transaction_type, trace_parent=trace_parent, start=start, auto_activate=auto_activate
+                transaction_type, trace_parent=trace_parent, start=start, auto_activate=auto_activate, links=links
             )
 
     def end_transaction(self, name=None, result="", duration=None):
diff --git a/elasticapm/contrib/opentelemetry/trace.py b/elasticapm/contrib/opentelemetry/trace.py
--- a/elasticapm/contrib/opentelemetry/trace.py
+++ b/elasticapm/contrib/opentelemetry/trace.py
@@ -36,7 +36,7 @@
 
 from opentelemetry import trace as trace_api
 from opentelemetry.sdk import trace as oteltrace
-from opentelemetry.trace import Context, SpanKind
+from opentelemetry.trace import Context, Link, SpanKind
 from opentelemetry.trace.propagation import _SPAN_KEY
 from opentelemetry.trace.status import Status, StatusCode
 from opentelemetry.util import types
@@ -76,7 +76,7 @@ def start_span(
         context: Optional[Context] = None,
         kind: SpanKind = SpanKind.INTERNAL,
         attributes: types.Attributes = None,
-        links: Optional[Sequence[Any]] = None,
+        links: Optional[Sequence[Link]] = None,
         start_time: Optional[int] = None,
         record_exception: bool = True,
         set_status_on_exception: bool = True,
@@ -116,8 +116,6 @@ def start_span(
         Returns:
             The newly-created span.
         """
-        if links:
-            logger.warning("The opentelemetry bridge does not support links at this time.")
         if not record_exception:
             logger.warning("record_exception was set to False, but exceptions will still be recorded for this span.")
 
@@ -130,6 +128,7 @@ def start_span(
         current_transaction = execution_context.get_transaction()
         client = self.client
 
+        elastic_links = tuple(get_traceparent(link.context) for link in links) if links else None
         if traceparent and current_transaction:
             logger.warning(
                 "Remote context included when a transaction was already active. "
@@ -137,7 +136,7 @@ def start_span(
             )
         elif traceparent:
             elastic_span = client.begin_transaction(
-                "otel", traceparent=traceparent, start=start_time, auto_activate=False
+                "otel", trace_parent=traceparent, start=start_time, auto_activate=False, links=elastic_links
             )
             span = Span(
                 name=name,
@@ -147,7 +146,7 @@ def start_span(
             )
             span.set_attributes(attributes)
         elif not current_transaction:
-            elastic_span = client.begin_transaction("otel", start=start_time, auto_activate=False)
+            elastic_span = client.begin_transaction("otel", start=start_time, auto_activate=False, links=elastic_links)
             span = Span(
                 name=name,
                 elastic_span=elastic_span,
@@ -156,7 +155,9 @@ def start_span(
             )
             span.set_attributes(attributes)
         else:
-            elastic_span = current_transaction.begin_span(name, "otel", start=start_time, auto_activate=False)
+            elastic_span = current_transaction.begin_span(
+                name, "otel", start=start_time, auto_activate=False, links=elastic_links
+            )
             span = Span(
                 name=name,
                 elastic_span=elastic_span,
@@ -176,7 +177,7 @@ def start_as_current_span(
         context: Optional[Context] = None,
         kind: SpanKind = SpanKind.INTERNAL,
         attributes: types.Attributes = None,
-        links: Optional[Sequence[Any]] = None,
+        links: Optional[Sequence[Link]] = None,
         start_time: Optional[int] = None,
         record_exception: bool = True,
         set_status_on_exception: bool = True,
diff --git a/elasticapm/traces.py b/elasticapm/traces.py
--- a/elasticapm/traces.py
+++ b/elasticapm/traces.py
@@ -38,7 +38,7 @@
 from collections import defaultdict
 from datetime import timedelta
 from types import TracebackType
-from typing import Any, Callable, Dict, List, Optional, Tuple, Type, Union
+from typing import Any, Callable, Dict, List, Optional, Sequence, Tuple, Type, Union
 
 import elasticapm
 from elasticapm.conf import constants
@@ -91,7 +91,7 @@ def duration(self) -> timedelta:
 
 
 class BaseSpan(object):
-    def __init__(self, labels=None, start=None):
+    def __init__(self, labels=None, start=None, links: Optional[Sequence[TraceParent]] = None):
         self._child_durations = ChildDuration(self)
         self.labels = {}
         self.outcome: Optional[str] = None
@@ -100,7 +100,10 @@ def __init__(self, labels=None, start=None):
         self.start_time: float = time_to_perf_counter(start) if start is not None else _time_func()
         self.ended_time: Optional[float] = None
         self.duration: Optional[timedelta] = None
-        self.links: List[Dict[str, str]] = []
+        self.links: Optional[List[Dict[str, str]]] = None
+        if links:
+            for trace_parent in links:
+                self.add_link(trace_parent)
         if labels:
             self.label(**labels)
 
@@ -149,6 +152,8 @@ def add_link(self, trace_parent: TraceParent) -> None:
         """
         Causally link this span/transaction to another span/transaction
         """
+        if self.links is None:
+            self.links = []
         self.links.append({"trace_id": trace_parent.trace_id, "span_id": trace_parent.span_id})
 
     def set_success(self):
@@ -175,6 +180,7 @@ def __init__(
         is_sampled: bool = True,
         start: Optional[float] = None,
         sample_rate: Optional[float] = None,
+        links: Optional[Sequence[TraceParent]] = None,
     ):
         """
         tracer
@@ -193,6 +199,8 @@ def __init__(
             Sample rate which was used to decide whether to sample this transaction.
             This is reported to the APM server so that unsampled transactions can
             be extrapolated.
+        links:
+            A list of traceparents to link this transaction causally
         """
         self.id = self.get_dist_tracing_id()
         if not trace_parent:
@@ -233,7 +241,10 @@ def __init__(
             )
         except (LookupError, AttributeError):
             self._breakdown = None
-        super(Transaction, self).__init__(start=start)
+        super().__init__(start=start)
+        if links:
+            for trace_parent in links:
+                self.add_link(trace_parent)
 
     def end(self, skip_frames: int = 0, duration: Optional[timedelta] = None):
         super().end(skip_frames, duration)
@@ -271,6 +282,7 @@ def _begin_span(
         sync=None,
         start=None,
         auto_activate=True,
+        links: Optional[Sequence[TraceParent]] = None,
     ):
         parent_span = execution_context.get_span()
         tracer = self.tracer
@@ -293,6 +305,7 @@ def _begin_span(
                 span_action=span_action,
                 sync=sync,
                 start=start,
+                links=links,
             )
             span.frames = tracer.frames_collector_func()
             self._span_counter += 1
@@ -312,6 +325,7 @@ def begin_span(
         sync=None,
         start=None,
         auto_activate=True,
+        links: Optional[Sequence[TraceParent]] = None,
     ):
         """
         Begin a new span
@@ -325,6 +339,7 @@ def begin_span(
         :param sync: indicate if the span is synchronous or not. In most cases, `None` should be used
         :param start: timestamp, mostly useful for testing
         :param auto_activate: whether to set this span in execution_context
+        :param links: an optional list of traceparents to link this span with
         :return: the Span object
         """
         return self._begin_span(
@@ -339,6 +354,7 @@ def begin_span(
             sync=sync,
             start=start,
             auto_activate=auto_activate,
+            links=links,
         )
 
     def end_span(self, skip_frames: int = 0, duration: Optional[float] = None, outcome: str = "unknown"):
@@ -503,6 +519,7 @@ def __init__(
         span_action: Optional[str] = None,
         sync: Optional[bool] = None,
         start: Optional[int] = None,
+        links: Optional[Sequence[TraceParent]] = None,
     ):
         """
         Create a new Span
@@ -538,7 +555,7 @@ def __init__(
         self.dist_tracing_propagated = False
         self.composite: Dict[str, Any] = {}
         self._cancelled: bool = False
-        super(Span, self).__init__(labels=labels, start=start)
+        super().__init__(labels=labels, start=start, links=links)
         self.timestamp = transaction.timestamp + (self.start_time - transaction.start_time)
         if self.transaction._breakdown:
             p = self.parent if self.parent else self.transaction
@@ -870,7 +887,14 @@ def span_stack_trace_min_duration(self) -> timedelta:
             else:
                 return self.config.span_frames_min_duration
 
-    def begin_transaction(self, transaction_type, trace_parent=None, start=None, auto_activate=True):
+    def begin_transaction(
+        self,
+        transaction_type: str,
+        trace_parent: Optional[TraceParent] = None,
+        start: Optional[float] = None,
+        auto_activate: bool = True,
+        links: Optional[Sequence[TraceParent]] = None,
+    ):
         """
         Start a new transactions and bind it in a thread-local variable
 
@@ -878,6 +902,7 @@ def begin_transaction(self, transaction_type, trace_parent=None, start=None, aut
         :param trace_parent: an optional TraceParent object
         :param start: override the start timestamp, mostly useful for testing
         :param auto_activate: whether to set this transaction in execution_context
+        :param list of traceparents to causally link this transaction to
 
         :returns the Transaction object
         """
@@ -900,6 +925,7 @@ def begin_transaction(self, transaction_type, trace_parent=None, start=None, aut
             is_sampled=is_sampled,
             start=start,
             sample_rate=sample_rate,
+            links=links,
         )
         if trace_parent is None:
             transaction.trace_parent.add_tracestate(constants.TRACESTATE.SAMPLE_RATE, sample_rate)
@@ -949,6 +975,7 @@ class capture_span(object):
         "duration",
         "start",
         "sync",
+        "links",
     )
 
     def __init__(
@@ -964,6 +991,7 @@ def __init__(
         start: Optional[int] = None,
         duration: Optional[Union[float, timedelta]] = None,
         sync: Optional[bool] = None,
+        links: Optional[Sequence[TraceParent]] = None,
     ):
         self.name = name
         if span_subtype is None and "." in span_type:
@@ -985,6 +1013,7 @@ def __init__(
             duration = timedelta(seconds=duration)
         self.duration = duration
         self.sync = sync
+        self.links = links
 
     def __call__(self, func: Callable) -> Callable:
         self.name = self.name or get_name_from_func(func)
@@ -1017,6 +1046,7 @@ def handle_enter(self, sync: bool) -> Optional[SpanType]:
                 span_action=self.action,
                 start=self.start,
                 sync=sync,
+                links=self.links,
             )
         return None
 


Test patch:
diff --git a/tests/client/transaction_tests.py b/tests/client/transaction_tests.py
--- a/tests/client/transaction_tests.py
+++ b/tests/client/transaction_tests.py
@@ -446,3 +446,23 @@ def test_backdating_transaction(elasticapm_client):
     elasticapm_client.end_transaction()
     transaction = elasticapm_client.events[constants.TRANSACTION][0]
     assert 1000 < transaction["duration"] < 2000
+
+
+def test_transaction_span_links(elasticapm_client):
+    tp1 = TraceParent.from_string("00-aabbccddeeff00112233445566778899-0011223344556677-01")
+    tp2 = TraceParent.from_string("00-00112233445566778899aabbccddeeff-aabbccddeeff0011-01")
+    elasticapm_client.begin_transaction("a", links=[tp1, tp2])
+    with elasticapm.capture_span("a", links=(tp1, tp2)):
+        pass
+    elasticapm_client.end_transaction("foo")
+    transaction = elasticapm_client.events[constants.TRANSACTION][0]
+    span = elasticapm_client.events[constants.SPAN][0]
+    assert transaction["links"][0]["trace_id"] == "aabbccddeeff00112233445566778899"
+    assert transaction["links"][0]["span_id"] == "0011223344556677"
+    assert transaction["links"][1]["trace_id"] == "00112233445566778899aabbccddeeff"
+    assert transaction["links"][1]["span_id"] == "aabbccddeeff0011"
+
+    assert span["links"][0]["trace_id"] == "aabbccddeeff00112233445566778899"
+    assert span["links"][0]["span_id"] == "0011223344556677"
+    assert span["links"][1]["trace_id"] == "00112233445566778899aabbccddeeff"
+    assert span["links"][1]["span_id"] == "aabbccddeeff0011"
diff --git a/tests/contrib/opentelemetry/tests.py b/tests/contrib/opentelemetry/tests.py
--- a/tests/contrib/opentelemetry/tests.py
+++ b/tests/contrib/opentelemetry/tests.py
@@ -32,7 +32,7 @@
 
 pytest.importorskip("opentelemetry.sdk")  # isort:skip
 
-from opentelemetry.trace import SpanContext, SpanKind, TraceFlags
+from opentelemetry.trace import Link, SpanContext, SpanKind, TraceFlags
 from opentelemetry.trace.propagation import _SPAN_KEY
 
 import elasticapm.contrib.opentelemetry.context as context
@@ -138,4 +138,21 @@ def test_ot_spancontext(tracer: Tracer):
     assert span_context.trace_flags.sampled == TraceFlags.SAMPLED
 
 
+def test_span_links(tracer: Tracer):
+    span_context = SpanContext(
+        trace_id=int("aabbccddeeff00112233445566778899", 16), span_id=int("0011223344556677", 16), is_remote=False
+    )
+    link = Link(span_context)
+    with tracer.start_as_current_span("testtransaction", links=[link]):
+        with tracer.start_as_current_span("testspan", links=[link]):
+            pass
+    client = tracer.client
+    transaction = client.events[constants.TRANSACTION][0]
+    span = client.events[constants.SPAN][0]
+    assert transaction["links"][0]["trace_id"] == "aabbccddeeff00112233445566778899"
+    assert transaction["links"][0]["span_id"] == "0011223344556677"
+    assert span["links"][0]["trace_id"] == "aabbccddeeff00112233445566778899"
+    assert span["links"][0]["span_id"] == "0011223344556677"
+
+
 # TODO Add some span subtype testing?