Add support OpenMetrics format in default endpoints

Author: draincoder

src/asgi_monitor/integrations/fastapi.py

@@ -37,6 +37,9 @@ class MetricsConfig(BaseMetricsConfig):
     include_metrics_endpoint: bool = field(default=True)
     """Whether to include a /metrics endpoint."""
 
+    openmetrics_format: bool = field(default=False)
+    """A flag indicating whether to generate metrics in OpenMetrics format."""
+
 
 @dataclass(slots=True, frozen=True)
 class TracingConfig(BaseTracingConfig):
@@ -92,6 +95,7 @@ def setup_metrics(app: FastAPI, config: MetricsConfig) -> None:
     )
     if config.include_metrics_endpoint:
         app.state.metrics_registry = config.registry
+        app.state.openmetrics_format = config.openmetrics_format
         app.add_route(
             path="/metrics",
             route=get_metrics,

src/asgi_monitor/integrations/litestar.py

@@ -154,7 +154,8 @@ async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
 @get(path="/metrics", summary="Get Prometheus metrics", include_in_schema=True)
 async def get_metrics(request: Request) -> Response:
     registry = request.app.state.metrics_registry
-    response = get_latest_metrics(registry, openmetrics_format=False)
+    openmetrics_format = request.app.state.openmetrics_format
+    response = get_latest_metrics(registry, openmetrics_format=openmetrics_format)
     return Response(
         content=response.payload,
         status_code=response.status_code,
@@ -195,14 +196,16 @@ def build_metrics_middleware(config: MetricsConfig) -> DefineMiddleware:
     )
 
 
-def add_metrics_endpoint(app: Litestar, registry: CollectorRegistry) -> None:
+def add_metrics_endpoint(app: Litestar, registry: CollectorRegistry, *, openmetrics_format: bool = False) -> None:
     """
     Add CollectorRegistry in state and register /metrics endpoint.
 
     :param Litestar app: The Litestar application instance.
     :param CollectorRegistry registry: The registry for the metrics.
+    :param bool openmetrics_format: A flag indicating whether to generate metrics in OpenMetrics format.
     :returns: None
     """
 
     app.state.metrics_registry = registry
+    app.state.openmetrics_format = openmetrics_format
     app.register(get_metrics)

src/asgi_monitor/integrations/starlette.py

@@ -98,6 +98,9 @@ class MetricsConfig(BaseMetricsConfig):
     include_metrics_endpoint: bool = field(default=True)
     """Whether to include a /metrics endpoint."""
 
+    openmetrics_format: bool = field(default=False)
+    """A flag indicating whether to generate metrics in OpenMetrics format."""
+
 
 class TracingMiddleware:
     __slots__ = ("app", "open_telemetry_middleware")
@@ -179,7 +182,8 @@ async def dispatch(self, request: Request, call_next: RequestResponseEndpoint) -
 
 async def get_metrics(request: Request) -> Response:
     registry = request.app.state.metrics_registry
-    response = get_latest_metrics(registry, openmetrics_format=False)
+    openmetrics_format = request.app.state.openmetrics_format
+    response = get_latest_metrics(registry, openmetrics_format=openmetrics_format)
     return Response(
         content=response.payload,
         status_code=response.status_code,
@@ -220,6 +224,7 @@ def setup_metrics(app: Starlette, config: MetricsConfig) -> None:
     )
     if config.include_metrics_endpoint:
         app.state.metrics_registry = config.registry
+        app.state.openmetrics_format = config.openmetrics_format
         app.add_route(
             path="/metrics",
             route=get_metrics,

tests/integration/fastapi/test_middlewares.py

@@ -19,8 +19,14 @@ async def index() -> dict:
 
 async def test_metrics() -> None:
     # Arrange
+    expected_content_type = "text/plain; version=0.0.4; charset=utf-8"
     app = FastAPI()
-    metrics_config = MetricsConfig(app_name="test", include_metrics_endpoint=True, include_trace_exemplar=False)
+    metrics_config = MetricsConfig(
+        app_name="test",
+        include_metrics_endpoint=True,
+        include_trace_exemplar=False,
+        openmetrics_format=False,
+    )
     setup_metrics(app=app, config=metrics_config)
 
     # Act
@@ -29,6 +35,7 @@ async def test_metrics() -> None:
 
         # Assert
         assert response.status_code == 200
+        assert response.headers["content-type"] == expected_content_type
         assert_that(response.content.decode()).contains(
             'fastapi_app_info{app_name="test"} 1.0',
             'fastapi_requests_total{app_name="test",method="GET",path="/metrics"} 1.0',
@@ -87,6 +94,39 @@ async def test_metrics_with_tracing() -> None:
         assert_that(metrics.payload.decode()).contains('fastapi_app_info{app_name="test"} 1.0')
 
 
+async def test_metrics_openmetrics_with_tracing() -> None:
+    # Arrange
+    expected_content_type = "application/openmetrics-text; version=1.0.0; charset=utf-8"
+    app = FastAPI()
+    app.include_router(router)
+    trace_config, _ = build_fastapi_tracing_config()
+    metrics_config = MetricsConfig(
+        app_name="test",
+        include_metrics_endpoint=True,
+        include_trace_exemplar=True,
+        openmetrics_format=True,
+    )
+
+    setup_metrics(app=app, config=metrics_config)
+    setup_tracing(app=app, config=trace_config)
+
+    # Act
+    async with fastapi_app(app) as client:
+        response = client.get("/")
+        metrics = client.get("/metrics")
+
+        # Assert
+        assert response.status_code == 200
+        assert metrics.status_code == 200
+        assert metrics.headers["content-type"] == expected_content_type
+        pattern = (
+            r"fastapi_request_duration_seconds_bucket\{"
+            r'app_name="test",le="([\d.]+)",method="GET",path="\/"}\ 1.0 # \{TraceID="(\w+)"\} (\d+\.\d+) (\d+\.\d+)'
+        )
+        assert_that(metrics.content.decode()).matches(pattern)
+        assert_that(metrics.content.decode()).contains('fastapi_app_info{app_name="test"} 1.0')
+
+
 async def test_metrics_get_path() -> None:
     # Arrange
     app = FastAPI()

tests/integration/litestar/test_middlewares.py

@@ -67,16 +67,18 @@ async def test_tracing() -> None:
 
 async def test_metrics() -> None:
     # Arrange
+    expected_content_type = "text/plain; version=0.0.4; charset=utf-8"
     metrics_config = MetricsConfig(app_name="test", include_trace_exemplar=False)
     app = Litestar(middleware=[build_metrics_middleware(metrics_config)])
-    add_metrics_endpoint(app, metrics_config.registry)
+    add_metrics_endpoint(app, metrics_config.registry, openmetrics_format=False)
 
     # Act
     async with litestar_app(app) as client:
         response = client.get("/metrics")
 
         # Assert
         assert response.status_code == 200
+        assert response.headers["content-type"] == expected_content_type
         assert_that(response.content.decode()).contains(
             'litestar_app_info{app_name="test"} 1.0',
             'litestar_requests_total{app_name="test",method="GET",path="/metrics"} 1.0',
@@ -129,3 +131,28 @@ async def test_metrics_with_tracing() -> None:
             r'app_name="test",le="([\d.]+)",method="GET",path="\/"}\ 1.0 # \{TraceID="(\w+)"\} (\d+\.\d+) (\d+\.\d+)'
         )
         assert_that(metrics.payload.decode()).matches(pattern)
+
+
+async def test_metrics_openmetrics_with_tracing() -> None:
+    # Arrange
+    expected_content_type = "application/openmetrics-text; version=1.0.0; charset=utf-8"
+    trace_config, _ = build_litestar_tracing_config()
+    metrics_config = MetricsConfig(app_name="test", include_trace_exemplar=True)
+    middlewares = [build_tracing_middleware(trace_config), build_metrics_middleware(metrics_config)]
+    app = Litestar([index], middleware=middlewares)
+    add_metrics_endpoint(app, metrics_config.registry, openmetrics_format=True)
+
+    # Act
+    async with litestar_app(app) as client:
+        response = client.get("/")
+        metrics = client.get("/metrics")
+
+        # Assert
+        assert response.status_code == 200
+        assert metrics.status_code == 200
+        assert metrics.headers["content-type"] == expected_content_type
+        pattern = (
+            r"litestar_request_duration_seconds_bucket\{"
+            r'app_name="test",le="([\d.]+)",method="GET",path="\/"}\ 1.0 # \{TraceID="(\w+)"\} (\d+\.\d+) (\d+\.\d+)'
+        )
+        assert_that(metrics.content.decode()).matches(pattern)

tests/integration/starlette/test_middlewares.py

@@ -135,8 +135,14 @@ async def test_tracing_partial_match() -> None:
 
 async def test_metrics() -> None:
     # Arrange
+    expected_content_type = "text/plain; version=0.0.4; charset=utf-8"
     app = Starlette()
-    metrics_config = MetricsConfig(app_name="test", include_metrics_endpoint=True, include_trace_exemplar=False)
+    metrics_config = MetricsConfig(
+        app_name="test",
+        include_metrics_endpoint=True,
+        include_trace_exemplar=False,
+        openmetrics_format=False,
+    )
     setup_metrics(app=app, config=metrics_config)
 
     # Act
@@ -145,6 +151,7 @@ async def test_metrics() -> None:
 
         # Assert
         assert response.status_code == 200
+        assert response.headers["content-type"] == expected_content_type
         assert_that(response.content.decode()).contains(
             'starlette_app_info{app_name="test"} 1.0',
             'starlette_requests_total{app_name="test",method="GET",path="/metrics"} 1.0',
@@ -201,3 +208,34 @@ async def test_metrics_with_tracing() -> None:
             r'app_name="test",le="([\d.]+)",method="GET",path="\/"}\ 1.0 # \{TraceID="(\w+)"\} (\d+\.\d+) (\d+\.\d+)'
         )
         assert_that(metrics.payload.decode()).matches(pattern)
+
+
+async def test_metrics_openmetrics_with_tracing() -> None:
+    # Arrange
+    expected_content_type = "application/openmetrics-text; version=1.0.0; charset=utf-8"
+    trace_config, _ = build_starlette_tracing_config()
+    metrics_config = MetricsConfig(
+        app_name="test",
+        include_metrics_endpoint=True,
+        include_trace_exemplar=True,
+        openmetrics_format=True,
+    )
+    app = Starlette(routes=[Route("/", endpoint=index, methods=["GET"])])
+
+    setup_metrics(app=app, config=metrics_config)
+    setup_tracing(app=app, config=trace_config)
+
+    # Act
+    async with starlette_app(app) as client:
+        response = client.get("/")
+        metrics = client.get("/metrics")
+
+        # Assert
+        assert response.status_code == 200
+        assert metrics.status_code == 200
+        assert metrics.headers["content-type"] == expected_content_type
+        pattern = (
+            r"starlette_request_duration_seconds_bucket\{"
+            r'app_name="test",le="([\d.]+)",method="GET",path="\/"}\ 1.0 # \{TraceID="(\w+)"\} (\d+\.\d+) (\d+\.\d+)'
+        )
+        assert_that(metrics.content.decode()).matches(pattern)