Implement tracing decorator and add it to example service

Author: febus982

.idea/inspectionProfiles/Project_Default.xml

@@ -0,0 +1,56 @@
+import asyncio
+from functools import wraps
+
+from opentelemetry import trace
+
+# Get the _tracer instance (You can set your own _tracer name)
+tracer = trace.get_tracer(__name__)
+
+
+def trace_function(trace_attributes: bool = True, trace_result: bool = True):
+    """
+    Decorator to trace callables using OpenTelemetry spans.
+
+    Parameters:
+    - trace_attributes (bool): If False, disables adding function arguments to the span.
+    - trace_result (bool): If False, disables adding the function's result to the span.
+    """
+
+    def decorator(func):
+        @wraps(func)
+        def sync_or_async_wrapper(*args, **kwargs):
+            with tracer.start_as_current_span(func.__name__) as span:
+                try:
+                    # Set function arguments as attributes
+                    if trace_attributes:
+                        span.set_attribute("function.args", str(args))
+                        span.set_attribute("function.kwargs", str(kwargs))
+
+                    async def async_handler():
+                        result = await func(*args, **kwargs)
+                        # Add result to span
+                        if trace_result:
+                            span.set_attribute("function.result", str(result))
+                        return result
+
+                    def sync_handler():
+                        result = func(*args, **kwargs)
+                        # Add result to span
+                        if trace_result:
+                            span.set_attribute("function.result", str(result))
+                        return result
+
+                    if asyncio.iscoroutinefunction(func):
+                        return async_handler()
+                    else:
+                        return sync_handler()
+
+                except Exception as e:
+                    # Record the exception in the span
+                    span.record_exception(e)
+                    span.set_status(trace.status.Status(trace.status.StatusCode.ERROR))
+                    raise
+
+        return sync_or_async_wrapper
+
+    return decorator

src/common/tracing.py

@@ -4,13 +4,16 @@
 from dependency_injector.wiring import Provide, inject
 from structlog import get_logger
 
+from common.tracing import trace_function
+from common.utils import apply_decorator_to_methods
 from ._gateway_interfaces import BookEventGatewayInterface, BookRepositoryInterface
 from ._models import BookModel
 from ._tasks import book_cpu_intensive_task
 from .dto import Book, BookData
 from .events import BookCreatedV1, BookCreatedV1Data
 
 
+@apply_decorator_to_methods(trace_function())
 class BookService:
     _book_repository: BookRepositoryInterface
     _event_gateway: BookEventGatewayInterface

src/domains/books/_service.py

@@ -0,0 +1,156 @@
+import asyncio
+from unittest.mock import MagicMock, call, patch
+
+import pytest
+
+from common.tracing import trace_function
+
+
+@pytest.fixture
+def mock_tracer():
+    """
+    Fixture to mock the OpenTelemetry tracer and span.
+    """
+    mock_tracer = MagicMock()
+    mock_span = MagicMock()
+    mock_tracer.start_as_current_span.return_value.__enter__.return_value = mock_span
+
+    with (
+        patch("opentelemetry.trace.get_tracer", return_value=mock_tracer),
+        patch("common.tracing.tracer", mock_tracer),
+    ):
+        yield mock_tracer, mock_span
+
+
+def test_sync_function_default_params(mock_tracer):
+    """
+    Test a synchronous function with default decorator parameters.
+    """
+    mock_tracer, mock_span = mock_tracer
+
+    # Define a sync function to wrap with the decorator
+    @trace_function()
+    def add_nums(a, b):
+        return a + b
+
+    # Call the function
+    result = add_nums(2, 3)
+
+    # Assertions
+    assert result == 5
+    mock_tracer.start_as_current_span.assert_called_once_with("add_nums")
+    mock_span.set_attribute.assert_any_call("function.args", "(2, 3)")
+    mock_span.set_attribute.assert_any_call("function.result", "5")
+
+
+@pytest.mark.asyncio
+async def test_async_function_default_params(mock_tracer):
+    """
+    Test an asynchronous function with default decorator parameters.
+    """
+    mock_tracer, mock_span = mock_tracer
+
+    # Define an async function to wrap with the decorator
+    @trace_function()
+    async def async_func(a, b):
+        await asyncio.sleep(0.1)
+        return a * b
+
+    # Run the async function
+    result = await async_func(4, 5)
+
+    # Assertions
+    assert result == 20
+    mock_tracer.start_as_current_span.assert_called_once_with("async_func")
+    mock_span.set_attribute.assert_any_call("function.args", "(4, 5)")
+    mock_span.set_attribute.assert_any_call("function.result", "20")
+
+
+def test_disable_function_attributes(mock_tracer):
+    """
+    Test a synchronous function with `add_function_attributes` set to False.
+    """
+    mock_tracer, mock_span = mock_tracer
+
+    # Define a sync function with attributes disabled
+    @trace_function(trace_attributes=False)
+    def sync_func(a, b):
+        return a - b
+
+    # Call the function
+    result = sync_func(10, 6)
+
+    # Assertions
+    assert result == 4
+    mock_tracer.start_as_current_span.assert_called_once_with("sync_func")
+    mock_span.set_attribute.assert_any_call("function.result", "4")
+    assert (
+        call("function.args", "(10, 6)") not in mock_span.set_attribute.call_args_list
+    )
+
+
+def test_disable_result_in_span_sync(mock_tracer):
+    """
+    Test an asynchronous function with `add_result_to_span` set to False.
+    """
+    mock_tracer, mock_span = mock_tracer
+
+    # Define an async function with result disabled
+    @trace_function(trace_result=False)
+    def sync_func(a, b):
+        return a / b
+
+    # Run the async function
+    result = sync_func(10, 2)
+
+    # Assertions
+    assert result == 5.0
+    mock_tracer.start_as_current_span.assert_called_once_with("sync_func")
+    mock_span.set_attribute.assert_any_call("function.args", "(10, 2)")
+    assert call("function.result") not in mock_span.set_attribute.call_args_list
+
+
+@pytest.mark.asyncio
+async def test_disable_result_in_span(mock_tracer):
+    """
+    Test an asynchronous function with `add_result_to_span` set to False.
+    """
+    mock_tracer, mock_span = mock_tracer
+
+    # Define an async function with result disabled
+    @trace_function(trace_result=False)
+    async def async_func(a, b):
+        await asyncio.sleep(0.1)
+        return a / b
+
+    # Run the async function
+    result = await async_func(10, 2)
+
+    # Assertions
+    assert result == 5.0
+    mock_tracer.start_as_current_span.assert_called_once_with("async_func")
+    mock_span.set_attribute.assert_any_call("function.args", "(10, 2)")
+    assert call("function.result") not in mock_span.set_attribute.call_args_list
+
+
+def test_exception_in_function(mock_tracer):
+    """
+    Test behavior when the function raises an exception.
+    """
+    mock_tracer, mock_span = mock_tracer
+
+    # Define a failing function
+    @trace_function()
+    def failing_func(a, b):
+        if b == 0:
+            raise ValueError("Division by zero!")
+        return a / b
+
+    # Use pytest to assert the exception is raised
+    with pytest.raises(ValueError, match="Division by zero!"):
+        failing_func(10, 0)
+
+    # Assertions
+    mock_tracer.start_as_current_span.assert_called_once_with("failing_func")
+    mock_span.record_exception.assert_called_once()
+    mock_span.set_status.assert_called_once()