diff options
Diffstat (limited to 'venv/lib/python3.11/site-packages/litestar/contrib/prometheus')
8 files changed, 303 insertions, 0 deletions
diff --git a/venv/lib/python3.11/site-packages/litestar/contrib/prometheus/__init__.py b/venv/lib/python3.11/site-packages/litestar/contrib/prometheus/__init__.py new file mode 100644 index 0000000..1ccb494 --- /dev/null +++ b/venv/lib/python3.11/site-packages/litestar/contrib/prometheus/__init__.py @@ -0,0 +1,5 @@ +from .config import PrometheusConfig +from .controller import PrometheusController +from .middleware import PrometheusMiddleware + +__all__ = ("PrometheusMiddleware", "PrometheusConfig", "PrometheusController") diff --git a/venv/lib/python3.11/site-packages/litestar/contrib/prometheus/__pycache__/__init__.cpython-311.pyc b/venv/lib/python3.11/site-packages/litestar/contrib/prometheus/__pycache__/__init__.cpython-311.pyc Binary files differnew file mode 100644 index 0000000..c6c5558 --- /dev/null +++ b/venv/lib/python3.11/site-packages/litestar/contrib/prometheus/__pycache__/__init__.cpython-311.pyc diff --git a/venv/lib/python3.11/site-packages/litestar/contrib/prometheus/__pycache__/config.cpython-311.pyc b/venv/lib/python3.11/site-packages/litestar/contrib/prometheus/__pycache__/config.cpython-311.pyc Binary files differnew file mode 100644 index 0000000..a998104 --- /dev/null +++ b/venv/lib/python3.11/site-packages/litestar/contrib/prometheus/__pycache__/config.cpython-311.pyc diff --git a/venv/lib/python3.11/site-packages/litestar/contrib/prometheus/__pycache__/controller.cpython-311.pyc b/venv/lib/python3.11/site-packages/litestar/contrib/prometheus/__pycache__/controller.cpython-311.pyc Binary files differnew file mode 100644 index 0000000..012b4be --- /dev/null +++ b/venv/lib/python3.11/site-packages/litestar/contrib/prometheus/__pycache__/controller.cpython-311.pyc diff --git a/venv/lib/python3.11/site-packages/litestar/contrib/prometheus/__pycache__/middleware.cpython-311.pyc b/venv/lib/python3.11/site-packages/litestar/contrib/prometheus/__pycache__/middleware.cpython-311.pyc Binary files differnew file mode 100644 index 0000000..2c08508 --- /dev/null +++ b/venv/lib/python3.11/site-packages/litestar/contrib/prometheus/__pycache__/middleware.cpython-311.pyc diff --git a/venv/lib/python3.11/site-packages/litestar/contrib/prometheus/config.py b/venv/lib/python3.11/site-packages/litestar/contrib/prometheus/config.py new file mode 100644 index 0000000..b77dab0 --- /dev/null +++ b/venv/lib/python3.11/site-packages/litestar/contrib/prometheus/config.py @@ -0,0 +1,64 @@ +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import TYPE_CHECKING, Callable, Mapping, Sequence + +from litestar.contrib.prometheus.middleware import ( + PrometheusMiddleware, +) +from litestar.exceptions import MissingDependencyException +from litestar.middleware.base import DefineMiddleware + +__all__ = ("PrometheusConfig",) + + +try: + import prometheus_client # noqa: F401 +except ImportError as e: + raise MissingDependencyException("prometheus_client", "prometheus-client", "prometheus") from e + + +if TYPE_CHECKING: + from litestar.connection.request import Request + from litestar.types import Method, Scopes + + +@dataclass +class PrometheusConfig: + """Configuration class for the PrometheusConfig middleware.""" + + app_name: str = field(default="litestar") + """The name of the application to use in the metrics.""" + prefix: str = "litestar" + """The prefix to use for the metrics.""" + labels: Mapping[str, str | Callable] | None = field(default=None) + """A mapping of labels to add to the metrics. The values can be either a string or a callable that returns a string.""" + exemplars: Callable[[Request], dict] | None = field(default=None) + """A callable that returns a list of exemplars to add to the metrics. Only supported in opementrics-text exposition format.""" + buckets: list[str | float] | None = field(default=None) + """A list of buckets to use for the histogram.""" + excluded_http_methods: Method | Sequence[Method] | None = field(default=None) + """A list of http methods to exclude from the metrics.""" + exclude_unhandled_paths: bool = field(default=False) + """Whether to ignore requests for unhandled paths from the metrics.""" + exclude: str | list[str] | None = field(default=None) + """A pattern or list of patterns for routes to exclude from the metrics.""" + exclude_opt_key: str | None = field(default=None) + """A key or list of keys in ``opt`` with which a route handler can "opt-out" of the middleware.""" + scopes: Scopes | None = field(default=None) + """ASGI scopes processed by the middleware, if None both ``http`` and ``websocket`` will be processed.""" + middleware_class: type[PrometheusMiddleware] = field(default=PrometheusMiddleware) + """The middleware class to use. + """ + + @property + def middleware(self) -> DefineMiddleware: + """Create an instance of :class:`DefineMiddleware <litestar.middleware.base.DefineMiddleware>` that wraps with. + + [PrometheusMiddleware][litestar.contrib.prometheus.PrometheusMiddleware]. or a subclass + of this middleware. + + Returns: + An instance of ``DefineMiddleware``. + """ + return DefineMiddleware(self.middleware_class, config=self) diff --git a/venv/lib/python3.11/site-packages/litestar/contrib/prometheus/controller.py b/venv/lib/python3.11/site-packages/litestar/contrib/prometheus/controller.py new file mode 100644 index 0000000..15f5bf1 --- /dev/null +++ b/venv/lib/python3.11/site-packages/litestar/contrib/prometheus/controller.py @@ -0,0 +1,53 @@ +from __future__ import annotations + +import os + +from litestar import Controller, get +from litestar.exceptions import MissingDependencyException +from litestar.response import Response + +try: + import prometheus_client # noqa: F401 +except ImportError as e: + raise MissingDependencyException("prometheus_client", "prometheus-client", "prometheus") from e + +from prometheus_client import ( + CONTENT_TYPE_LATEST, + REGISTRY, + CollectorRegistry, + generate_latest, + multiprocess, +) +from prometheus_client.openmetrics.exposition import ( + CONTENT_TYPE_LATEST as OPENMETRICS_CONTENT_TYPE_LATEST, +) +from prometheus_client.openmetrics.exposition import ( + generate_latest as openmetrics_generate_latest, +) + +__all__ = [ + "PrometheusController", +] + + +class PrometheusController(Controller): + """Controller for Prometheus endpoints.""" + + path: str = "/metrics" + """The path to expose the metrics on.""" + openmetrics_format: bool = False + """Whether to expose the metrics in OpenMetrics format.""" + + @get() + async def get(self) -> Response: + registry = REGISTRY + if "prometheus_multiproc_dir" in os.environ or "PROMETHEUS_MULTIPROC_DIR" in os.environ: + registry = CollectorRegistry() + multiprocess.MultiProcessCollector(registry) # type: ignore[no-untyped-call] + + if self.openmetrics_format: + headers = {"Content-Type": OPENMETRICS_CONTENT_TYPE_LATEST} + return Response(openmetrics_generate_latest(registry), status_code=200, headers=headers) # type: ignore[no-untyped-call] + + headers = {"Content-Type": CONTENT_TYPE_LATEST} + return Response(generate_latest(registry), status_code=200, headers=headers) diff --git a/venv/lib/python3.11/site-packages/litestar/contrib/prometheus/middleware.py b/venv/lib/python3.11/site-packages/litestar/contrib/prometheus/middleware.py new file mode 100644 index 0000000..50bc7cb --- /dev/null +++ b/venv/lib/python3.11/site-packages/litestar/contrib/prometheus/middleware.py @@ -0,0 +1,181 @@ +from __future__ import annotations + +import time +from functools import wraps +from typing import TYPE_CHECKING, Any, Callable, ClassVar, cast + +from litestar.connection.request import Request +from litestar.enums import ScopeType +from litestar.exceptions import MissingDependencyException +from litestar.middleware.base import AbstractMiddleware + +__all__ = ("PrometheusMiddleware",) + +from litestar.status_codes import HTTP_500_INTERNAL_SERVER_ERROR + +try: + import prometheus_client # noqa: F401 +except ImportError as e: + raise MissingDependencyException("prometheus_client", "prometheus-client", "prometheus") from e + +from prometheus_client import Counter, Gauge, Histogram + +if TYPE_CHECKING: + from prometheus_client.metrics import MetricWrapperBase + + from litestar.contrib.prometheus import PrometheusConfig + from litestar.types import ASGIApp, Message, Receive, Scope, Send + + +class PrometheusMiddleware(AbstractMiddleware): + """Prometheus Middleware.""" + + _metrics: ClassVar[dict[str, MetricWrapperBase]] = {} + + def __init__(self, app: ASGIApp, config: PrometheusConfig) -> None: + """Middleware that adds Prometheus instrumentation to the application. + + Args: + app: The ``next`` ASGI app to call. + config: An instance of :class:`PrometheusConfig <.contrib.prometheus.PrometheusConfig>` + """ + super().__init__(app=app, scopes=config.scopes, exclude=config.exclude, exclude_opt_key=config.exclude_opt_key) + self._config = config + self._kwargs: dict[str, Any] = {} + + if self._config.buckets is not None: + self._kwargs["buckets"] = self._config.buckets + + def request_count(self, labels: dict[str, str | int | float]) -> Counter: + metric_name = f"{self._config.prefix}_requests_total" + + if metric_name not in PrometheusMiddleware._metrics: + PrometheusMiddleware._metrics[metric_name] = Counter( + name=metric_name, + documentation="Total requests", + labelnames=[*labels.keys()], + ) + + return cast("Counter", PrometheusMiddleware._metrics[metric_name]) + + def request_time(self, labels: dict[str, str | int | float]) -> Histogram: + metric_name = f"{self._config.prefix}_request_duration_seconds" + + if metric_name not in PrometheusMiddleware._metrics: + PrometheusMiddleware._metrics[metric_name] = Histogram( + name=metric_name, + documentation="Request duration, in seconds", + labelnames=[*labels.keys()], + **self._kwargs, + ) + return cast("Histogram", PrometheusMiddleware._metrics[metric_name]) + + def requests_in_progress(self, labels: dict[str, str | int | float]) -> Gauge: + metric_name = f"{self._config.prefix}_requests_in_progress" + + if metric_name not in PrometheusMiddleware._metrics: + PrometheusMiddleware._metrics[metric_name] = Gauge( + name=metric_name, + documentation="Total requests currently in progress", + labelnames=[*labels.keys()], + multiprocess_mode="livesum", + ) + return cast("Gauge", PrometheusMiddleware._metrics[metric_name]) + + def requests_error_count(self, labels: dict[str, str | int | float]) -> Counter: + metric_name = f"{self._config.prefix}_requests_error_total" + + if metric_name not in PrometheusMiddleware._metrics: + PrometheusMiddleware._metrics[metric_name] = Counter( + name=metric_name, + documentation="Total errors in requests", + labelnames=[*labels.keys()], + ) + return cast("Counter", PrometheusMiddleware._metrics[metric_name]) + + def _get_extra_labels(self, request: Request[Any, Any, Any]) -> dict[str, str]: + """Get extra labels provided by the config and if they are callable, parse them. + + Args: + request: The request object. + + Returns: + A dictionary of extra labels. + """ + + return {k: str(v(request) if callable(v) else v) for k, v in (self._config.labels or {}).items()} + + def _get_default_labels(self, request: Request[Any, Any, Any]) -> dict[str, str | int | float]: + """Get default label values from the request. + + Args: + request: The request object. + + Returns: + A dictionary of default labels. + """ + + return { + "method": request.method if request.scope["type"] == ScopeType.HTTP else request.scope["type"], + "path": request.url.path, + "status_code": 200, + "app_name": self._config.app_name, + } + + async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None: + """ASGI callable. + + Args: + scope: The ASGI connection scope. + receive: The ASGI receive function. + send: The ASGI send function. + + Returns: + None + """ + + request = Request[Any, Any, Any](scope, receive) + + if self._config.excluded_http_methods and request.method in self._config.excluded_http_methods: + await self.app(scope, receive, send) + return + + labels = {**self._get_default_labels(request), **self._get_extra_labels(request)} + + request_span = {"start_time": time.perf_counter(), "end_time": 0, "duration": 0, "status_code": 200} + + wrapped_send = self._get_wrapped_send(send, request_span) + + self.requests_in_progress(labels).labels(*labels.values()).inc() + + try: + await self.app(scope, receive, wrapped_send) + finally: + extra: dict[str, Any] = {} + if self._config.exemplars: + extra["exemplar"] = self._config.exemplars(request) + + self.requests_in_progress(labels).labels(*labels.values()).dec() + + labels["status_code"] = request_span["status_code"] + label_values = [*labels.values()] + + if request_span["status_code"] >= HTTP_500_INTERNAL_SERVER_ERROR: + self.requests_error_count(labels).labels(*label_values).inc(**extra) + + self.request_count(labels).labels(*label_values).inc(**extra) + self.request_time(labels).labels(*label_values).observe(request_span["duration"], **extra) + + def _get_wrapped_send(self, send: Send, request_span: dict[str, float]) -> Callable: + @wraps(send) + async def wrapped_send(message: Message) -> None: + if message["type"] == "http.response.start": + request_span["status_code"] = message["status"] + + if message["type"] == "http.response.body": + end = time.perf_counter() + request_span["duration"] = end - request_span["start_time"] + request_span["end_time"] = end + await send(message) + + return wrapped_send |