diff --git a/microbootstrap/bootstrappers/fastapi.py b/microbootstrap/bootstrappers/fastapi.py index b567bf4..628a953 100644 --- a/microbootstrap/bootstrappers/fastapi.py +++ b/microbootstrap/bootstrappers/fastapi.py @@ -5,7 +5,7 @@ from fastapi.middleware.cors import CORSMiddleware from fastapi_offline_docs import enable_offline_docs from opentelemetry.instrumentation.fastapi import FastAPIInstrumentor -from prometheus_fastapi_instrumentator import Instrumentator +from prometheus_fastapi_instrumentator import Instrumentator, metrics from microbootstrap.bootstrappers.base import ApplicationBootstrapper from microbootstrap.config.fastapi import FastApiConfig @@ -113,7 +113,11 @@ def bootstrap_after(self, application: ApplicationT) -> ApplicationT: @FastApiBootstrapper.use_instrument() class FastApiPrometheusInstrument(PrometheusInstrument[FastApiPrometheusConfig]): def bootstrap_after(self, application: ApplicationT) -> ApplicationT: - Instrumentator(**self.instrument_config.prometheus_instrumentator_params).instrument( + Instrumentator(**self.instrument_config.prometheus_instrumentator_params).add( + metrics.default( + custom_labels=self.instrument_config.prometheus_custom_labels, + ), + ).instrument( application, **self.instrument_config.prometheus_instrument_params, ).expose( diff --git a/microbootstrap/bootstrappers/faststream.py b/microbootstrap/bootstrappers/faststream.py index 23e3d5d..7cd725e 100644 --- a/microbootstrap/bootstrappers/faststream.py +++ b/microbootstrap/bootstrappers/faststream.py @@ -91,7 +91,10 @@ def bootstrap_before(self) -> dict[str, typing.Any]: def bootstrap_after(self, application: AsgiFastStream) -> AsgiFastStream: # type: ignore[override] if self.instrument_config.prometheus_middleware_cls and application.broker: application.broker.add_middleware( - self.instrument_config.prometheus_middleware_cls(registry=prometheus_client.REGISTRY), + self.instrument_config.prometheus_middleware_cls( + registry=prometheus_client.REGISTRY, + custom_labels=self.instrument_config.prometheus_custom_labels, + ), ) return application diff --git a/microbootstrap/instruments/prometheus_instrument.py b/microbootstrap/instruments/prometheus_instrument.py index dc84795..9751d72 100644 --- a/microbootstrap/instruments/prometheus_instrument.py +++ b/microbootstrap/instruments/prometheus_instrument.py @@ -30,6 +30,7 @@ class FastApiPrometheusConfig(BasePrometheusConfig): prometheus_instrumentator_params: dict[str, typing.Any] = pydantic.Field(default_factory=dict) prometheus_instrument_params: dict[str, typing.Any] = pydantic.Field(default_factory=dict) prometheus_expose_params: dict[str, typing.Any] = pydantic.Field(default_factory=dict) + prometheus_custom_labels: dict[str, typing.Any] = pydantic.Field(default_factory=dict) @typing.runtime_checkable @@ -41,11 +42,13 @@ def __init__( app_name: str = ..., metrics_prefix: str = "faststream", received_messages_size_buckets: typing.Sequence[float] | None = None, + custom_labels: dict[str, str | typing.Callable[[typing.Any], str]] | None = None, ) -> None: ... class FastStreamPrometheusConfig(BasePrometheusConfig): prometheus_middleware_cls: type[FastStreamPrometheusMiddlewareProtocol] | None = None + prometheus_custom_labels: dict[str, typing.Any] = pydantic.Field(default_factory=dict) class PrometheusInstrument(Instrument[PrometheusConfigT]): diff --git a/pyproject.toml b/pyproject.toml index 0a61aa1..83125d7 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -72,7 +72,7 @@ litestar = [ "prometheus-client>=0.20", ] granian = ["granian[reload]>=1"] -faststream = ["faststream~=0.5", "prometheus-client>=0.20"] +faststream = ["faststream~=0.6.2", "prometheus-client>=0.20"] [dependency-groups] dev = [ diff --git a/tests/conftest.py b/tests/conftest.py index 53faf25..630fd56 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -5,11 +5,13 @@ import litestar import pytest +from prometheus_client import REGISTRY from sentry_sdk.transport import Transport as SentryTransport import microbootstrap.settings from microbootstrap import ( FastApiPrometheusConfig, + FastStreamPrometheusConfig, LitestarPrometheusConfig, LoggingConfig, OpentelemetryConfig, @@ -74,6 +76,11 @@ def minimal_litestar_prometheus_config() -> LitestarPrometheusConfig: return LitestarPrometheusConfig() +@pytest.fixture +def minimal_faststream_prometheus_config() -> FastStreamPrometheusConfig: + return FastStreamPrometheusConfig() + + @pytest.fixture def minimal_swagger_config() -> SwaggerConfig: return SwaggerConfig() @@ -132,3 +139,9 @@ def reset_reloaded_settings_module() -> typing.Iterator[None]: @pytest.fixture(autouse=True) def patch_out_entry_points(monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.setattr(opentelemetry_instrument, "entry_points", MagicMock(retrun_value=[])) + + +@pytest.fixture(autouse=True) +def clean_prometheus_registry() -> None: + REGISTRY._names_to_collectors.clear() # noqa: SLF001 + REGISTRY._collector_to_names.clear() # noqa: SLF001 diff --git a/tests/instruments/test_prometheus.py b/tests/instruments/test_prometheus.py index c170c09..1769027 100644 --- a/tests/instruments/test_prometheus.py +++ b/tests/instruments/test_prometheus.py @@ -2,15 +2,34 @@ import fastapi import litestar +import pytest from fastapi.testclient import TestClient as FastAPITestClient +from faststream.redis import RedisBroker, TestRedisBroker +from faststream.redis.prometheus import RedisPrometheusMiddleware from litestar import status_codes from litestar.middleware.base import DefineMiddleware from litestar.testing import TestClient as LitestarTestClient +from prometheus_client import REGISTRY -from microbootstrap import FastApiPrometheusConfig, LitestarPrometheusConfig +from microbootstrap import FastApiPrometheusConfig, FastStreamSettings, LitestarPrometheusConfig from microbootstrap.bootstrappers.fastapi import FastApiPrometheusInstrument +from microbootstrap.bootstrappers.faststream import FastStreamBootstrapper from microbootstrap.bootstrappers.litestar import LitestarPrometheusInstrument -from microbootstrap.instruments.prometheus_instrument import BasePrometheusConfig, PrometheusInstrument +from microbootstrap.config.faststream import FastStreamConfig +from microbootstrap.instruments.prometheus_instrument import ( + BasePrometheusConfig, + FastStreamPrometheusConfig, + PrometheusInstrument, +) + + +def check_is_metrics_has_labels(custom_labels_keys: set[str]) -> bool: + for metric in REGISTRY.collect(): + for sample in metric.samples: + label_keys = set(sample.labels.keys()) + if custom_labels_keys & label_keys: + return True + return False def test_prometheus_is_ready(minimal_base_prometheus_config: BasePrometheusConfig) -> None: @@ -85,3 +104,69 @@ def test_fastapi_prometheus_bootstrap_working(minimal_fastapi_prometheus_config: ) assert response.status_code == status_codes.HTTP_200_OK assert response.text + + +@pytest.mark.parametrize( + ("custom_labels", "expected_label_keys"), + [ + ({"test_label": "test_value"}, {"test_label"}), + ({}, {"method", "handler", "status"}), + ], +) +def test_fastapi_prometheus_custom_labels( + minimal_fastapi_prometheus_config: FastApiPrometheusConfig, + custom_labels: dict[str, str], + expected_label_keys: set[str], +) -> None: + minimal_fastapi_prometheus_config.prometheus_custom_labels = custom_labels + prometheus_instrument: typing.Final = FastApiPrometheusInstrument(minimal_fastapi_prometheus_config) + + fastapi_application = fastapi.FastAPI() + fastapi_application = prometheus_instrument.bootstrap_after(fastapi_application) + + response: typing.Final = FastAPITestClient(app=fastapi_application).get( + minimal_fastapi_prometheus_config.prometheus_metrics_path + ) + + assert response.status_code == status_codes.HTTP_200_OK + assert check_is_metrics_has_labels(expected_label_keys) + + +@pytest.mark.parametrize( + ("custom_labels", "expected_label_keys"), + [ + ({"test_label": "test_value"}, {"test_label"}), + ({}, {"app_name", "broker", "handler"}), + ], +) +async def test_faststream_prometheus_custom_labels( + minimal_faststream_prometheus_config: FastStreamPrometheusConfig, + custom_labels: dict[str, str], + expected_label_keys: set[str], +) -> None: + minimal_faststream_prometheus_config.prometheus_custom_labels = custom_labels + minimal_faststream_prometheus_config.prometheus_middleware_cls = RedisPrometheusMiddleware # type: ignore[assignment] + + broker: typing.Final = RedisBroker() + ( + FastStreamBootstrapper(FastStreamSettings()) + .configure_application(FastStreamConfig(broker=broker)) + .configure_instrument(minimal_faststream_prometheus_config) + .bootstrap() + ) + + def create_test_redis_subscriber( + broker: RedisBroker, + topic: str, + ) -> typing.Callable[[dict[str, str]], typing.Coroutine[typing.Any, typing.Any, None]]: + @broker.subscriber(topic) + async def test_subscriber(payload: dict[str, str]) -> None: + pass + + return test_subscriber + + create_test_redis_subscriber(broker, topic="test-topic") + + async with TestRedisBroker(broker) as tb: + await tb.publish({"foo": "bar"}, "test-topic") + assert check_is_metrics_has_labels(expected_label_keys)