From 56b3e2073f47213792f6ef3af12da40b21b81571 Mon Sep 17 00:00:00 2001 From: WenyXu Date: Wed, 9 Sep 2026 17:57:49 +0000 Subject: [PATCH] feat: propagate insert WAL policy through query context Signed-off-by: WenyXu --- src/operator/src/insert.rs | 26 ++++- .../src/req_convert/common/partitioner.rs | 3 +- .../src/req_convert/insert/row_to_region.rs | 76 +++++++++++++- .../src/req_convert/insert/stmt_to_region.rs | 2 +- .../src/req_convert/insert/table_to_region.rs | 91 +++++++++-------- src/servers/src/grpc/context_auth.rs | 73 ++++++++++++-- src/servers/src/grpc/greptime_handler.rs | 61 +++++++++++- src/servers/src/hint_headers.rs | 16 +++ src/servers/src/http.rs | 4 +- src/servers/src/http/header.rs | 5 + src/servers/src/http/skip_wal.rs | 56 +++++++++++ src/servers/tests/http/http_handler_test.rs | 99 +++++++++++++++++++ src/session/src/context.rs | 49 +++++++++ src/session/src/hints.rs | 3 + src/session/src/lib.rs | 3 + 15 files changed, 511 insertions(+), 56 deletions(-) create mode 100644 src/servers/src/http/skip_wal.rs diff --git a/src/operator/src/insert.rs b/src/operator/src/insert.rs index 52a73892dc..676638dd73 100644 --- a/src/operator/src/insert.rs +++ b/src/operator/src/insert.rs @@ -265,6 +265,8 @@ impl Inserter { accommodate_existing_schema: bool, is_single_value: bool, ) -> Result { + let skip_wal = ctx.skip_wal(); + // remove empty requests requests.inserts.retain(|req| { req.rows @@ -297,7 +299,7 @@ impl Inserter { instant_table_ids, self.partition_manager.as_ref(), ) - .convert(requests) + .convert(requests, skip_wal) .await?; self.do_request(inserts, &table_infos, &ctx).await @@ -311,6 +313,8 @@ impl Inserter { statement_executor: &StatementExecutor, physical_table: String, ) -> Result { + let skip_wal = ctx.skip_wal(); + // remove empty requests requests.inserts.retain(|req| { req.rows @@ -343,7 +347,7 @@ impl Inserter { .map(|info| (info.name.clone(), info.clone())) .collect::>(); let inserts = RowToRegion::new(name_to_info, instant_table_ids, &self.partition_manager) - .convert(requests) + .convert(requests, skip_wal) .await?; self.do_request(inserts, &table_infos, &ctx).await @@ -354,6 +358,7 @@ impl Inserter { request: TableInsertRequest, ctx: QueryContextRef, ) -> Result { + let skip_wal = ctx.skip_wal(); let catalog = request.catalog_name.as_str(); let schema = request.schema_name.as_str(); let table_name = request.table_name.as_str(); @@ -364,7 +369,7 @@ impl Inserter { let table_info = table.table_info(); let inserts = TableToRegion::new(&table_info, &self.partition_manager) - .convert(request) + .convert(request, skip_wal) .await?; let table_infos = HashMap::from_iter([(table_info.table_id(), table_info.clone())]); @@ -1743,6 +1748,21 @@ mod tests { assert!(table_is_native_histogram(&table)); } + #[test] + fn test_skip_wal_does_not_change_table_options() { + for skip_wal in [false, true] { + let ctx = Arc::new(QueryContext::with( + DEFAULT_CATALOG_NAME, + DEFAULT_SCHEMA_NAME, + )); + ctx.set_skip_wal(skip_wal); + let mut options = Default::default(); + fill_table_options_for_create(&mut options, &AutoCreateTableType::Physical, &ctx); + assert!(!options.contains_key(session::hints::INSERT_SKIP_WAL_HINT)); + assert!(!options.contains_key("skip_wal")); + } + } + #[test] fn test_last_non_null_create_options_preserve_default_without_append_mode() { let ctx = Arc::new(QueryContext::with( diff --git a/src/operator/src/req_convert/common/partitioner.rs b/src/operator/src/req_convert/common/partitioner.rs index 513f2ff030..3353e2fcda 100644 --- a/src/operator/src/req_convert/common/partitioner.rs +++ b/src/operator/src/req_convert/common/partitioner.rs @@ -34,6 +34,7 @@ impl<'a> Partitioner<'a> { &self, table_info: &TableInfo, rows: Rows, + skip_wal: bool, ) -> Result> { let table_id = table_info.table_id(); let requests = self @@ -44,7 +45,7 @@ impl<'a> Partitioner<'a> { .into_iter() .map( |(region_number, (rows, partition_expr_version))| InsertRequest { - skip_wal: false, + skip_wal, region_id: RegionId::new(table_id, region_number).into(), rows: Some(rows), partition_expr_version: partition_expr_version diff --git a/src/operator/src/req_convert/insert/row_to_region.rs b/src/operator/src/req_convert/insert/row_to_region.rs index 9a03cd7121..3e105f760b 100644 --- a/src/operator/src/req_convert/insert/row_to_region.rs +++ b/src/operator/src/req_convert/insert/row_to_region.rs @@ -45,6 +45,7 @@ impl<'a> RowToRegion<'a> { pub async fn convert( &self, requests: RowInsertRequests, + skip_wal: bool, ) -> Result { let mut region_request = Vec::with_capacity(requests.inserts.len()); let mut instant_request = Vec::with_capacity(requests.inserts.len()); @@ -55,7 +56,7 @@ impl<'a> RowToRegion<'a> { let table_id = table_info.table_id(); let requests = Partitioner::new(self.partition_manager) - .partition_insert_requests(table_info, rows) + .partition_insert_requests(table_info, rows, skip_wal) .await?; if self.instant_table_ids.contains(&table_id) { @@ -81,3 +82,76 @@ impl<'a> RowToRegion<'a> { .context(TableNotFoundSnafu { table_name }) } } + +#[cfg(test)] +mod tests { + use std::sync::Arc; + + use api::v1::helper::tag_column_schema; + use api::v1::value::ValueData; + use api::v1::{ColumnDataType, Row, RowInsertRequest, Rows, Value}; + + use super::*; + use crate::tests::{ + create_partition_rule_manager, new_test_table_info, prepare_mocked_backend, + }; + + #[tokio::test] + async fn test_partitioned_insert_skip_wal_normal_and_instant() { + let backend = prepare_mocked_backend().await; + let partition_manager = create_partition_rule_manager(backend).await; + let table_info = Arc::new(new_test_table_info(1, "table_1", [1, 2, 3].into_iter())); + for instant in [false, true] { + let instant_table_ids = if instant { + HashSet::from_iter([1]) + } else { + HashSet::default() + }; + let converter = RowToRegion::new( + HashMap::from_iter([("table_1".to_string(), table_info.clone())]), + instant_table_ids, + &partition_manager, + ); + for skip_wal in [false, true] { + let requests = RowInsertRequests { + inserts: vec![RowInsertRequest { + table_name: "table_1".to_string(), + rows: Some(Rows { + schema: vec![tag_column_schema("a", ColumnDataType::Int32)], + rows: [1, 11, 101] + .into_iter() + .map(|value| Row { + values: vec![Value { + value_data: Some(ValueData::I32Value(value)), + }], + }) + .collect(), + }), + }], + }; + let result = converter.convert(requests, skip_wal).await.unwrap(); + let (selected, other) = if instant { + (result.instant_requests, result.normal_requests) + } else { + (result.normal_requests, result.instant_requests) + }; + assert!(other.requests.is_empty()); + assert_eq!(selected.requests.len(), 3); + assert!( + selected + .requests + .iter() + .all(|request| request.skip_wal == skip_wal) + ); + assert_eq!( + selected + .requests + .iter() + .map(|request| request.rows.as_ref().unwrap().rows.len()) + .sum::(), + 3 + ); + } + } + } +} diff --git a/src/operator/src/req_convert/insert/stmt_to_region.rs b/src/operator/src/req_convert/insert/stmt_to_region.rs index b2424d3a17..e811facf74 100644 --- a/src/operator/src/req_convert/insert/stmt_to_region.rs +++ b/src/operator/src/req_convert/insert/stmt_to_region.rs @@ -154,7 +154,7 @@ impl<'a> StatementToRegion<'a> { } let requests = Partitioner::new(self.partition_manager) - .partition_insert_requests(&table_info, Rows { schema, rows }) + .partition_insert_requests(&table_info, Rows { schema, rows }, query_ctx.skip_wal()) .await?; let requests = RegionInsertRequests { requests }; if table_info.is_ttl_instant_table() { diff --git a/src/operator/src/req_convert/insert/table_to_region.rs b/src/operator/src/req_convert/insert/table_to_region.rs index 2cddea1b2d..7e7e0c1070 100644 --- a/src/operator/src/req_convert/insert/table_to_region.rs +++ b/src/operator/src/req_convert/insert/table_to_region.rs @@ -39,6 +39,7 @@ impl<'a> TableToRegion<'a> { pub async fn convert( &self, request: TableInsertRequest, + skip_wal: bool, ) -> Result { let row_count = row_count(&request.columns_values)?; let schema = column_schema(self.table_info, &request.columns_values)?; @@ -46,7 +47,7 @@ impl<'a> TableToRegion<'a> { let rows = Rows { schema, rows }; let requests = Partitioner::new(self.partition_manager) - .partition_insert_requests(self.table_info, rows) + .partition_insert_requests(self.table_info, rows, skip_wal) .await?; let requests = RegionInsertRequests { requests }; @@ -100,49 +101,56 @@ mod tests { let converter = TableToRegion::new(&table_info, &partition_manager); - let table_request = build_table_request(Arc::new(Int32Vector::from(vec![ - Some(1), - None, - Some(11), - Some(101), - ]))); - let versions = partition_manager - .find_physical_partition_info(1) - .await - .unwrap() - .partitions - .iter() - .map(|p| (p.id.as_u64(), p.partition_expr_version)) - .collect::>(); + for skip_wal in [false, true] { + let table_request = build_table_request(Arc::new(Int32Vector::from(vec![ + Some(1), + None, + Some(11), + Some(101), + ]))); + let versions = partition_manager + .find_physical_partition_info(1) + .await + .unwrap() + .partitions + .iter() + .map(|p| (p.id.as_u64(), p.partition_expr_version)) + .collect::>(); - let region_requests = converter.convert(table_request).await.unwrap(); - let mut region_id_to_region_requests = region_requests - .normal_requests - .requests - .into_iter() - .map(|r| (r.region_id, r)) - .collect::>(); + let region_requests = converter.convert(table_request, skip_wal).await.unwrap(); + let mut region_id_to_region_requests = region_requests + .normal_requests + .requests + .into_iter() + .map(|r| (r.region_id, r)) + .collect::>(); - let region_id = RegionId::new(1, 1).as_u64(); - let region_request = region_id_to_region_requests.remove(®ion_id).unwrap(); - assert_eq!( - region_request, - build_region_request(vec![Some(101)], region_id, versions[®ion_id]) - ); + let region_id = RegionId::new(1, 1).as_u64(); + let region_request = region_id_to_region_requests.remove(®ion_id).unwrap(); + assert_eq!( + region_request, + build_region_request(vec![Some(101)], region_id, versions[®ion_id], skip_wal) + ); - let region_id = RegionId::new(1, 2).as_u64(); - let region_request = region_id_to_region_requests.remove(®ion_id).unwrap(); - assert_eq!( - region_request, - build_region_request(vec![Some(11)], region_id, versions[®ion_id]) - ); + let region_id = RegionId::new(1, 2).as_u64(); + let region_request = region_id_to_region_requests.remove(®ion_id).unwrap(); + assert_eq!( + region_request, + build_region_request(vec![Some(11)], region_id, versions[®ion_id], skip_wal) + ); - let region_id = RegionId::new(1, 3).as_u64(); - let region_request = region_id_to_region_requests.remove(®ion_id).unwrap(); - assert_eq!( - region_request, - build_region_request(vec![Some(1), None], region_id, versions[®ion_id]) - ); + let region_id = RegionId::new(1, 3).as_u64(); + let region_request = region_id_to_region_requests.remove(®ion_id).unwrap(); + assert_eq!( + region_request, + build_region_request( + vec![Some(1), None], + region_id, + versions[®ion_id], + skip_wal + ) + ); + } } fn build_table_request(vector: VectorRef) -> TableInsertRequest { @@ -158,9 +166,10 @@ mod tests { rows: Vec>, region_id: u64, version: Option, + skip_wal: bool, ) -> RegionInsertRequest { RegionInsertRequest { - skip_wal: false, + skip_wal, region_id, rows: Some(Rows { schema: vec![tag_column_schema("a", ColumnDataType::Int32)], diff --git a/src/servers/src/grpc/context_auth.rs b/src/servers/src/grpc/context_auth.rs index 39c4fc5c88..0a44c0cb47 100644 --- a/src/servers/src/grpc/context_auth.rs +++ b/src/servers/src/grpc/context_auth.rs @@ -21,6 +21,7 @@ use common_catalog::consts::{DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME}; use common_catalog::parse_catalog_and_schema_from_db_string; use common_error::ext::ErrorExt; use session::context::{Channel, QueryContextBuilder, QueryContextRef}; +use session::hints::INSERT_SKIP_WAL_HINT; use snafu::{OptionExt, ResultExt}; use tonic::Status; use tonic::metadata::MetadataMap; @@ -28,6 +29,7 @@ use tonic::metadata::MetadataMap; use crate::error::Error::UnsupportedAuthScheme; use crate::error::{AuthSnafu, InvalidParameterSnafu, NotFoundAuthHeaderSnafu, Result}; use crate::grpc::TonicResult; +use crate::hint_headers; use crate::http::AUTHORIZATION_HEADER; use crate::http::header::constants::GREPTIME_DB_HEADER_NAME; use crate::metrics::METRIC_AUTH_FAILURE; @@ -45,13 +47,25 @@ pub fn create_query_context_from_grpc_metadata( ) }; - Ok(Arc::new( - QueryContextBuilder::default() - .current_catalog(catalog) - .current_schema(schema) - .channel(Channel::Grpc) - .build(), - )) + let ctx = QueryContextBuilder::default() + .current_catalog(catalog) + .current_schema(schema) + .channel(Channel::Grpc) + .build(); + // OTEL Arrow uses ordinary inserts. Accept only its request-level WAL hint, + // leaving unrelated hints and reserved internal extensions unchanged. + for (key, value) in hint_headers::extract_hints(headers) { + if key == INSERT_SKIP_WAL_HINT { + let skip_wal = value.parse::().map_err(|_| { + InvalidParameterSnafu { + reason: format!("Invalid {key} hint: expected true or false, got {value:?}"), + } + .build() + })?; + ctx.set_skip_wal(skip_wal); + } + } + Ok(Arc::new(ctx)) } /// Helper function to extract a header from the metadata map. @@ -161,3 +175,48 @@ pub async fn auth( .inc(); }) } + +#[cfg(test)] +mod tests { + use session::hints::{HINTS_KEY, REMOTE_QUERY_ID_EXTENSION_KEY, RESERVED_EXTENSION_KEYS}; + + use super::*; + + #[test] + fn test_arrow_insert_hint_does_not_accept_reserved_extensions() { + let mut headers = MetadataMap::new(); + assert_eq!( + create_query_context_from_grpc_metadata(&headers) + .unwrap() + .extension(INSERT_SKIP_WAL_HINT), + None + ); + for (value, expected) in [("true", true), ("false", false)] { + let mut hints = format!("insert_skip_wal={value},ttl=7d"); + for key in RESERVED_EXTENSION_KEYS { + hints.push_str(&format!(",{key}=external")); + } + headers.insert(HINTS_KEY, hints.parse().unwrap()); + let ctx = create_query_context_from_grpc_metadata(&headers).unwrap(); + assert_eq!(ctx.skip_wal(), expected); + assert_eq!(ctx.extension(INSERT_SKIP_WAL_HINT), None); + assert_eq!(ctx.extension("ttl"), None); + for key in RESERVED_EXTENSION_KEYS { + if key == REMOTE_QUERY_ID_EXTENSION_KEY { + // The builder generates this ID; external hints must not replace it. + assert!(ctx.remote_query_id().is_some()); + assert_ne!(ctx.extension(key), Some("external")); + } else { + assert_eq!(ctx.extension(key), None); + } + } + } + for value in ["", "TRUE", "1", "invalid"] { + headers.insert( + HINTS_KEY, + format!("insert_skip_wal={value}").parse().unwrap(), + ); + assert!(create_query_context_from_grpc_metadata(&headers).is_err()); + } + } +} diff --git a/src/servers/src/grpc/greptime_handler.rs b/src/servers/src/grpc/greptime_handler.rs index 98786d30ad..7ae59a024b 100644 --- a/src/servers/src/grpc/greptime_handler.rs +++ b/src/servers/src/grpc/greptime_handler.rs @@ -38,7 +38,7 @@ use common_telemetry::{debug, error, tracing, warn}; use common_time::timezone::parse_timezone; use futures_util::StreamExt; use session::context::{Channel, QueryContextBuilder, QueryContextRef}; -use session::hints::{READ_PREFERENCE_HINT, is_reserved_extension_key}; +use session::hints::{INSERT_SKIP_WAL_HINT, READ_PREFERENCE_HINT, is_reserved_extension_key}; use snafu::{OptionExt, ResultExt}; use tokio::sync::mpsc; use tokio::sync::mpsc::error::TrySendError; @@ -255,6 +255,16 @@ pub(crate) fn create_query_context( } for (key, value) in extensions { + if key == INSERT_SKIP_WAL_HINT { + let skip_wal = value.parse::().map_err(|_| { + UnknownHintSnafu { + hint: format!("{key}={value}"), + } + .build() + })?; + ctx_builder = ctx_builder.skip_wal(skip_wal); + continue; + } if is_reserved_extension_key(&key) { debug!( key = key.as_str(), @@ -322,6 +332,55 @@ mod tests { use super::*; use crate::error::{ExecuteGrpcRequestSnafu, InvalidParameterSnafu}; + #[test] + fn test_create_query_context_typed_skip_wal() { + let ctx = create_query_context(Channel::Grpc, None, vec![], HashMap::new()).unwrap(); + assert!(!ctx.skip_wal()); + let legacy = create_query_context( + Channel::Grpc, + None, + vec![("skip_wal".to_string(), "true".to_string())], + HashMap::new(), + ) + .unwrap(); + assert!(!legacy.skip_wal()); + assert_eq!(legacy.extension("skip_wal"), Some("true")); + for (value, expected) in [("true", true), ("false", false)] { + let ctx = create_query_context( + Channel::Grpc, + None, + vec![(INSERT_SKIP_WAL_HINT.to_string(), value.to_string())], + HashMap::new(), + ) + .unwrap(); + assert_eq!(ctx.skip_wal(), expected); + assert_eq!(ctx.extension(INSERT_SKIP_WAL_HINT), None); + } + for value in ["", "TRUE", "1", "invalid"] { + assert!( + create_query_context( + Channel::Grpc, + None, + vec![(INSERT_SKIP_WAL_HINT.to_string(), value.to_string())], + HashMap::new() + ) + .is_err() + ); + } + let ctx = create_query_context( + Channel::Grpc, + None, + vec![ + (INSERT_SKIP_WAL_HINT.to_string(), "true".to_string()), + (INSERT_SKIP_WAL_HINT.to_string(), "false".to_string()), + ], + HashMap::new(), + ) + .unwrap(); + assert!(!ctx.skip_wal()); + assert_eq!(ctx.extension(INSERT_SKIP_WAL_HINT), None); + } + #[test] fn test_create_query_context() { let header = RequestHeader { diff --git a/src/servers/src/hint_headers.rs b/src/servers/src/hint_headers.rs index 5f64be584f..3d13dd9bf5 100644 --- a/src/servers/src/hint_headers.rs +++ b/src/servers/src/hint_headers.rs @@ -59,6 +59,22 @@ mod tests { use super::*; + #[test] + fn test_extract_skip_wal_hint() { + use session::hints::INSERT_SKIP_WAL_HINT; + + let mut headers = HeaderMap::new(); + headers.insert(HINTS_KEY, HeaderValue::from_static("insert_skip_wal=true")); + let mut metadata = MetadataMap::new(); + metadata.insert( + HINTS_KEY, + MetadataValue::from_static("insert_skip_wal=true"), + ); + let expected = vec![(INSERT_SKIP_WAL_HINT.to_string(), "true".to_string())]; + assert_eq!(extract_hints(&headers), expected); + assert_eq!(extract_hints(&metadata), expected); + } + #[test] fn test_extract_hints_with_full_header_map() { let mut headers = HeaderMap::new(); diff --git a/src/servers/src/http.rs b/src/servers/src/http.rs index 311ed48d91..40b7bb515b 100644 --- a/src/servers/src/http.rs +++ b/src/servers/src/http.rs @@ -117,6 +117,7 @@ mod client_ip; use crate::prom_remote_write::validation::PromValidationMode; mod hints; mod read_preference; +mod skip_wal; #[cfg(any(test, feature = "testing"))] pub mod test_helpers; @@ -1059,7 +1060,8 @@ impl HttpServer { .layer(middleware::from_fn(client_ip::log_error_with_client_ip)) .layer(middleware::from_fn( read_preference::extract_read_preference, - )), + )) + .layer(middleware::from_fn(skip_wal::extract_skip_wal)), ); // Debug handlers are part of the complete router; the API listener hides diff --git a/src/servers/src/http/header.rs b/src/servers/src/http/header.rs index 1dc94bb76e..625bee41ee 100644 --- a/src/servers/src/http/header.rs +++ b/src/servers/src/http/header.rs @@ -44,6 +44,7 @@ pub mod constants { pub const GREPTIME_DB_HEADER_METRICS: &str = "x-greptime-metrics"; pub const GREPTIME_DB_HEADER_NAME: &str = "x-greptime-db-name"; pub const GREPTIME_DB_HEADER_READ_PREFERENCE: &str = "x-greptime-read-preference"; + pub const GREPTIME_INSERT_SKIP_WAL_HEADER_NAME: &str = "x-greptime-insert-skip-wal"; pub const GREPTIME_TIMEZONE_HEADER_NAME: &str = "x-greptime-timezone"; pub const GREPTIME_DB_HEADER_ERROR_CODE: &str = common_error::GREPTIME_DB_HEADER_ERROR_CODE; @@ -94,6 +95,10 @@ pub static GREPTIME_TIMEZONE_HEADER_NAME: HeaderName = pub static GREPTIME_DB_HEADER_READ_PREFERENCE: HeaderName = HeaderName::from_static(constants::GREPTIME_DB_HEADER_READ_PREFERENCE); +/// Request-level WAL policy, independent of the table-level skip_wal option. +pub static GREPTIME_INSERT_SKIP_WAL_HEADER_NAME: HeaderName = + HeaderName::from_static(constants::GREPTIME_INSERT_SKIP_WAL_HEADER_NAME); + pub static CONTENT_TYPE_PROTOBUF_STR: &str = "application/x-protobuf"; pub static CONTENT_TYPE_PROTOBUF: HeaderValue = HeaderValue::from_static(CONTENT_TYPE_PROTOBUF_STR); pub static CONTENT_ENCODING_SNAPPY: HeaderValue = HeaderValue::from_static("snappy"); diff --git a/src/servers/src/http/skip_wal.rs b/src/servers/src/http/skip_wal.rs new file mode 100644 index 0000000000..ea91552789 --- /dev/null +++ b/src/servers/src/http/skip_wal.rs @@ -0,0 +1,56 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use axum::body::Body; +use axum::http::Request; +use axum::middleware::Next; +use axum::response::{IntoResponse, Response}; +use session::context::QueryContext; + +use crate::error::InvalidParameterSnafu; +use crate::http::header::GREPTIME_INSERT_SKIP_WAL_HEADER_NAME; +use crate::http::result::error_result::ErrorResponse; + +/// Extract the request-level WAL policy from the dedicated HTTP header. +pub async fn extract_skip_wal(mut request: Request, next: Next) -> Response { + let skip_wal = match request.headers().get(&GREPTIME_INSERT_SKIP_WAL_HEADER_NAME) { + None => false, + Some(value) => match value + .to_str() + .ok() + .and_then(|value| value.parse::().ok()) + { + Some(skip_wal) => skip_wal, + None => { + return ( + http::StatusCode::BAD_REQUEST, + ErrorResponse::from_error( + InvalidParameterSnafu { + reason: format!( + "{} must be true or false", + GREPTIME_INSERT_SKIP_WAL_HEADER_NAME + ), + } + .build(), + ), + ) + .into_response(); + } + }, + }; + if let Some(query_ctx) = request.extensions_mut().get_mut::() { + query_ctx.set_skip_wal(skip_wal); + } + next.run(request).await +} diff --git a/src/servers/tests/http/http_handler_test.rs b/src/servers/tests/http/http_handler_test.rs index 7456d28b36..42f53c370c 100644 --- a/src/servers/tests/http/http_handler_test.rs +++ b/src/servers/tests/http/http_handler_test.rs @@ -980,3 +980,102 @@ async fn get_body(response: Response) -> Bytes { .await .unwrap() } + +#[tokio::test] +async fn test_http_skip_wal_header() { + struct SkipWalSqlHandler { + inner: ServerSqlQueryHandlerRef, + observed: AtomicUsize, + } + + #[async_trait] + impl SqlQueryHandler for SkipWalSqlHandler { + async fn do_query(&self, query: &str, ctx: QueryContextRef) -> Vec> { + self.observed + .store(usize::from(ctx.skip_wal()), Ordering::Relaxed); + self.inner.do_query(query, ctx).await + } + + async fn do_analyze_stream_query( + &self, + query: &str, + ctx: QueryContextRef, + ) -> Result { + self.inner.do_analyze_stream_query(query, ctx).await + } + + async fn do_exec_plan( + &self, + plan: LogicalPlan, + stmt: Option, + ctx: QueryContextRef, + ) -> Result { + self.inner.do_exec_plan(plan, stmt, ctx).await + } + + async fn do_promql_query( + &self, + query: &PromQuery, + ctx: QueryContextRef, + ) -> Vec> { + self.inner.do_promql_query(query, ctx).await + } + + async fn do_describe( + &self, + stmt: Statement, + ctx: QueryContextRef, + ) -> Result> { + self.inner.do_describe(stmt, ctx).await + } + + async fn is_valid_schema(&self, catalog: &str, schema: &str) -> Result { + self.inner.is_valid_schema(catalog, schema).await + } + } + + let handler = Arc::new(SkipWalSqlHandler { + inner: create_testing_sql_query_handler(MemTable::default_numbers_table()), + observed: AtomicUsize::new(2), + }); + let server = HttpServerBuilder::new(HttpOptions::default()) + .with_sql_handler(handler.clone()) + .build(); + let client = TestClient::new(server.build(server.make_app()).unwrap()).await; + for (header, expected) in [ + (Some(("x-greptime-insert-skip-wal", "true")), Some(true)), + (Some(("x-greptime-insert-skip-wal", "false")), Some(false)), + (None, Some(false)), + (Some(("x-greptime-insert-skip-wal", "yes")), None), + (Some(("x-greptime-insert-skip-wal", "1")), None), + (Some(("x-greptime-insert-skip-wal", "")), None), + (Some(("x-greptime-skip-wal", "true")), Some(false)), + (Some(("x-greptime-hints", "skip_wal=true")), Some(false)), + ( + Some(("x-greptime-hints", "insert_skip_wal=true")), + Some(false), + ), + ] { + // Sentinel also proves invalid values are rejected before the SQL handler. + handler.observed.store(2, Ordering::Relaxed); + let mut request = client.get("/v1/sql?sql=SELECT%201"); + if let Some((key, value)) = header { + request = request.header(key, value); + } + let response = request.send().await; + match expected { + Some(expected) => { + assert_eq!(response.status(), StatusCode::OK, "{header:?}"); + assert_eq!( + handler.observed.load(Ordering::Relaxed), + usize::from(expected), + "{header:?}" + ); + } + None => { + assert_eq!(response.status(), StatusCode::BAD_REQUEST, "{header:?}"); + assert_eq!(handler.observed.load(Ordering::Relaxed), 2); + } + } + } +} diff --git a/src/session/src/context.rs b/src/session/src/context.rs index d383aef477..2e8b49dae3 100644 --- a/src/session/src/context.rs +++ b/src/session/src/context.rs @@ -153,6 +153,15 @@ impl QueryContextBuilder { self } + pub fn skip_wal(mut self, skip_wal: bool) -> Self { + self.mutable_session_data + .get_or_insert_default() + .write() + .unwrap() + .skip_wal = skip_wal; + self + } + pub fn read_preference(mut self, read_preference: ReadPreference) -> Self { self.mutable_session_data .get_or_insert_default() @@ -348,6 +357,15 @@ impl QueryContext { self.mutable_session_data.write().unwrap().timezone = timezone; } + /// Returns whether ordinary inserts in this request should skip WAL. + pub fn skip_wal(&self) -> bool { + self.mutable_session_data.read().unwrap().skip_wal + } + + pub fn set_skip_wal(&self, skip_wal: bool) { + self.mutable_session_data.write().unwrap().skip_wal = skip_wal; + } + pub fn read_preference(&self) -> ReadPreference { self.mutable_session_data.read().unwrap().read_preference } @@ -739,6 +757,37 @@ mod test { assert_eq!(fork.current_schema(), "private"); } + #[test] + fn test_skip_wal_default_builder_and_fork() { + let default_context = QueryContext::with(DEFAULT_CATALOG_NAME, "public"); + assert!(!default_context.skip_wal()); + assert!(!QueryContextBuilder::default().build().skip_wal()); + let context = QueryContextBuilder::default().skip_wal(true).build(); + assert!(context.skip_wal()); + let fork = context.fork(); + assert!(fork.skip_wal()); + fork.set_skip_wal(false); + assert!(context.skip_wal()); + assert!(!fork.skip_wal()); + context.set_skip_wal(false); + fork.set_skip_wal(true); + assert!(!context.skip_wal()); + assert!(fork.skip_wal()); + } + + #[test] + fn test_skip_wal_is_not_serialized_in_query_context() { + let context = QueryContextBuilder::default().skip_wal(true).build(); + let api_context: api::v1::QueryContext = context.into(); + assert!( + !api_context + .extensions + .contains_key(crate::hints::INSERT_SKIP_WAL_HINT) + ); + let restored: QueryContext = api_context.into(); + assert!(!restored.skip_wal()); + } + #[test] fn test_api_query_context_roundtrip_with_sequences() { let api_ctx = api::v1::QueryContext { diff --git a/src/session/src/hints.rs b/src/session/src/hints.rs index e656b34c82..971426f4cc 100644 --- a/src/session/src/hints.rs +++ b/src/session/src/hints.rs @@ -23,6 +23,9 @@ pub const SUPPORT_FLIGHT_METRICS_BEFORE_BATCH_EXTENSION_KEY: &str = "query.support_flight_metrics_before_batch"; pub const LIVE_ANALYZE_METRICS_EXTENSION_KEY: &str = "query.live_analyze_metrics"; +/// Skip WAL for this insert only; never persisted as a table option. +pub const INSERT_SKIP_WAL_HINT: &str = "insert_skip_wal"; + pub const READ_PREFERENCE_HINT: &str = "read_preference"; pub const RESERVED_EXTENSION_KEYS: [&str; 4] = [ REMOTE_QUERY_ID_EXTENSION_KEY, diff --git a/src/session/src/lib.rs b/src/session/src/lib.rs index c1414fc5a3..40d2e3e67f 100644 --- a/src/session/src/lib.rs +++ b/src/session/src/lib.rs @@ -60,6 +60,8 @@ pub(crate) struct MutableInner { timezone: Timezone, query_timeout: Option, read_preference: ReadPreference, + /// Request-level WAL policy for ordinary inserts. + skip_wal: bool, #[debug(skip)] pub(crate) cursors: HashMap>, /// Warning messages for MySQL SHOW WARNINGS support @@ -74,6 +76,7 @@ impl Default for MutableInner { timezone: get_timezone(None).clone(), query_timeout: None, read_preference: ReadPreference::Leader, + skip_wal: false, cursors: HashMap::with_capacity(0), warnings: VecDeque::new(), }