Fixing RichConsoleExporter to allow for multiple traces at once (open-telemetry#1336)

Author: rbizos

CHANGELOG.md

@@ -12,6 +12,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - `opentelemetry-instrumentation-django` Fixed bug where auto-instrumentation fails when django is installed and settings are not configured.
   ([#1369](https://github.com/open-telemetry/opentelemetry-python-contrib/pull/1369))
 - `opentelemetry-instrumentation-system-metrics` add supports to collect system thread count. ([#1339](https://github.com/open-telemetry/opentelemetry-python-contrib/pull/1339))
+- `opentelemetry-exporter-richconsole` Fixing RichConsoleExpoter to allow multiple traces, fixing duplicate spans and include resources ([#1336](https://github.com/open-telemetry/opentelemetry-python-contrib/pull/1336))
 
 ## [1.13.0-0.34b0](https://github.com/open-telemetry/opentelemetry-python/releases/tag/v1.13.0-0.34b0) - 2022-09-26
 

exporter/opentelemetry-exporter-richconsole/src/opentelemetry/exporter/richconsole/__init__.py

@@ -54,7 +54,7 @@
 
 import datetime
 import typing
-from typing import Optional
+from typing import Dict, Optional
 
 from rich.console import Console
 from rich.syntax import Syntax
@@ -76,6 +76,11 @@ def _child_to_tree(child: Tree, span: ReadableSpan):
     child.add(
         Text.from_markup(f"[bold cyan]Kind :[/bold cyan] {span.kind.name}")
     )
+    _add_status(child, span)
+    _child_add_optional_attributes(child, span)
+
+
+def _add_status(child: Tree, span: ReadableSpan):
     if not span.status.is_unset:
         if not span.status.is_ok:
             child.add(
@@ -96,6 +101,8 @@ def _child_to_tree(child: Tree, span: ReadableSpan):
             )
         )
 
+
+def _child_add_optional_attributes(child: Tree, span: ReadableSpan):
     if span.events:
         events = child.add(
             label=Text.from_markup("[bold cyan]Events :[/bold cyan] ")
@@ -122,6 +129,16 @@ def _child_to_tree(child: Tree, span: ReadableSpan):
                         f"[bold cyan]{attribute} :[/bold cyan] {span.attributes[attribute]}"
                     )
                 )
+    if span.resource:
+        resources = child.add(
+            label=Text.from_markup("[bold cyan]Resources :[/bold cyan] ")
+        )
+        for resource in span.resource.attributes:
+            resources.add(
+                Text.from_markup(
+                    f"[bold cyan]{resource} :[/bold cyan] {span.resource.attributes[resource]}"
+                )
+            )
 
 
 class RichConsoleSpanExporter(SpanExporter):
@@ -141,35 +158,39 @@ def __init__(
     def export(self, spans: typing.Sequence[ReadableSpan]) -> SpanExportResult:
         if not spans:
             return SpanExportResult.SUCCESS
-        tree = Tree(
-            label=f"Trace {opentelemetry.trace.format_trace_id(spans[0].context.trace_id)}"
-        )
-        parents = {}
-        for span in spans:
-            child = tree.add(
-                label=Text.from_markup(
-                    f"[blue][{_ns_to_time(span.start_time)}][/blue] [bold]{span.name}[/bold], span {opentelemetry.trace.format_span_id(span.context.span_id)}"
-                )
-            )
-            parents[span.context.span_id] = child
-            _child_to_tree(child, span)
-
-        for span in spans:
-            if span.parent and span.parent.span_id in parents:
-                child = parents[span.parent.span_id].add(
-                    label=Text.from_markup(
-                        f"[blue][{_ns_to_time(span.start_time)}][/blue] [bold]{span.name}[/bold], span {opentelemetry.trace.format_span_id(span.context.span_id)}"
-                    )
-                )
-            else:
-                child = tree.add(
-                    label=Text.from_markup(
-                        f"[blue][{_ns_to_time(span.start_time)}][/blue] [bold]{span.name}[/bold], span {opentelemetry.trace.format_span_id(span.context.span_id)}"
-                    )
-                )
 
-            parents[span.context.span_id] = child
-            _child_to_tree(child, span)
+        for tree in self.spans_to_tree(spans).values():
+            self.console.print(tree)
 
-        self.console.print(tree)
         return SpanExportResult.SUCCESS
+
+    @staticmethod
+    def spans_to_tree(spans: typing.Sequence[ReadableSpan]) -> Dict[str, Tree]:
+        trees = {}
+        parents = {}
+        spans = list(spans)
+        while spans:
+            for span in spans:
+                if not span.parent:
+                    trace_id = opentelemetry.trace.format_trace_id(
+                        span.context.trace_id
+                    )
+                    trees[trace_id] = Tree(label=f"Trace {trace_id}")
+                    child = trees[trace_id].add(
+                        label=Text.from_markup(
+                            f"[blue][{_ns_to_time(span.start_time)}][/blue] [bold]{span.name}[/bold], span {opentelemetry.trace.format_span_id(span.context.span_id)}"
+                        )
+                    )
+                    parents[span.context.span_id] = child
+                    _child_to_tree(child, span)
+                    spans.remove(span)
+                elif span.parent and span.parent.span_id in parents:
+                    child = parents[span.parent.span_id].add(
+                        label=Text.from_markup(
+                            f"[blue][{_ns_to_time(span.start_time)}][/blue] [bold]{span.name}[/bold], span {opentelemetry.trace.format_span_id(span.context.span_id)}"
+                        )
+                    )
+                    parents[span.context.span_id] = child
+                    _child_to_tree(child, span)
+                    spans.remove(span)
+        return trees

exporter/opentelemetry-exporter-richconsole/tests/test_rich_exporter.py

@@ -13,7 +13,9 @@
 # limitations under the License.
 
 import pytest
+from rich.tree import Tree
 
+import opentelemetry.trace
 from opentelemetry.exporter.richconsole import RichConsoleSpanExporter
 from opentelemetry.sdk import trace
 from opentelemetry.sdk.trace.export import BatchSpanProcessor
@@ -40,8 +42,57 @@ def fixture_tracer_provider(span_processor):
 def test_span_exporter(tracer_provider, span_processor, capsys):
     tracer = tracer_provider.get_tracer(__name__)
     span = tracer.start_span("test_span")
+
     span.set_attribute("key", "V4LuE")
     span.end()
     span_processor.force_flush()
     captured = capsys.readouterr()
+
     assert "V4LuE" in captured.out
+
+
+def walk_tree(root: Tree) -> int:
+    # counts the amount of spans in a tree that contains a span
+    return sum(walk_tree(child) for child in root.children) + int(
+        "span" in root.label
+    )
+
+
+def test_multiple_traces(tracer_provider):
+    exporter = RichConsoleSpanExporter()
+    tracer = tracer_provider.get_tracer(__name__)
+    with tracer.start_as_current_span("parent_1") as parent_1:
+        with tracer.start_as_current_span("child_1") as child_1:
+            pass
+
+    with tracer.start_as_current_span("parent_2") as parent_2:
+        pass
+
+    trees = exporter.spans_to_tree((parent_2, parent_1, child_1))
+    # asserts that we have all traces
+    assert len(trees) == 2
+    traceid_1 = opentelemetry.trace.format_trace_id(parent_1.context.trace_id)
+
+    assert traceid_1 in trees
+
+    assert (
+        opentelemetry.trace.format_trace_id(parent_2.context.trace_id) in trees
+    )
+
+    # asserts that we have exactly the number of spans we exported
+    assert sum(walk_tree(tree) for tree in trees.values()) == 3
+
+    # assert that the relationship is correct
+    assert parent_1.name in trees[traceid_1].children[0].label
+    assert any(
+        child_1.name in child.label
+        for child in trees[traceid_1].children[0].children
+    )
+    assert not any(
+        parent_1.name in child.label
+        for child in trees[traceid_1].children[0].children
+    )
+    assert not any(
+        parent_2.name in child.label
+        for child in trees[traceid_1].children[0].children
+    )