use crate::collectors::registry::CollectorRegistry;
use axum::{
extract::Extension,
http::{HeaderMap, HeaderValue, StatusCode},
response::{IntoResponse, Response},
};
use prometheus::Encoder;
use sqlx::PgPool;
use tracing::{debug, error, instrument};
#[instrument(skip(pool, registry), fields(http.route="/metrics"))]
pub async fn metrics(
Extension(pool): Extension<PgPool>,
Extension(registry): Extension<CollectorRegistry>,
) -> Response {
let mut headers = HeaderMap::new();
headers.insert(
"content-type",
HeaderValue::from_static("text/plain; charset=utf-8"),
);
match registry.collect_all_bytes(&pool).await {
Ok(metrics) => {
debug!("Successfully collected metrics");
(StatusCode::OK, headers, metrics).into_response()
}
Err(e) => {
error!("Failed to collect metrics: {}", e);
let mut buffer = Vec::new();
let encoder = prometheus::TextEncoder::new();
let metric_families = registry.registry().gather();
if let Err(encode_err) = encoder.encode(&metric_families, &mut buffer) {
error!("Failed to encode metrics on error path: {}", encode_err);
return (
StatusCode::OK,
headers,
format!(
"# Error collecting metrics: {e}\n# Error encoding metrics: {encode_err}"
),
)
.into_response();
}
(StatusCode::OK, headers, buffer).into_response()
}
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
#[allow(clippy::unwrap_used)]
fn test_metrics_response_headers() {
let mut headers = HeaderMap::new();
headers.insert(
"content-type",
HeaderValue::from_static("text/plain; charset=utf-8"),
);
assert_eq!(
headers.get("content-type").unwrap(),
"text/plain; charset=utf-8"
);
}
#[test]
#[allow(clippy::unwrap_used)]
fn test_header_value_creation() {
let header_val = HeaderValue::from_static("text/plain; charset=utf-8");
assert_eq!(header_val.to_str().unwrap(), "text/plain; charset=utf-8");
}
}