Source code for opentelemetry.instrumentation.aio_pika.aio_pika_instrumentor

# Copyright The OpenTelemetry Authors
# SPDX-License-Identifier: Apache-2.0
from typing import Any, Callable, Collection

import wrapt
from aio_pika import Exchange, Queue
from aio_pika.abc import AbstractIncomingMessage

from opentelemetry import trace
from opentelemetry.instrumentation.aio_pika.callback_decorator import (
    CallbackDecorator,
)
from opentelemetry.instrumentation.aio_pika.package import _instruments
from opentelemetry.instrumentation.aio_pika.publish_decorator import (
    PublishDecorator,
)
from opentelemetry.instrumentation.aio_pika.version import __version__
from opentelemetry.instrumentation.instrumentor import BaseInstrumentor
from opentelemetry.instrumentation.utils import unwrap
from opentelemetry.trace import Tracer

_INSTRUMENTATION_MODULE_NAME = "opentelemetry.instrumentation.aio_pika"


class AioPikaInstrumentor(BaseInstrumentor):
    @staticmethod
    def _instrument_queue(tracer: Tracer):
        async def wrapper(wrapped, instance, args, kwargs):
            async def consume(
                callback: Callable[[AbstractIncomingMessage], Any],
                *fargs,
                **fkwargs,
            ):
                decorated_callback = CallbackDecorator(
                    tracer, instance
                ).decorate(callback)
                return await wrapped(decorated_callback, *fargs, **fkwargs)

            return await consume(*args, **kwargs)

        wrapt.wrap_function_wrapper(Queue, "consume", wrapper)

    @staticmethod
    def _instrument_exchange(tracer: Tracer):
        async def wrapper(wrapped, instance, args, kwargs):
            decorated_publish = PublishDecorator(tracer, instance).decorate(
                wrapped
            )
            return await decorated_publish(*args, **kwargs)

        wrapt.wrap_function_wrapper(Exchange, "publish", wrapper)

    def _instrument(self, **kwargs):
        tracer_provider = kwargs.get("tracer_provider", None)
        tracer = trace.get_tracer(
            _INSTRUMENTATION_MODULE_NAME,
            __version__,
            tracer_provider,
            schema_url="https://opentelemetry.io/schemas/1.11.0",
        )
        self._instrument_queue(tracer)
        self._instrument_exchange(tracer)

    @staticmethod
    def _uninstrument_queue():
        unwrap(Queue, "consume")

    @staticmethod
    def _uninstrument_exchange():
        unwrap(Exchange, "publish")

    def _uninstrument(self, **kwargs):
        self._uninstrument_queue()
        self._uninstrument_exchange()

    def instrumentation_dependencies(self) -> Collection[str]:
        return _instruments