diff --git a/src/servers/src/http/handler.rs b/src/servers/src/http/handler.rs index 0598fe03935..3b3544bc7de 100644 --- a/src/servers/src/http/handler.rs +++ b/src/servers/src/http/handler.rs @@ -17,7 +17,7 @@ use std::time::Instant; use aide::transform::TransformOperation; use axum::extract::{Json, Query, State}; -use axum::Extension; +use axum::{Extension, Form}; use common_error::status_code::StatusCode; use common_telemetry::metric; use schemars::JsonSchema; @@ -36,14 +36,18 @@ pub struct SqlQuery { #[axum_macros::debug_handler] pub async fn sql( State(state): State, - Query(params): Query, + Query(query_params): Query, // TODO(fys): pass _user_info into query context _user_info: Extension, + Form(form_params): Form, ) -> Json { let sql_handler = &state.sql_handler; let start = Instant::now(); - let resp = if let Some(sql) = ¶ms.sql { - match super::query_context_from_db(sql_handler.clone(), params.db) { + let sql = query_params.sql.or(form_params.sql); + let db = query_params.db.or(form_params.db); + + let resp = if let Some(sql) = &sql { + match super::query_context_from_db(sql_handler.clone(), db) { Ok(query_ctx) => { JsonResponse::from_output(sql_handler.do_query(sql, query_ctx).await).await } @@ -59,6 +63,7 @@ pub async fn sql( Json(resp.with_execution_time(start.elapsed().as_millis())) } +// TODO(ruihang): add db param and form data support #[derive(Debug, Default, Serialize, Deserialize, JsonSchema)] pub struct PromqlQuery { pub query: String, diff --git a/src/servers/tests/http/http_handler_test.rs b/src/servers/tests/http/http_handler_test.rs index 5ce7d07ea85..401418908d2 100644 --- a/src/servers/tests/http/http_handler_test.rs +++ b/src/servers/tests/http/http_handler_test.rs @@ -16,6 +16,7 @@ use std::collections::HashMap; use axum::body::Body; use axum::extract::{Json, Query, RawBody, State}; +use axum::Form; use common_telemetry::metric; use metrics::counter; use servers::http::{handler as http_handler, script as script_handler, ApiState, JsonOutput}; @@ -34,6 +35,7 @@ async fn test_sql_not_provided() { }), Query(http_handler::SqlQuery::default()), axum::Extension(UserInfo::default()), + Form(http_handler::SqlQuery::default()), ) .await; assert!(!json.success()); @@ -58,6 +60,34 @@ async fn test_sql_output_rows() { }), query, axum::Extension(UserInfo::default()), + Form(http_handler::SqlQuery::default()), + ) + .await; + assert!(json.success(), "{json:?}"); + assert!(json.error().is_none()); + match &json.output().expect("assertion failed")[0] { + JsonOutput::Records(records) => { + assert_eq!(1, records.num_rows()); + } + _ => unreachable!(), + } +} + +#[tokio::test] +async fn test_sql_form() { + common_telemetry::init_default_ut_logging(); + + let form = create_form(); + let sql_handler = create_testing_sql_query_handler(MemTable::default_numbers_table()); + + let Json(json) = http_handler::sql( + State(ApiState { + sql_handler, + script_handler: None, + }), + Query(http_handler::SqlQuery::default()), + axum::Extension(UserInfo::default()), + form, ) .await; assert!(json.success(), "{json:?}"); @@ -143,6 +173,13 @@ fn create_query() -> Query { }) } +fn create_form() -> Form { + Form(http_handler::SqlQuery { + sql: Some("select sum(uint32s) from numbers limit 20".to_string()), + db: None, + }) +} + /// Currently the payload of response should be simply an empty json "{}"; #[tokio::test] async fn test_health() {