Add span decorator

Author: draincoder

src/asgi_monitor/tracing/__init__.py

@@ -1,3 +1,4 @@
 from .config import BaseTracingConfig
+from .decorator import span
 
-__all__ = ("BaseTracingConfig",)
+__all__ = ("span", "BaseTracingConfig")

src/asgi_monitor/tracing/decorator.py

@@ -0,0 +1,76 @@
+import asyncio
+from functools import wraps
+from typing import Callable, Mapping, ParamSpec, Sequence, TypeVar, cast, overload
+
+from opentelemetry import trace
+from opentelemetry.trace import Tracer
+
+__all__ = ("span",)
+
+
+F_Spec = ParamSpec("F_Spec")
+F_Return = TypeVar("F_Return")
+Func = Callable[F_Spec, F_Return]
+Attributes = Mapping[str, str | bool | int | float | Sequence[str] | Sequence[bool] | Sequence[int] | Sequence[float]]
+
+
+def _span_wrapper(
+    span_name: str | None,
+    attributes: Attributes | None,
+    tracer: Tracer | None,
+) -> Callable[[Func], Func]:
+    def span_decorator(func: Func) -> Func:
+        name: str = span_name or func.__name__
+        tracer_ = tracer or trace.get_tracer(__name__)
+
+        @wraps(func)
+        def sync_span_wrapper(*args: F_Spec.args, **kwargs: F_Spec.kwargs) -> F_Return:  # type: ignore[type-var]
+            with tracer_.start_as_current_span(name=name, attributes=attributes):
+                result = func(*args, **kwargs)
+            return cast(F_Return, result)
+
+        @wraps(func)
+        async def async_span_wrapper(*args: F_Spec.args, **kwargs: F_Spec.kwargs) -> F_Return:
+            with tracer_.start_as_current_span(name=name, attributes=attributes):
+                result = await func(*args, **kwargs)
+            return cast(F_Return, result)
+
+        if asyncio.iscoroutinefunction(func):
+            return async_span_wrapper
+        return sync_span_wrapper
+
+    return span_decorator
+
+
+@overload
+def span(
+    call: Func,
+    *,
+    span_name: None = None,
+    attributes: None = None,
+    tracer: None = None,
+) -> Func: ...
+
+
+@overload
+def span(
+    call: None = None,
+    *,
+    span_name: str | None = None,
+    attributes: Attributes | None = None,
+    tracer: Tracer | None = None,
+) -> Callable[[Func], Func]: ...
+
+
+def span(
+    call: Func | None = None,
+    *,
+    span_name: str | None = None,
+    attributes: Attributes | None = None,
+    tracer: Tracer | None = None,
+) -> Callable[[Func], Func] | Func:
+    wrap_decorator = _span_wrapper(span_name, attributes, tracer)
+    if call is None:
+        return wrap_decorator
+    else:
+        return wrap_decorator(call)

tests/integration/factory.py

@@ -54,12 +54,12 @@ def build_fastapi_tracing_config() -> tuple[FastAPITraceConfig, InMemorySpanExpo
             "service.name": "fastapi",
         },
     )
-    tracer = TracerProvider(resource=resource)
-    trace.set_tracer_provider(tracer)
+    tracer_provider = TracerProvider(resource=resource)
+    trace.set_tracer_provider(tracer_provider)
     exporter = InMemorySpanExporter()
-    tracer.add_span_processor(SimpleSpanProcessor(exporter))
+    tracer_provider.add_span_processor(SimpleSpanProcessor(exporter))
 
-    return FastAPITraceConfig(tracer_provider=tracer), exporter
+    return FastAPITraceConfig(tracer_provider=tracer_provider), exporter
 
 
 def build_starlette_tracing_config() -> tuple[StarletteTraceConfig, InMemorySpanExporter]:
@@ -68,12 +68,12 @@ def build_starlette_tracing_config() -> tuple[StarletteTraceConfig, InMemorySpan
             "service.name": "starlette",
         },
     )
-    tracer = TracerProvider(resource=resource)
-    trace.set_tracer_provider(tracer)
+    tracer_provider = TracerProvider(resource=resource)
+    trace.set_tracer_provider(tracer_provider)
     exporter = InMemorySpanExporter()
-    tracer.add_span_processor(SimpleSpanProcessor(exporter))
+    tracer_provider.add_span_processor(SimpleSpanProcessor(exporter))
 
-    return StarletteTraceConfig(tracer_provider=tracer), exporter
+    return StarletteTraceConfig(tracer_provider=tracer_provider), exporter
 
 
 def build_litestar_tracing_config() -> tuple[LitestarTraceConfig, InMemorySpanExporter]:
@@ -82,9 +82,9 @@ def build_litestar_tracing_config() -> tuple[LitestarTraceConfig, InMemorySpanEx
             "service.name": "litestar",
         },
     )
-    tracer = TracerProvider(resource=resource)
-    trace.set_tracer_provider(tracer)
+    tracer_provider = TracerProvider(resource=resource)
+    trace.set_tracer_provider(tracer_provider)
     exporter = InMemorySpanExporter()
-    tracer.add_span_processor(SimpleSpanProcessor(exporter))
+    tracer_provider.add_span_processor(SimpleSpanProcessor(exporter))
 
-    return LitestarTraceConfig(tracer_provider=tracer), exporter
+    return LitestarTraceConfig(tracer_provider=tracer_provider), exporter

tests/unit/tracing/conftest.py

@@ -0,0 +1,23 @@
+import pytest
+from opentelemetry.sdk.resources import Resource
+from opentelemetry.sdk.trace import TracerProvider
+from opentelemetry.sdk.trace.export import SimpleSpanProcessor
+from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter
+
+
+@pytest.fixture
+def tracer_provider() -> TracerProvider:
+    resource = Resource.create(
+        attributes={
+            "service.name": "action",
+        },
+    )
+    return TracerProvider(resource=resource)
+
+
+@pytest.fixture
+def exporter(tracer_provider: TracerProvider) -> InMemorySpanExporter:
+    exporter = InMemorySpanExporter()
+    processor = SimpleSpanProcessor(exporter)
+    tracer_provider.add_span_processor(processor)
+    return exporter

tests/unit/tracing/test_span_decorator.py

@@ -0,0 +1,122 @@
+import asyncio
+import time
+
+from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter
+from opentelemetry.trace import TracerProvider
+
+from asgi_monitor.tracing import span as span_decorator
+
+
+def test_simple_span_decorator() -> None:
+    # Arrange
+    expected = 1
+
+    @span_decorator
+    def action() -> int:
+        return expected
+
+    # Act
+    result = action()
+
+    # Assert
+    assert result == expected
+
+
+def test_empty_span_decorator(exporter: InMemorySpanExporter, tracer_provider: TracerProvider) -> None:
+    # Arrange
+    @span_decorator(tracer=tracer_provider.get_tracer(__name__))
+    def action() -> None:
+        time.sleep(0.01)
+
+    # Act
+    action()
+
+    # Assert
+    span = exporter.get_finished_spans()[0]
+    assert span.name == action.__name__
+    assert span.attributes == {}
+
+
+def test_span_decorator(exporter: InMemorySpanExporter, tracer_provider: TracerProvider) -> None:
+    # Arrange
+    @span_decorator(span_name="test", attributes={"test": "test"}, tracer=tracer_provider.get_tracer(__name__))
+    def action() -> None:
+        time.sleep(0.01)
+
+    # Act
+    action()
+
+    # Assert
+    span = exporter.get_finished_spans()[0]
+    assert span.name == "test"
+    assert span.attributes == {"test": "test"}
+
+
+async def test_async_simple_span_decorator(exporter: InMemorySpanExporter, tracer_provider: TracerProvider) -> None:
+    # Arrange
+    expected = 1
+
+    @span_decorator
+    async def action() -> int:
+        return expected
+
+    # Act
+    result = await action()
+
+    # Assert
+    assert result == expected
+
+
+async def test_async_empty_span_decorator(exporter: InMemorySpanExporter, tracer_provider: TracerProvider) -> None:
+    # Arrange
+    @span_decorator(tracer=tracer_provider.get_tracer(__name__))
+    async def action() -> None:
+        await asyncio.sleep(0.01)
+
+    # Act
+    await action()
+
+    # Assert
+    span = exporter.get_finished_spans()[0]
+    assert span.name == action.__name__
+    assert span.attributes == {}
+
+
+async def test_async_span_decorator(exporter: InMemorySpanExporter, tracer_provider: TracerProvider) -> None:
+    # Arrange
+    @span_decorator(span_name="test", attributes={"test": "test"}, tracer=tracer_provider.get_tracer(__name__))
+    async def action() -> None:
+        await asyncio.sleep(0.01)
+
+    # Act
+    await action()
+
+    # Assert
+    span = exporter.get_finished_spans()[0]
+    assert span.name == "test"
+    assert span.attributes == {"test": "test"}
+
+
+async def test_nested_spans(exporter: InMemorySpanExporter, tracer_provider: TracerProvider) -> None:
+    # Arrange
+    tracer = tracer_provider.get_tracer(__name__)
+
+    @span_decorator(tracer=tracer)
+    async def count() -> int:
+        return 1
+
+    @span_decorator(tracer=tracer)
+    async def action() -> None:
+        num = await count()
+        assert num == 1
+
+    # Act
+    await action()
+
+    # Assert
+    count_sleep, span_action = exporter.get_finished_spans()
+    assert span_action.name == action.__name__
+    assert span_action.attributes == {}
+    assert count_sleep.name == count.__name__
+    assert count_sleep.attributes == {}
+    assert count_sleep.parent.span_id == span_action.context.span_id  # type: ignore[union-attr]