Sweet Syntax Sugar To Make OTel Delightful
Manual OTel is honestly very laborious. Typically you end up with using the context manager followed by a bunch of span.set_attribute and span.add_event calls. It makes the code explicit and clear but it's just a lot of boilerplate. It looks like this:
def foo(a: int, b: int, c: Dict[str, str]):
with tracer.start_as_current_span("foo") as span:
span.set_attribute("a", a)
span.set_attribute("b", b)
span.set_attribute("c", c)
logger.info("Hello, World!", c=c)
return a + b
Recently I've been working on a project that requires a lot of attributes population. To be more DRY I ended up with creating a traceit decorator that can be used to wrap functions.
Here is the code snippet for the traceit decorator:
import json
from functools import wraps
from opentelemetry import trace
import inspect
from typing import Callable
tracer = trace.get_tracer("opsmate")
def traceit(*args, name: str = None, exclude: list = []):
"""
Decorator to trace function calls.
Usage:
@traceit # all arguments will be traced as attributes
def my_function(a, b, c):
pass
@traceit(exclude=["b"]) # b will not be traced as an attribute
def my_function(a, b, c):
pass
@traceit # Span can be accessed in the function
def my_function(a, b, c, span: Span = None):
span.add_event("my_event", {"key": "value"})
return a + b + c
"""
if len(args) == 1 and callable(args[0]):
return _traceit(args[0], name, exclude)
elif len(args) == 0:
def decorator(func: Callable):
return _traceit(func, name, exclude)
return decorator
else:
raise ValueError("Invalid arguments")
def _traceit(func: Callable, name: str = None, exclude: list = []):
@wraps(func)
def wrapper(*args, **kwargs):
kvs = {}
parameters = inspect.signature(func).parameters
parameter_items = list(parameters.values())
for idx, val in enumerate(args):
if parameter_items[idx].name in exclude:
continue
if parameter_items[idx].annotation in (int, str, bool, float):
kvs[parameter_items[idx].name] = val
elif parameter_items[idx].annotation in (dict, list):
kvs[parameter_items[idx].name] = json.dumps(val)
for k, v in kwargs.items():
if k in exclude:
continue
if isinstance(k, (int, str, bool, float)):
kvs[k] = v
elif isinstance(k, (dict, list)):
kvs[k] = json.dumps(v)
span_name = name or func.__qualname__
with tracer.start_as_current_span(span_name) as span:
for k, v in kvs.items():
span.set_attribute(f"{span_name}.{k}", v)
if parameters.get("span") is not None:
kwargs["span"] = span
return func(*args, **kwargs)
return wrapper
Like most of the decorator code, it's not the easiest code to read but it gets the job done.
Essentially the following code equivalent:
# before
def foo(a: int, b: int, c: Dict[str, str]):
with tracer.start_as_current_span("foo") as span:
span.set_attribute("foo.a", a)
span.set_attribute("foo.b", b)
span.set_attribute("foo.c", c)
span.add_event("result", {"result": a + b})
return a + b
# after
@traceit(name="foo")
def foo(a: int, b: int, c: Dict[str, str], span):
span.add_event("result", {"result": a + b})
return a + b
Here are some of the pytest tests to validate the decorator. It also makes sure that we have a baseline safety net to expand the functionality in the future:
import pytest
from opentelemetry import trace
from opentelemetry.sdk.trace import TracerProvider
from opentelemetry.sdk.trace.export import SimpleSpanProcessor
from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter
from opsmate.libs.core.trace import traceit
@pytest.fixture(scope="module", autouse=True)
def tracer_provider():
exporter = InMemorySpanExporter()
tracer_provider = TracerProvider()
span_processor = SimpleSpanProcessor(exporter)
tracer_provider.add_span_processor(span_processor)
trace.set_tracer_provider(tracer_provider)
yield tracer_provider, exporter
# teardown
tracer_provider.shutdown()
exporter.clear()
span_processor.shutdown()
def test_traceit_with_args(tracer_provider):
tracer_provider, exporter = tracer_provider
@traceit
def func(a: int, b: str, c: dict):
return a + len(b) + len(c)
result = func(1, "world", {"y": 2})
assert result == 7
spans = exporter.get_finished_spans()
span = spans[-1]
assert span.name == "test_traceit_with_args.<locals>.func"
assert span.attributes["test_traceit_with_args.<locals>.func.a"] == 1
assert span.attributes["test_traceit_with_args.<locals>.func.b"] == "world"
assert span.attributes["test_traceit_with_args.<locals>.func.c"] == '{"y": 2}'
def test_traceit_with_name(tracer_provider):
tracer_provider, exporter = tracer_provider
@traceit(name="da_func")
def func(a: int, b: str, c: dict):
return a + len(b) + len(c)
result = func(1, "world", {"y": 2})
assert result == 7
spans = exporter.get_finished_spans()
span = spans[-1]
assert span.name == "da_func"
assert span.attributes["da_func.a"] == 1
assert span.attributes["da_func.b"] == "world"
assert span.attributes["da_func.c"] == '{"y": 2}'
def test_traceit_with_exclude(tracer_provider):
tracer_provider, exporter = tracer_provider
@traceit(name="da_func", exclude=["b"])
def func(a: int, b: str, c: dict):
return a + len(b) + len(c)
result = func(1, "hello", {"x": 1})
assert result == 7
spans = exporter.get_finished_spans()
span = spans[-1]
assert span.name == "da_func"
assert span.attributes["da_func.a"] == 1
assert span.attributes["da_func.c"] == '{"x": 1}'
assert "da_func.b" not in span.attributes
def test_traceit_with_span_arg(tracer_provider):
tracer_provider, exporter = tracer_provider
@traceit(name="da_func")
def func(a: int, span: trace.Span = None):
span.add_event("test_event")
return a
result = func(1)
assert result == 1
spans = exporter.get_finished_spans()
span = spans[-1]
assert span.name == "da_func"
assert span.attributes["da_func.a"] == 1
assert span.events[0].name == "test_event"
In real life usage I ended up with replacing all the previous manual instrumentation with this decorator. It generates a pretty comprehensive trace graph that looks like this:
It is not exactly a life changer but certainly reduces the amount of toils you need to deal with in order to get the observability right.