otel instrumentation for outgoing calls using generated clientlibs

setup.py

@@ -16,6 +16,7 @@
     python_requires=">=3.7",
     install_requires=[
         'py_zipkin>=0.10.1',
+        'opentelemetry-sdk>=0.26.1',
     ],
     keywords='zipkin',
     classifiers=[

swagger_zipkin/otel_decorator.py

@@ -0,0 +1,165 @@
+from __future__ import annotations
+
+import importlib
+from contextlib import contextmanager
+from typing import Any
+from typing import TYPE_CHECKING
+from typing import TypeVar
+
+from opentelemetry import trace
+from opentelemetry.trace.propagation.tracecontext import TraceContextTextMapPropagator
+from opentelemetry.trace.span import format_span_id
+from opentelemetry.trace.span import format_trace_id
+from opentelemetry.trace.span import TraceFlags
+from typing_extensions import ParamSpec
+
+from swagger_zipkin.decorate_client import Client
+from swagger_zipkin.decorate_client import decorate_client
+from swagger_zipkin.decorate_client import Resource
+
+if TYPE_CHECKING:
+    import pyramid.request.Request   # type: ignore # noqa: F401
+
+T = TypeVar('T', covariant=True)
+P = ParamSpec('P')
+
+tracer = trace.get_tracer("otel_decorator")
+
+
+class OtelResourceDecorator:
+    """A wrapper to the swagger resource.
+
+    :param resource: A resource object. eg. `client.pet`, `client.store`.
+    :type resource: :class:`swaggerpy.client.Resource` or :class:`bravado_core.resource.Resource`
+    """
+
+    def __init__(self, resource: Client, client_identifier: str, smartstack_namespace: str) -> None:
+        self.resource = resource
+        self.client_identifier = client_identifier
+        self.smartstack_namespace = smartstack_namespace
+
+    def __getattr__(self, name: str) -> Resource:
+        return decorate_client(self.resource, self.with_headers, name)
+
+    def with_headers(self, call_name: str, *args: Any, **kwargs: Any) -> Any:
+        kwargs.setdefault('_request_options', {})
+        request_options: dict = kwargs['_request_options']
+        request_options.setdefault('headers', {})
+
+        # what is the right way to get the Request object. can we use contruct_request
+        # https://github.com/Yelp/bravado/blob/master/bravado/client.py#L283C5-L283C22
+        # this would create a bravado dependency.
+        request = get_pyramid_current_request()
+        http_route = getattr(request, "matched_route", "")
+        http_request_method = getattr(request, "method", "")
+
+        parent_span = trace.get_current_span()
+        span_name = f"{http_request_method} {http_route}"
+        with tracer.start_as_current_span(
+            span_name, kind=trace.SpanKind.CLIENT
+        ) as span:
+            self.inject_otel_headers(kwargs, span)
+            self.inject_zipkin_headers(kwargs, span, parent_span)
+
+            # ideally the exception should be scoped for self.with_headers
+            # handle_exception should handle general excetion and the specific exception HTTPError
+            # But this would create a dependency on bravado package. Is this ok?
+            # Assuming resource.operation does throw HTTPError
+            # https://github.com/Yelp/bravado/blob/master/bravado/exception.py
+            with self.handle_exception():
+                span.set_attribute("url.path", getattr(request, "path", ""))
+                span.set_attribute("http.route", http_route)
+                span.set_attribute("http.request.method", http_request_method)
+
+                span.set_attribute("client.namespace", self.client_identifier)
+                span.set_attribute("peer.service", self.smartstack_namespace)
+                span.set_attribute("server.namespace", self.smartstack_namespace)
+                span.set_attribute("http.response.status_code", "200")
+
+                return getattr(self.resource, call_name)(*args, **kwargs)
+
+    @contextmanager
+    def handle_exception(
+        self,
+    ) -> Any:
+        try:
+            yield
+        except Exception as e:
+            span = trace.get_current_span()
+            span.set_attribute("error.type", e.__class__.__name__)
+            span.set_attribute("http.response.status_code", "500")
+            raise e
+
+    def __dir__(self) -> list[str]:
+        return dir(self.resource)  # pragma: no cover
+
+    def inject_otel_headers(
+        self, kwargs: dict[str, Any], current_span: trace.Span
+    ) -> None:
+        propagator = TraceContextTextMapPropagator()
+        carrier = kwargs['_request_options']["headers"]
+        propagator.inject(carrier=carrier, context=trace.set_span_in_context(current_span))
+
+    def inject_zipkin_headers(
+            self, kwargs: dict[str, Any], current_span: trace.Span, parent_span: trace.Span
+    ) -> None:
+        current_span_context = current_span.get_span_context()
+        kwargs["_request_options"]["headers"]["X-B3-TraceId"] = format_trace_id(
+            current_span_context.trace_id
+        )
+        kwargs["_request_options"]["headers"]["X-B3-SpanId"] = format_span_id(
+            current_span_context.span_id
+        )
+        if parent_span is not None and parent_span.is_recording():
+            parent_span_context = parent_span.get_span_context()
+            kwargs["_request_options"]["headers"]["X-B3-ParentSpanId"] = format_span_id(
+                parent_span_context.span_id)
+
+        kwargs["_request_options"]["headers"]["X-B3-Sampled"] = (
+            "1"
+            if (current_span_context.trace_flags & TraceFlags.SAMPLED == TraceFlags.SAMPLED)
+            else "0"
+        )
+        kwargs["_request_options"]["headers"]["X-B3-Flags"] = "0"
+
+
+class OtelClientDecorator:
+    """A wrapper to swagger client (swagger-py or bravado) to pass on otel and zipkin
+    headers to the service call. It will also generate a CLIENT span for the outgoing call.
+
+    Even though client is initialised once, all the calls made will have
+    independent spans.
+
+    :param client: Swagger Client
+    :type client: :class:`swaggerpy.client.SwaggerClient` or :class:`bravado.client.SwaggerClient`.
+    :param client_identifier: the name of the service that is using this
+            generated clientlib
+    :type client_identifier: string
+    :param smartstack_namespace: the smartstack name of the paasta instance
+            this generated clientlib is hitting
+    :type smartstack_namespace: string
+    """
+
+    def __init__(self, client: Client, client_identifier: str, smartstack_namespace: str):
+        self._client = client
+        self.client_identifier = client_identifier
+        self.smartstack_namespace = smartstack_namespace
+
+    def __getattr__(self, name: str) -> Client:
+        return OtelResourceDecorator(
+            getattr(self._client, name),
+            client_identifier=self.client_identifier,
+            smartstack_namespace=self.smartstack_namespace,
+        )
+
+    def __dir__(self) -> list[str]:
+        return dir(self._client)  # pragma: no cover
+
+
+def get_pyramid_current_request() -> pyramid.request.Request | None:
+    try:
+        threadlocal = importlib.import_module("pyramid.threadlocal")
+    except ImportError:
+        return None
+
+    return threadlocal.get_current_request()

tests/otel_decorator_test.py

@@ -0,0 +1,187 @@
+from unittest import mock
+
+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 opentelemetry.trace import SpanKind
+from opentelemetry.trace.span import format_span_id
+from opentelemetry.trace.span import format_trace_id
+
+from swagger_zipkin.otel_decorator import OtelClientDecorator
+from swagger_zipkin.otel_decorator import OtelResourceDecorator
+
+memory_exporter = InMemorySpanExporter()
+span_processor = SimpleSpanProcessor(memory_exporter)
+trace.set_tracer_provider(TracerProvider())
+trace.get_tracer_provider().add_span_processor(span_processor)
+
+client_identifier = "test_client"
+smartstack_namespace = "smartstack_namespace"
+tracer = trace.get_tracer("otel_decorator")
+
+
+@pytest.fixture
+def setup():
+    memory_exporter.clear()
+
+
+@pytest.fixture
+def get_request():
+    mock_request = mock.Mock()
+    mock_request.path = "/sample-url"
+    mock_request.method = "GET"
+    mock_request.matched_route = "sample-view"
+    return mock_request
+
+
+def create_request_options(parent_span: trace.Span, exported_span: trace.Span):
+    trace_id = format_trace_id(exported_span.get_span_context().trace_id)
+    span_id = format_span_id(exported_span.get_span_context().span_id)
+
+    headers = {}
+    headers['headers'] = {
+        'traceparent': f'00-{trace_id}-{span_id}-01',
+        'X-B3-TraceId': trace_id,
+        'X-B3-SpanId': span_id,
+        'X-B3-Flags': '0',
+        'X-B3-Sampled': '1',
+    }
+    if parent_span is not None:
+        headers['headers']['X-B3-ParentSpanId'] = format_span_id(parent_span.get_span_context().span_id)
+
+    return headers
+
+
+@mock.patch(
+    "swagger_zipkin.otel_decorator.get_pyramid_current_request", autospec=True
+)
+def test_client_request(mock_request, get_request, setup):
+    mock_request.return_value = get_request
+
+    with tracer.start_as_current_span(
+        "parent_span", kind=trace.SpanKind.SERVER
+    ) as parent_span:
+        client = mock.Mock()
+        wrapped_client = OtelClientDecorator(
+            client,
+            client_identifier=client_identifier,
+            smartstack_namespace=smartstack_namespace
+        )
+        resource = wrapped_client.resource
+        param1 = mock.Mock()
+        resource.operation(param1)
+
+        assert len(memory_exporter.get_finished_spans()) == 1
+        exported_span = memory_exporter.get_finished_spans()[0]
+
+        client.resource.operation.assert_called_with(
+            param1,
+            _request_options=create_request_options(parent_span, exported_span)
+        )
+
+        assert exported_span.kind == SpanKind.CLIENT
+        assert exported_span.name == f"{get_request.method} {get_request.matched_route}"
+        assert exported_span.attributes["url.path"] == get_request.path
+        assert exported_span.attributes["http.request.method"] == get_request.method
+        assert exported_span.attributes["http.route"] == get_request.matched_route
+        assert exported_span.attributes["client.namespace"] == client_identifier
+        assert exported_span.attributes["peer.service"] == smartstack_namespace
+        assert exported_span.attributes["server.namespace"] == smartstack_namespace
+        assert exported_span.attributes["http.response.status_code"] == "200"
+
+        param2 = mock.Mock()
+        resource.operation(param2)
+
+        assert len(memory_exporter.get_finished_spans()) == 2
+        exported_span = memory_exporter.get_finished_spans()[1]
+
+        client.resource.operation.assert_called_with(
+            param2,
+            _request_options=create_request_options(parent_span, exported_span)
+        )
+
+        assert exported_span.name == f"{get_request.method} {get_request.matched_route}"
+        assert exported_span.attributes["url.path"] == get_request.path
+        assert exported_span.attributes["http.request.method"] == get_request.method
+        assert exported_span.attributes["http.route"] == get_request.matched_route
+        assert exported_span.attributes["client.namespace"] == client_identifier
+        assert exported_span.attributes["peer.service"] == smartstack_namespace
+        assert exported_span.attributes["server.namespace"] == smartstack_namespace
+        assert exported_span.attributes["http.response.status_code"] == "200"
+
+
+@mock.patch(
+    "swagger_zipkin.otel_decorator.get_pyramid_current_request", autospec=True
+)
+def test_client_request_no_parent_span(mock_request, get_request, setup):
+    mock_request.return_value = get_request
+
+    client = mock.Mock()
+    wrapped_client = OtelClientDecorator(
+        client,
+        client_identifier=client_identifier,
+        smartstack_namespace=smartstack_namespace
+    )
+    resource = wrapped_client.resource
+    param = mock.Mock()
+    resource.operation(param)
+
+    assert len(memory_exporter.get_finished_spans()) == 1
+    exported_span = memory_exporter.get_finished_spans()[0]
+
+    client.resource.operation.assert_called_with(
+        param,
+        _request_options=create_request_options(None, exported_span)
+    )
+
+    assert exported_span.kind == SpanKind.CLIENT
+    assert exported_span.name == f"{get_request.method} {get_request.matched_route}"
+    assert exported_span.attributes["url.path"] == get_request.path
+    assert exported_span.attributes["http.request.method"] == get_request.method
+    assert exported_span.attributes["http.route"] == get_request.matched_route
+    assert exported_span.attributes["client.namespace"] == client_identifier
+    assert exported_span.attributes["peer.service"] == smartstack_namespace
+    assert exported_span.attributes["server.namespace"] == smartstack_namespace
+    assert exported_span.attributes["http.response.status_code"] == "200"
+
+
+@mock.patch(
+    "swagger_zipkin.otel_decorator.get_pyramid_current_request", autospec=True
+)
+def test_with_headers_exception(mock_request, get_request, setup):
+    mock_request.return_value = get_request
+
+    # Create a mock resource and configure it to raise an exception
+    mock_resource = mock.MagicMock()
+    mock_method = mock.MagicMock(side_effect=Exception("simulated exception"))
+    setattr(mock_resource, 'test_operation', mock_method)
+
+    decorator = OtelResourceDecorator(resource=mock_resource, client_identifier="test_client",
+                                      smartstack_namespace="smartstack_namespace")
+
+    # Prepare arguments
+    args = ()
+    kwargs = {'_request_options': {'headers': {}}}
+
+    with pytest.raises(Exception):
+        decorator.with_headers("test_operation", *args, **kwargs)
+
+    assert len(memory_exporter.get_finished_spans()) == 1
+    exported_span = memory_exporter.get_finished_spans()[0]
+
+    expected_headers = kwargs['_request_options']['headers']
+    actual_headers = create_request_options(None, exported_span)['headers']
+    assert expected_headers == actual_headers
+
+    assert exported_span.kind == SpanKind.CLIENT
+    assert exported_span.name == f"{get_request.method} {get_request.matched_route}"
+    assert exported_span.attributes["url.path"] == get_request.path
+    assert exported_span.attributes["http.request.method"] == get_request.method
+    assert exported_span.attributes["http.route"] == get_request.matched_route
+    assert exported_span.attributes["client.namespace"] == client_identifier
+    assert exported_span.attributes["peer.service"] == smartstack_namespace
+    assert exported_span.attributes["server.namespace"] == smartstack_namespace
+    assert exported_span.attributes["error.type"] == "Exception"
+    assert exported_span.attributes["http.response.status_code"] == "500"