diff --git a/crates/sdk/src/default_main.rs b/crates/sdk/src/default_main.rs index 7d10adc..048aac5 100644 --- a/crates/sdk/src/default_main.rs +++ b/crates/sdk/src/default_main.rs @@ -1,7 +1,12 @@ -use std::error::Error; -use std::net; -use std::path::{Path, PathBuf}; - +use crate::check_health; +use crate::connector::{ + Connector, ConnectorSetup, ExplainError, FetchMetricsError, HealthError, MutationError, + QueryError, SchemaError, +}; +use crate::fetch_metrics::fetch_metrics; +use crate::json_rejection::JsonRejection; +use crate::json_response::JsonResponse; +use crate::tracing::{init_tracing, make_span, on_response}; use axum::{ body::Body, extract::State, @@ -12,22 +17,15 @@ use axum::{ }; use axum_extra::extract::WithRejection; use clap::{Parser, Subcommand}; -use prometheus::{Registry, TextEncoder}; -use tower_http::{trace::TraceLayer, validate_request::ValidateRequestHeaderLayer}; - use ndc_models::{ CapabilitiesResponse, ErrorResponse, ExplainResponse, MutationRequest, MutationResponse, QueryRequest, QueryResponse, SchemaResponse, }; - -use crate::check_health; -use crate::connector::{ - Connector, ConnectorSetup, ExplainError, FetchMetricsError, HealthError, MutationError, - QueryError, SchemaError, -}; -use crate::json_rejection::JsonRejection; -use crate::json_response::JsonResponse; -use crate::tracing::{init_tracing, make_span, on_response}; +use prometheus::Registry; +use std::error::Error; +use std::net; +use std::path::{Path, PathBuf}; +use tower_http::{trace::TraceLayer, validate_request::ValidateRequestHeaderLayer}; #[derive(Parser)] struct CliArgs { @@ -376,15 +374,7 @@ fn auth_handler( async fn get_metrics( State(state): State>, ) -> Result { - let encoder = TextEncoder::new(); - - C::fetch_metrics(&state.configuration, &state.state)?; - - let metric_families = &state.metrics.gather(); - - encoder - .encode_to_string(metric_families) - .map_err(|_| FetchMetricsError::new("Unable to encode metrics")) + fetch_metrics::(&state.configuration, &state.state, &state.metrics) } async fn get_capabilities() -> JsonResponse { diff --git a/crates/sdk/src/fetch_metrics.rs b/crates/sdk/src/fetch_metrics.rs new file mode 100644 index 0000000..654f371 --- /dev/null +++ b/crates/sdk/src/fetch_metrics.rs @@ -0,0 +1,18 @@ +use crate::connector::{Connector, FetchMetricsError}; +use prometheus::{Registry, TextEncoder}; + +pub fn fetch_metrics( + configuration: &C::Configuration, + state: &C::State, + metrics: &Registry, +) -> Result { + let encoder = TextEncoder::new(); + + C::fetch_metrics(configuration, state)?; + + let metric_families = &metrics.gather(); + + encoder + .encode_to_string(metric_families) + .map_err(|_| FetchMetricsError::new("Unable to encode metrics")) +} diff --git a/crates/sdk/src/lib.rs b/crates/sdk/src/lib.rs index ee29aec..a4741cd 100644 --- a/crates/sdk/src/lib.rs +++ b/crates/sdk/src/lib.rs @@ -1,6 +1,7 @@ pub mod check_health; pub mod connector; pub mod default_main; +pub mod fetch_metrics; pub mod json_rejection; pub mod json_response; pub mod tracing;