Merge pull request #126 from m1khalchenko/prometheus-custom-labels

Add Prometheus custom labels support for FastAPI and FastStream

Author: insani7y

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(

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
 

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]):

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 = [

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

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)