diff --git a/src/servers/benches/to_http_output.rs b/src/servers/benches/to_http_output.rs index f3a9149083a..674c0c912bd 100644 --- a/src/servers/benches/to_http_output.rs +++ b/src/servers/benches/to_http_output.rs @@ -190,6 +190,7 @@ async fn limited_response( ..Default::default() }), Extension(QueryContext::with_db_name(None)), + None, Form(SqlQuery::default()), ) .await; diff --git a/src/servers/src/http/authorize.rs b/src/servers/src/http/authorize.rs index 5c5211ce79f..325efd129db 100644 --- a/src/servers/src/http/authorize.rs +++ b/src/servers/src/http/authorize.rs @@ -30,7 +30,7 @@ use common_time::Timezone; use common_time::timezone::parse_timezone; use headers::Header; use serde::Serialize; -use session::context::QueryContextBuilder; +use session::context::{QueryContext, QueryContextBuilder}; use snafu::{OptionExt, ResultExt, ensure}; use crate::error::{ @@ -113,6 +113,7 @@ pub async fn inner_auth( { Ok(userinfo) => { query_ctx.set_current_user(userinfo); + req.extensions_mut().insert(user_provider); let _ = req.extensions_mut().insert(query_ctx); return Ok(req); } @@ -163,6 +164,7 @@ pub async fn inner_auth( { Ok(userinfo) => { query_ctx.set_current_user(userinfo); + req.extensions_mut().insert(user_provider); let _ = req.extensions_mut().insert(query_ctx); Ok(req) } @@ -180,6 +182,34 @@ pub async fn inner_auth( } } +pub(crate) async fn select_database( + query_ctx: &mut QueryContext, + db: Option<&str>, + user_provider: Option<&UserProviderRef>, +) -> auth::error::Result<()> { + let Some(db) = db else { + return Ok(()); + }; + let (catalog, schema) = parse_catalog_and_schema_from_db_string(db); + if catalog == query_ctx.current_catalog() && schema == query_ctx.current_schema() { + return Ok(()); + } + // Middleware authorized the initial scope; form and query parameters can change it. + if let Some(provider) = user_provider + && let Err(error) = provider + .authorize(&catalog, &schema, &query_ctx.current_user()) + .await + { + crate::metrics::METRIC_AUTH_FAILURE + .with_label_values(&[error.status_code().as_ref()]) + .inc(); + return Err(error); + } + query_ctx.set_current_catalog(&catalog); + query_ctx.set_current_schema(&schema); + Ok(()) +} + pub async fn check_http_auth( State(auth_state): State, req: Request, @@ -454,13 +484,54 @@ fn extract_influxdb_user_from_query(query: &str) -> (Option<&str>, Option<&str>) } #[cfg(test)] -mod tests { +pub(crate) mod tests { use std::assert_matches; use common_base::secrets::ExposeSecret; use super::*; + pub(crate) struct ScopedUserProvider; + + #[async_trait::async_trait] + impl auth::UserProvider for ScopedUserProvider { + fn name(&self) -> &str { + "scoped-test" + } + + async fn authenticate( + &self, + _: auth::Identity<'_>, + _: auth::Password<'_>, + ) -> auth::error::Result { + Ok(auth::userinfo_by_name(Some("reader".to_string()))) + } + + async fn authorize( + &self, + catalog: &str, + schema: &str, + _: &auth::UserInfoRef, + ) -> auth::error::Result<()> { + if catalog == "greptime" && matches!(schema, "public" | "%64enied") { + Ok(()) + } else { + auth::error::PermissionDeniedSnafu.fail() + } + } + + async fn auth_bearer_token( + &self, + _: &str, + catalog: &str, + schema: &str, + ) -> auth::error::Result { + let user = auth::userinfo_by_name(Some("reader".to_string())); + self.authorize(catalog, schema, &user).await?; + Ok(user) + } + } + #[test] fn test_need_auth() { let req = Request::builder() diff --git a/src/servers/src/http/handler.rs b/src/servers/src/http/handler.rs index 8cbfc25ff07..fd85d22402b 100644 --- a/src/servers/src/http/handler.rs +++ b/src/servers/src/http/handler.rs @@ -19,12 +19,12 @@ use std::sync::Arc; use std::sync::atomic::{AtomicU64, Ordering}; use std::time::{Duration, Instant}; +use auth::UserProviderRef; use axum::extract::rejection::FormRejection; use axum::extract::{Json, Query, State}; use axum::response::sse::{Event, KeepAlive, Sse}; use axum::response::{IntoResponse, Response}; use axum::{Extension, Form}; -use common_catalog::parse_catalog_and_schema_from_db_string; use common_error::ext::ErrorExt; use common_error::status_code::StatusCode; use common_plugins::GREPTIME_EXEC_WRITE_COST; @@ -48,6 +48,7 @@ use crate::error::{ CollectRecordbatchSnafu, FailedToParseQuerySnafu, InvalidQuerySnafu, Result, log_error_if_necessary, }; +use crate::http::authorize::select_database; use crate::http::header::collect_plan_metrics; use crate::http::result::arrow_result::ArrowResponse; use crate::http::result::csv_result::CsvResponse; @@ -143,14 +144,19 @@ pub async fn sql( State(state): State, Query(query_params): Query, Extension(mut query_ctx): Extension, + user_provider: Option>, Form(form_params): Form, ) -> HttpResponse { let start = Instant::now(); let sql_handler = &state.sql_handler; - if let Some(db) = &query_params.db.or(form_params.db) { - let (catalog, schema) = parse_catalog_and_schema_from_db_string(db); - query_ctx.set_current_catalog(&catalog); - query_ctx.set_current_schema(&schema); + if let Err(error) = select_database( + &mut query_ctx, + query_params.db.or(form_params.db).as_deref(), + user_provider.as_ref().map(|provider| &provider.0), + ) + .await + { + return HttpResponse::Error(ErrorResponse::from_error(error)); } let db = query_ctx.get_db_string(); @@ -243,6 +249,7 @@ pub struct LogicalTableDdl { pub async fn create_logical_tables( State(state): State, Extension(mut query_ctx): Extension, + user_provider: Option>, form: std::result::Result, FormRejection>, ) -> HttpResponse { let start = Instant::now(); @@ -267,10 +274,14 @@ pub async fn create_logical_tables( "logical-table DDL exceeds 1 MiB".to_string(), )); } - if let Some(db) = form.db { - let (catalog, schema) = parse_catalog_and_schema_from_db_string(&db); - query_ctx.set_current_catalog(&catalog); - query_ctx.set_current_schema(&schema); + if let Err(error) = select_database( + &mut query_ctx, + form.db.as_deref(), + user_provider.as_ref().map(|provider| &provider.0), + ) + .await + { + return HttpResponse::Error(ErrorResponse::from_error(error)); } query_ctx.set_channel(Channel::HttpSql); let query_ctx = Arc::new(query_ctx); @@ -374,6 +385,7 @@ pub async fn sql_analyze_stream( State(state): State, Query(query_params): Query, Extension(mut query_ctx): Extension, + user_provider: Option>, form_params: std::result::Result, FormRejection>, ) -> Response { let start = Instant::now(); @@ -392,10 +404,14 @@ pub async fn sql_analyze_stream( } }; let sql_handler = &state.sql_handler; - if let Some(db) = &query_params.db.or(form_params.db) { - let (catalog, schema) = parse_catalog_and_schema_from_db_string(db); - query_ctx.set_current_catalog(&catalog); - query_ctx.set_current_schema(&schema); + if let Err(error) = select_database( + &mut query_ctx, + query_params.db.or(form_params.db).as_deref(), + user_provider.as_ref().map(|provider| &provider.0), + ) + .await + { + return ErrorResponse::from_error(error).into_response(); } query_ctx.set_channel(Channel::HttpSql); query_ctx.enable_live_analyze_metrics(); @@ -1086,6 +1102,7 @@ mod tests { use snafu::IntoError; use super::*; + use crate::http::authorize::tests::ScopedUserProvider; use crate::query_handler::sql::SqlQueryHandler; struct TestSqlQueryHandler { @@ -1134,6 +1151,84 @@ mod tests { } } + #[tokio::test] + async fn test_sql_endpoints_authorize_effective_database() { + use axum::http::{Request, StatusCode as HttpStatusCode}; + use tower::ServiceExt; + + for authorization in ["Basic cmVhZGVyOnB3ZA==", "Bearer test-token"] { + for (path, query, body, method) in [ + ("/v1/sql", "", "sql=select+1&db=denied", "POST"), + ("/v1/sql", "?db=denied&sql=select+1", "", "GET"), + ("/v1/sql", "?db=%64enied&sql=select+1", "", "GET"), + ("/v1/sql", "?db=denied", "sql=select+1&db=public", "POST"), + ("/v1/ddl", "", "sql=create+table+t&db=denied", "POST"), + ("/v1/analyze", "", "sql=select+1&db=denied", "POST"), + ("/v1/analyze", "?db=denied&sql=select+1", "", "GET"), + ] { + let handler = Arc::new(TestSqlQueryHandler { + outputs: Mutex::new(vec![Ok(Output::new_with_affected_rows(0))]), + }); + let app = axum::Router::new() + .route("/v1/sql", axum::routing::get(sql).post(sql)) + .route("/v1/ddl", axum::routing::post(create_logical_tables)) + .route( + "/v1/analyze", + axum::routing::get(sql_analyze_stream).post(sql_analyze_stream), + ) + .with_state(ApiState { + sql_handler: handler.clone(), + }) + .layer(axum::middleware::from_fn_with_state( + crate::http::authorize::AuthState::new(Some(Arc::new(ScopedUserProvider))), + crate::http::authorize::check_http_auth, + )); + let mut request = Request::builder() + .method(method) + .uri(format!("{path}{query}")) + .header("authorization", authorization) + .header("content-type", "application/x-www-form-urlencoded"); + if !query.contains("%64") { + request = request.header("x-greptime-db-name", "public"); + } + let response = app + .oneshot(request.body(axum::body::Body::from(body)).unwrap()) + .await + .unwrap(); + assert_eq!( + HttpStatusCode::FORBIDDEN, + response.status(), + "{path}{query}" + ); + assert_eq!(1, handler.outputs.lock().unwrap().len()); + } + } + } + + #[tokio::test] + async fn test_select_database_authorizes_before_mutating_context() { + let provider: UserProviderRef = Arc::new(ScopedUserProvider); + let mut ctx = QueryContext::with_db_name(Some("initial")); + select_database(&mut ctx, Some("public"), Some(&provider)) + .await + .unwrap(); + assert_eq!("public", ctx.current_schema()); + let failures = crate::metrics::METRIC_AUTH_FAILURE + .with_label_values(&[StatusCode::PermissionDenied.as_ref()]); + let before = failures.get(); + assert!( + select_database(&mut ctx, Some("denied"), Some(&provider)) + .await + .is_err() + ); + assert!(failures.get() > before); + assert_eq!("public", ctx.current_schema()); + select_database(&mut ctx, Some("other"), None) + .await + .unwrap(); + assert_eq!("other", ctx.current_schema()); + } + async fn sql_response( outputs: Vec>, format: &str, @@ -1152,6 +1247,7 @@ mod tests { ..Default::default() }), Extension(QueryContext::with_db_name(None)), + None, Form(SqlQuery::default()), ) .await @@ -1181,6 +1277,7 @@ mod tests { }), }), Extension(QueryContext::with_db_name(None)), + None, Ok(Form(LogicalTableDdl { sql: Some(sql.to_string()), db: None, diff --git a/src/servers/src/http/prometheus.rs b/src/servers/src/http/prometheus.rs index 0773b337889..6ab9b053f5f 100644 --- a/src/servers/src/http/prometheus.rs +++ b/src/servers/src/http/prometheus.rs @@ -26,11 +26,10 @@ use arrow::datatypes::{ UInt8Type, UInt16Type, UInt32Type, UInt64Type, }; use arrow_schema::{DataType, IntervalUnit}; -use auth::{PermissionTableTarget, PermissionTableTargets}; +use auth::{PermissionTableTarget, PermissionTableTargets, UserProviderRef}; use axum::extract::{Path, Query, State}; use axum::{Extension, Form}; use catalog::CatalogManagerRef; -use common_catalog::parse_catalog_and_schema_from_db_string; use common_decimal::Decimal128; use common_error::ext::ErrorExt; use common_error::status_code::StatusCode; @@ -76,6 +75,7 @@ use crate::error::{ CollectRecordbatchSnafu, ConvertScalarValueSnafu, DataFusionSnafu, Error, InvalidQuerySnafu, NotSupportedSnafu, ParseTimestampSnafu, Result, TableNotFoundSnafu, UnexpectedResultSnafu, }; +use crate::http::authorize::select_database; use crate::http::header::collect_plan_metrics; use crate::otlp::metrics::ucum_to_openmetrics_unit; use crate::prom_store::{FIELD_NAME_LABEL, METRIC_NAME_LABEL, is_database_selection_label}; @@ -237,14 +237,20 @@ pub struct ExemplarQuery { fields(protocol = "prometheus", request_type = "query_exemplars") )] pub async fn query_exemplars( + user_provider: Option>, State(handler): State, Query(params): Query, Extension(mut ctx): Extension, Form(form): Form, ) -> PrometheusJsonResponse { - if let Some(db) = params.db.or(form.db) { - let (catalog, schema) = parse_catalog_and_schema_from_db_string(&db); - try_update_catalog_schema(&mut ctx, &catalog, &schema); + if let Err(error) = select_database( + &mut ctx, + params.db.or(form.db).as_deref(), + user_provider.as_ref().map(|Extension(provider)| provider), + ) + .await + { + return PrometheusJsonResponse::error(error.status_code(), error.output_msg()); } ctx.set_channel(Channel::Prometheus); let ctx = Arc::new(ctx); @@ -356,12 +362,20 @@ pub async fn build_info_query() -> PrometheusJsonResponse { fields(protocol = "prometheus", request_type = "metadata_query") )] pub async fn metadata_query( + user_provider: Option>, State(handler): State, Query(params): Query, Extension(mut query_ctx): Extension, ) -> PrometheusJsonResponse { - let (catalog, schema) = get_catalog_schema(¶ms.db, &query_ctx); - try_update_catalog_schema(&mut query_ctx, &catalog, &schema); + if let Err(error) = select_database( + &mut query_ctx, + params.db.as_deref(), + user_provider.as_ref().map(|Extension(provider)| provider), + ) + .await + { + return PrometheusJsonResponse::error(error.status_code(), error.output_msg()); + } let query_ctx = Arc::new(query_ctx); if let Err(err) = handler.check_query_permission(&[], &query_ctx).await { @@ -459,6 +473,7 @@ macro_rules! try_call_return_response { fields(protocol = "prometheus", request_type = "instant_query") )] pub async fn instant_query( + user_provider: Option>, State(handler): State, Query(params): Query, Extension(mut query_ctx): Extension, @@ -482,9 +497,14 @@ pub async fn instant_query( }; // update catalog and schema in query context if necessary - if let Some(db) = ¶ms.db { - let (catalog, schema) = parse_catalog_and_schema_from_db_string(db); - try_update_catalog_schema(&mut query_ctx, &catalog, &schema); + if let Err(error) = select_database( + &mut query_ctx, + params.db.as_deref(), + user_provider.as_ref().map(|Extension(provider)| provider), + ) + .await + { + return PrometheusJsonResponse::error(error.status_code(), error.output_msg()); } let query_ctx = Arc::new(query_ctx); let _timer = crate::metrics::METRIC_HTTP_PROMETHEUS_PROMQL_ELAPSED @@ -595,6 +615,7 @@ pub struct RangeQuery { fields(protocol = "prometheus", request_type = "range_query") )] pub async fn range_query( + user_provider: Option>, State(handler): State, Query(params): Query, Extension(mut query_ctx): Extension, @@ -613,9 +634,14 @@ pub async fn range_query( }; // update catalog and schema in query context if necessary - if let Some(db) = ¶ms.db { - let (catalog, schema) = parse_catalog_and_schema_from_db_string(db); - try_update_catalog_schema(&mut query_ctx, &catalog, &schema); + if let Err(error) = select_database( + &mut query_ctx, + params.db.as_deref(), + user_provider.as_ref().map(|Extension(provider)| provider), + ) + .await + { + return PrometheusJsonResponse::error(error.status_code(), error.output_msg()); } let query_ctx = Arc::new(query_ctx); let _timer = crate::metrics::METRIC_HTTP_PROMETHEUS_PROMQL_ELAPSED @@ -798,13 +824,21 @@ macro_rules! handle_schema_err { fields(protocol = "prometheus", request_type = "labels_query") )] pub async fn labels_query( + user_provider: Option>, State(handler): State, Query(params): Query, Extension(mut query_ctx): Extension, Form(form_params): Form, ) -> PrometheusJsonResponse { - let (catalog, schema) = get_catalog_schema(¶ms.db, &query_ctx); - try_update_catalog_schema(&mut query_ctx, &catalog, &schema); + if let Err(error) = select_database( + &mut query_ctx, + params.db.as_deref(), + user_provider.as_ref().map(|Extension(provider)| provider), + ) + .await + { + return PrometheusJsonResponse::error(error.status_code(), error.output_msg()); + } let query_ctx = Arc::new(query_ctx); let mut queries = params.matches.0; @@ -857,10 +891,13 @@ pub async fn labels_query( prom_queries }; + let catalog = query_ctx.current_catalog(); + let schema = query_ctx.current_schema(); + // Fetch all columns if no query matcher is provided if prom_queries.is_empty() { let (mut labels, table_names) = - match get_all_column_names(&catalog, &schema, &handler.catalog_manager()).await { + match get_all_column_names(catalog, &schema, &handler.catalog_manager()).await { Ok(result) => result, Err(e) => return PrometheusJsonResponse::error(e.status_code(), e.output_msg()), }; @@ -886,7 +923,7 @@ pub async fn labels_query( Err(e) => return PrometheusJsonResponse::error(e.status_code(), e.output_msg()), } } - None => match get_all_column_names(&catalog, &schema, &handler.catalog_manager()).await { + None => match get_all_column_names(catalog, &schema, &handler.catalog_manager()).await { Ok((labels, _)) => labels, Err(e) => return PrometheusJsonResponse::error(e.status_code(), e.output_msg()), }, @@ -1403,27 +1440,6 @@ pub(crate) fn retrieve_metric_name_and_result_type( ) } -/// Tries to get catalog and schema from an optional db param. And retrieves -/// them from [QueryContext] if they don't present. -pub(crate) fn get_catalog_schema(db: &Option, ctx: &QueryContext) -> (String, String) { - if let Some(db) = db { - parse_catalog_and_schema_from_db_string(db) - } else { - ( - ctx.current_catalog().to_string(), - ctx.current_schema().clone(), - ) - } -} - -/// Update catalog and schema in [QueryContext] if necessary. -pub(crate) fn try_update_catalog_schema(ctx: &mut QueryContext, catalog: &str, schema: &str) { - if ctx.current_catalog() != catalog || ctx.current_schema() != schema { - ctx.set_current_catalog(catalog); - ctx.set_current_schema(schema); - } -} - fn current_schema_metric_targets( query_ctx: &QueryContext, metric_names: &[String], @@ -1834,13 +1850,21 @@ pub struct LabelValueQuery { fields(protocol = "prometheus", request_type = "label_values_query") )] pub async fn label_values_query( + user_provider: Option>, State(handler): State, Path(label_name): Path, Extension(mut query_ctx): Extension, Query(params): Query, ) -> PrometheusJsonResponse { - let (catalog, schema) = get_catalog_schema(¶ms.db, &query_ctx); - try_update_catalog_schema(&mut query_ctx, &catalog, &schema); + if let Err(error) = select_database( + &mut query_ctx, + params.db.as_deref(), + user_provider.as_ref().map(|Extension(provider)| provider), + ) + .await + { + return PrometheusJsonResponse::error(error.status_code(), error.output_msg()); + } let query_ctx = Arc::new(query_ctx); let _timer = crate::metrics::METRIC_HTTP_PROMETHEUS_PROMQL_ELAPSED @@ -2495,6 +2519,7 @@ pub struct SeriesQuery { fields(protocol = "prometheus", request_type = "series_query") )] pub async fn series_query( + user_provider: Option>, State(handler): State, Query(params): Query, Extension(mut query_ctx): Extension, @@ -2524,9 +2549,14 @@ pub async fn series_query( .unwrap_or_else(|| DEFAULT_LOOKBACK_STRING.to_string()); // update catalog and schema in query context if necessary - if let Some(db) = ¶ms.db { - let (catalog, schema) = parse_catalog_and_schema_from_db_string(db); - try_update_catalog_schema(&mut query_ctx, &catalog, &schema); + if let Err(error) = select_database( + &mut query_ctx, + params.db.as_deref(), + user_provider.as_ref().map(|Extension(provider)| provider), + ) + .await + { + return PrometheusJsonResponse::error(error.status_code(), error.output_msg()); } let query_ctx = Arc::new(query_ctx); @@ -2875,6 +2905,105 @@ mod tests { ); } + #[tokio::test] + async fn test_prometheus_endpoints_authorize_effective_database() { + use axum::http::{Request, StatusCode as HttpStatusCode}; + use axum::response::IntoResponse; + use tower::ServiceExt; + + for authorization in ["Basic cmVhZGVyOnB3ZA==", "Bearer test-token"] { + for (path, form_db, encoded_db) in [ + ("/v1/prometheus/api/v1/query", false, false), + ("/v1/prometheus/api/v1/query_range", false, false), + ("/v1/prometheus/api/v1/metadata", false, false), + ("/v1/prometheus/api/v1/labels", false, false), + ("/v1/prometheus/api/v1/label/job/values", false, false), + ("/v1/prometheus/api/v1/series", false, false), + ("/v1/prometheus/api/v1/query_exemplars", false, false), + ("/v1/prometheus/api/v1/query_exemplars", true, false), + ("/v1/prometheus/api/v1/query", false, true), + ] { + let handler = Arc::new(TestPrometheusHandler { + catalog_manager: MemoryCatalogManager::new(), + deny_operation: false, + denied_table: None, + metric_names: Vec::new(), + label_metric_names: Vec::new(), + label_lookups: Mutex::new(Vec::new()), + queries: Mutex::new(Vec::new()), + ordered_outputs: Mutex::new(Vec::new()), + }); + let app = axum::Router::new() + .route( + "/v1/prometheus/api/v1/query", + axum::routing::get(instant_query), + ) + .route( + "/v1/prometheus/api/v1/query_range", + axum::routing::get(range_query), + ) + .route( + "/v1/prometheus/api/v1/metadata", + axum::routing::get(metadata_query), + ) + .route( + "/v1/prometheus/api/v1/labels", + axum::routing::get(labels_query), + ) + .route( + "/v1/prometheus/api/v1/label/{name}/values", + axum::routing::get(label_values_query), + ) + .route( + "/v1/prometheus/api/v1/series", + axum::routing::get(series_query), + ) + .route( + "/v1/prometheus/api/v1/query_exemplars", + axum::routing::get(query_exemplars).post(query_exemplars), + ) + .with_state(handler.clone() as PrometheusHandlerRef) + .layer(axum::middleware::from_fn_with_state( + crate::http::authorize::AuthState::new(Some(Arc::new( + crate::http::authorize::tests::ScopedUserProvider, + ))), + crate::http::authorize::check_http_auth, + )); + let db = if encoded_db { "%64enied" } else { "denied" }; + let query = if form_db { + String::new() + } else { + format!("&db={db}") + }; + let mut request = Request::builder() + .method(if form_db { "POST" } else { "GET" }) + .uri(format!( + "{path}?query=requests&start=0&end=1&step=1&match[]=requests{query}" + )) + .header("authorization", authorization) + .header("content-type", "application/x-www-form-urlencoded"); + if !encoded_db { + request = request.header("x-greptime-db-name", "public"); + } + let body = if form_db { "db=denied" } else { "" }; + let response = tokio::time::timeout( + std::time::Duration::from_secs(5), + app.oneshot(request.body(axum::body::Body::from(body)).unwrap()), + ) + .await + .unwrap() + .unwrap(); + assert_eq!( + HttpStatusCode::FORBIDDEN, + response.into_response().status(), + "{path}" + ); + assert!(handler.queries.lock().unwrap().is_empty()); + assert!(handler.label_lookups.lock().unwrap().is_empty()); + } + } + } + #[tokio::test] async fn test_exemplar_timeout_cancels_scan() { for explicit in [false, true] { @@ -2893,6 +3022,7 @@ mod tests { ctx.set_query_timeout(std::time::Duration::from_millis(1)); } let response = query_exemplars( + None, State(handler.clone()), Query(ExemplarQuery { query: Some("requests".into()), @@ -2924,6 +3054,7 @@ mod tests { }); let state: PrometheusHandlerRef = handler.clone(); instant_query( + None, State(state.clone()), Query(InstantQuery { query: Some("sort(vector(1))".to_string()), @@ -2941,6 +3072,7 @@ mod tests { // Both a single-point and a multi-step range query take the same path. for end in ["0", "1"] { range_query( + None, State(state.clone()), Query(RangeQuery { query: Some("sort(vector(1))".to_string()), @@ -2978,6 +3110,7 @@ mod tests { ordered_outputs: Mutex::new(Vec::new()), }); let response = range_query( + None, State(handler.clone() as PrometheusHandlerRef), Query(RangeQuery { query: Some("vector(1)[5m:]".to_string()), @@ -3023,6 +3156,7 @@ mod tests { .with_label_values(&[db.as_str(), "instant_query"]); let instant_count = instant_histogram.get_sample_count(); let response = instant_query( + None, State(handler.clone()), Query(InstantQuery { query: Some("up{".to_string()), @@ -3039,6 +3173,7 @@ mod tests { .with_label_values(&[db.as_str(), "range_query"]); let range_count = range_histogram.get_sample_count(); let response = range_query( + None, State(handler), Query(RangeQuery { query: Some("up{".to_string()), @@ -3166,6 +3301,7 @@ mod tests { assert!(schemas.contains(&"empty".to_string())); assert!(!schemas.contains(&"denied_schema".to_string())); let response = label_values_query( + None, State(handler), Path(FIELD_NAME_LABEL.to_string()), Extension(QueryContext::with( @@ -3225,6 +3361,7 @@ mod tests { ) -> Vec { let state: PrometheusHandlerRef = handler; let response = label_values_query( + None, State(state), Path(METRIC_NAME_LABEL.to_string()), Extension(QueryContext::with( @@ -3362,6 +3499,7 @@ mod tests { let query_ctx = QueryContext::with(DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME); let target_ctx = Arc::new(query_ctx.clone()); let response = series_query( + None, State(state), Query(SeriesQuery { matches: Matches(vec![r#"{__name__=~"cpu_.*"}"#.to_string()]), @@ -4346,6 +4484,7 @@ mod tests { assert!(metadata.contains_key("allowed")); let response = metadata_query( + None, State(Arc::new(TestPrometheusHandler { catalog_manager: manager.clone(), deny_operation: true, @@ -4373,6 +4512,7 @@ mod tests { ordered_outputs: Mutex::new(Vec::new()), }); let response = metadata_query( + None, State(handler.clone()), Query(MetadataQuery::default()), Extension(query_ctx.clone()), @@ -4388,6 +4528,7 @@ mod tests { ); let response = metadata_query( + None, State(handler), Query(MetadataQuery { metric: Some("denied".to_string()), diff --git a/src/servers/tests/http/http_handler_test.rs b/src/servers/tests/http/http_handler_test.rs index 2c8064727a6..7492ec50ab8 100644 --- a/src/servers/tests/http/http_handler_test.rs +++ b/src/servers/tests/http/http_handler_test.rs @@ -266,6 +266,7 @@ async fn test_sql_not_provided() { State(api_state.clone()), Query(query), axum::Extension(ctx.clone()), + None, Form(http_handler::SqlQuery::default()), ) .await @@ -294,6 +295,7 @@ async fn test_sql_output_rows() { State(api_state.clone()), query, axum::Extension(ctx.clone()), + None, Form(http_handler::SqlQuery::default()), ) .await; @@ -397,6 +399,7 @@ async fn test_dashboard_sql_limit() { State(api_state.clone()), query, axum::Extension(ctx.clone()), + None, Form(http_handler::SqlQuery::default()), ) .await; @@ -441,6 +444,7 @@ async fn test_sql_form() { State(api_state.clone()), Query(http_handler::SqlQuery::default()), axum::Extension(ctx.clone()), + None, form, ) .await; @@ -750,6 +754,7 @@ async fn test_analyze_stream_worker_panic_emits_one_error_terminal() { ..Default::default() }), axum::Extension(ctx), + None, Ok(Form(http_handler::SqlQuery::default())), ) .await; @@ -791,6 +796,7 @@ async fn test_analyze_stream_worker_panic_after_metrics_uses_current_sequence() ..Default::default() }), axum::Extension(ctx), + None, Ok(Form(http_handler::SqlQuery::default())), ) .await;