kytos.core.apm

Complexity

Total Complexity

12

Size/Duplication

Total Lines	83
Duplicated Lines	0 %

Test Coverage

Coverage

97.78%

Importance

Changes

0

3 Methods

Rating	Name	Duplication	Size	Complexity
A	ElasticAPM.init_flask_app()	0	7	2
A	ElasticAPM.init_client()	0	21	3
A	ElasticAPM.get_client()	0	6	2

"""APM module."""
import os
from functools import wraps

from elasticapm import Client
from elasticapm.contrib.flask import ElasticAPM as FlaskAPM
from elasticapm.traces import execution_context

from kytos.core.exceptions import KytosAPMInitException


def init_apm(apm_backend="es", **kwargs):
    """Init APM backend."""
    backends = {"es": ElasticAPM}
    try:
        return backends[apm_backend].init_client(**kwargs)
    except KeyError:
        client_names = ",".join(list(backends.keys()))
        raise KytosAPMInitException(
            f"APM backend '{apm_backend}' isn't supported."
            f" Current supported APMs: {client_names}"
        )


def begin_span(func, span_type="custom"):
    """APM transaction begin_span decorator."""

    @wraps(func)
    def wrapper(*args, **kwds):
        transaction = execution_context.get_transaction()
        if transaction:
            transaction.begin_span(func.__name__, span_type)
        result = func(*args, **kwds)
        if transaction:
            transaction.end_span()
        return result

    return wrapper


class ElasticAPM:
    """ElasticAPM Client instance."""

    _flask_apm: FlaskAPM = None
    _client: Client = None

    @classmethod
    def get_client(cls, **kwargs) -> Client:
        """Get client."""
        if not cls._client:
            return cls.init_client(**kwargs)
        return cls._client

    @classmethod
    def init_flask_app(cls, app) -> FlaskAPM:
        """Init Flask APM Client."""
        if not cls._flask_apm:
            return None
        cls._flask_apm.init_app(app)
        return cls._flask_apm

    @classmethod
    def init_client(
        cls,
        service_name=os.environ.get("ELASTIC_APM_SERVICE_NAME", "kytos"),
        server_url=os.environ.get("ELASTIC_APM_URL", "http://localhost:8200"),
        secret_token=os.environ.get("ELASTIC_APM_SECRET_TOKEN",
                                    "elasticapm_token"),
        **kwargs,
    ) -> Client:
        """Init APM Client."""
        app = kwargs.pop("app", None)
        if not cls._client:
            cls._client = Client(
                service_name=service_name,
                server_url=server_url,
                secret_token=secret_token,
                **kwargs,
            )
        if not cls._flask_apm:
            cls._flask_apm = FlaskAPM(client=cls._client, app=app)
        return cls._client