From 9b4e8555e256eedf1bf7f2b6c56b077a7ac4e405 Mon Sep 17 00:00:00 2001 From: jeremyhi Date: Wed, 18 Dec 2024 16:17:34 +0800 Subject: [PATCH] feat: extract hints from http header (#5128) * feat: extract hints from http header * Update src/servers/src/http/hints.rs Co-authored-by: shuiyisong <113876041+shuiyisong@users.noreply.github.com> * chore: by comment * refactor: get instead of loop --------- Co-authored-by: shuiyisong <113876041+shuiyisong@users.noreply.github.com> --- src/servers/src/grpc/database.rs | 57 +---------- src/servers/src/hint_headers.rs | 170 +++++++++++++++++++++++++++++++ src/servers/src/http.rs | 4 +- src/servers/src/http/hints.rs | 30 ++++++ src/servers/src/lib.rs | 1 + 5 files changed, 207 insertions(+), 55 deletions(-) create mode 100644 src/servers/src/hint_headers.rs create mode 100644 src/servers/src/http/hints.rs diff --git a/src/servers/src/grpc/database.rs b/src/servers/src/grpc/database.rs index 572f3c66f4d2..121d8c6c8594 100644 --- a/src/servers/src/grpc/database.rs +++ b/src/servers/src/grpc/database.rs @@ -20,13 +20,11 @@ use common_error::status_code::StatusCode; use common_query::OutputData; use common_telemetry::{debug, warn}; use futures::StreamExt; -use tonic::metadata::{KeyAndValueRef, MetadataMap}; use tonic::{Request, Response, Status, Streaming}; use crate::grpc::greptime_handler::GreptimeRequestHandler; use crate::grpc::{cancellation, TonicResult}; - -pub const GREPTIME_DB_HEADER_HINT_PREFIX: &str = "x-greptime-hint-"; +use crate::hint_headers; pub(crate) struct DatabaseService { handler: GreptimeRequestHandler, @@ -45,7 +43,7 @@ impl GreptimeDatabase for DatabaseService { request: Request, ) -> TonicResult> { let remote_addr = request.remote_addr(); - let hints = extract_hints(request.metadata()); + let hints = hint_headers::extract_hints(request.metadata()); debug!( "GreptimeDatabase::Handle: request from {:?} with hints: {:?}", remote_addr, hints @@ -91,7 +89,7 @@ impl GreptimeDatabase for DatabaseService { request: Request>, ) -> Result, Status> { let remote_addr = request.remote_addr(); - let hints = extract_hints(request.metadata()); + let hints = hint_headers::extract_hints(request.metadata()); debug!( "GreptimeDatabase::HandleRequests: request from {:?} with hints: {:?}", remote_addr, hints @@ -142,52 +140,3 @@ impl GreptimeDatabase for DatabaseService { cancellation::with_cancellation_handler(request_future, cancellation_future).await } } - -fn extract_hints(metadata: &MetadataMap) -> Vec<(String, String)> { - metadata - .iter() - .filter_map(|kv| { - let KeyAndValueRef::Ascii(key, value) = kv else { - return None; - }; - let key = key.as_str(); - let new_key = key.strip_prefix(GREPTIME_DB_HEADER_HINT_PREFIX)?; - let Ok(value) = value.to_str() else { - // Simply return None for non-string values. - return None; - }; - Some((new_key.to_string(), value.trim().to_string())) - }) - .collect() -} - -#[cfg(test)] -mod tests { - use tonic::metadata::MetadataValue; - - use super::*; - - #[test] - fn test_extract_hints() { - let mut metadata = MetadataMap::new(); - let prev = metadata.insert( - "x-greptime-hint-append_mode", - MetadataValue::from_static("true"), - ); - metadata.insert("test-key", MetadataValue::from_static("test-value")); - assert!(prev.is_none()); - let hints = extract_hints(&metadata); - assert_eq!(hints, vec![("append_mode".to_string(), "true".to_string())]); - } - - #[test] - fn extract_hints_ignores_non_ascii_metadata() { - let mut metadata = MetadataMap::new(); - metadata.insert_bin( - "x-greptime-hint-merge_mode-bin", - MetadataValue::from_bytes(b"last_non_null"), - ); - let hints = extract_hints(&metadata); - assert!(hints.is_empty()); - } -} diff --git a/src/servers/src/hint_headers.rs b/src/servers/src/hint_headers.rs new file mode 100644 index 000000000000..6dafd45196b3 --- /dev/null +++ b/src/servers/src/hint_headers.rs @@ -0,0 +1,170 @@ +// 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 http::HeaderMap; +use tonic::metadata::MetadataMap; + +pub const HINT_KEYS: [&str; 5] = [ + "x-greptime-hint-auto_create_table", + "x-greptime-hint-ttl", + "x-greptime-hint-append_mode", + "x-greptime-hint-merge_mode", + "x-greptime-hint-physical_table", +]; + +pub(crate) fn extract_hints(headers: &T) -> Vec<(String, String)> { + let mut hints = Vec::new(); + for key in HINT_KEYS.iter() { + if let Some(value) = headers.get(key) { + let new_key = key.replace("x-greptime-hint-", ""); + hints.push((new_key, value.trim().to_string())); + } + } + hints +} + +pub(crate) trait ToHeaderMap { + fn get(&self, key: &str) -> Option<&str>; +} + +impl ToHeaderMap for MetadataMap { + fn get(&self, key: &str) -> Option<&str> { + self.get(key).and_then(|v| v.to_str().ok()) + } +} + +impl ToHeaderMap for HeaderMap { + fn get(&self, key: &str) -> Option<&str> { + self.get(key).and_then(|v| v.to_str().ok()) + } +} +#[cfg(test)] +mod tests { + use http::header::{HeaderMap, HeaderValue}; + use tonic::metadata::{MetadataMap, MetadataValue}; + + use super::*; + + #[test] + fn test_extract_hints_with_full_header_map() { + let mut headers = HeaderMap::new(); + headers.insert( + "x-greptime-hint-auto_create_table", + HeaderValue::from_static("true"), + ); + headers.insert("x-greptime-hint-ttl", HeaderValue::from_static("3600d")); + headers.insert( + "x-greptime-hint-append_mode", + HeaderValue::from_static("true"), + ); + headers.insert( + "x-greptime-hint-merge_mode", + HeaderValue::from_static("false"), + ); + headers.insert( + "x-greptime-hint-physical_table", + HeaderValue::from_static("table1"), + ); + + let hints = extract_hints(&headers); + + assert_eq!(hints.len(), 5); + assert_eq!( + hints[0], + ("auto_create_table".to_string(), "true".to_string()) + ); + assert_eq!(hints[1], ("ttl".to_string(), "3600d".to_string())); + assert_eq!(hints[2], ("append_mode".to_string(), "true".to_string())); + assert_eq!(hints[3], ("merge_mode".to_string(), "false".to_string())); + assert_eq!( + hints[4], + ("physical_table".to_string(), "table1".to_string()) + ); + } + + #[test] + fn test_extract_hints_with_missing_keys() { + let mut headers = HeaderMap::new(); + headers.insert( + "x-greptime-hint-auto_create_table", + HeaderValue::from_static("true"), + ); + headers.insert("x-greptime-hint-ttl", HeaderValue::from_static("3600d")); + + let hints = extract_hints(&headers); + + assert_eq!(hints.len(), 2); + assert_eq!( + hints[0], + ("auto_create_table".to_string(), "true".to_string()) + ); + assert_eq!(hints[1], ("ttl".to_string(), "3600d".to_string())); + } + + #[test] + fn test_extract_hints_with_metadata_map() { + let mut metadata = MetadataMap::new(); + metadata.insert( + "x-greptime-hint-auto_create_table", + MetadataValue::from_static("true"), + ); + metadata.insert("x-greptime-hint-ttl", MetadataValue::from_static("3600d")); + metadata.insert( + "x-greptime-hint-append_mode", + MetadataValue::from_static("true"), + ); + metadata.insert( + "x-greptime-hint-merge_mode", + MetadataValue::from_static("false"), + ); + metadata.insert( + "x-greptime-hint-physical_table", + MetadataValue::from_static("table1"), + ); + + let hints = extract_hints(&metadata); + + assert_eq!(hints.len(), 5); + assert_eq!( + hints[0], + ("auto_create_table".to_string(), "true".to_string()) + ); + assert_eq!(hints[1], ("ttl".to_string(), "3600d".to_string())); + assert_eq!(hints[2], ("append_mode".to_string(), "true".to_string())); + assert_eq!(hints[3], ("merge_mode".to_string(), "false".to_string())); + assert_eq!( + hints[4], + ("physical_table".to_string(), "table1".to_string()) + ); + } + + #[test] + fn test_extract_hints_with_partial_metadata_map() { + let mut metadata = MetadataMap::new(); + metadata.insert( + "x-greptime-hint-auto_create_table", + MetadataValue::from_static("true"), + ); + metadata.insert("x-greptime-hint-ttl", MetadataValue::from_static("3600d")); + + let hints = extract_hints(&metadata); + + assert_eq!(hints.len(), 2); + assert_eq!( + hints[0], + ("auto_create_table".to_string(), "true".to_string()) + ); + assert_eq!(hints[1], ("ttl".to_string(), "3600d".to_string())); + } +} diff --git a/src/servers/src/http.rs b/src/servers/src/http.rs index 1107870c9a25..9841f02d6ead 100644 --- a/src/servers/src/http.rs +++ b/src/servers/src/http.rs @@ -92,6 +92,7 @@ mod timeout; pub(crate) use timeout::DynamicTimeoutLayer; +mod hints; #[cfg(any(test, feature = "testing"))] pub mod test_helpers; @@ -703,7 +704,8 @@ impl HttpServer { .layer(middleware::from_fn_with_state( AuthState::new(self.user_provider.clone()), authorize::check_http_auth, - )), + )) + .layer(middleware::from_fn(hints::extract_hints)), ) // Handlers for debug, we don't expect a timeout. .nest( diff --git a/src/servers/src/http/hints.rs b/src/servers/src/http/hints.rs new file mode 100644 index 000000000000..4612201880eb --- /dev/null +++ b/src/servers/src/http/hints.rs @@ -0,0 +1,30 @@ +// 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::http::Request; +use axum::middleware::Next; +use axum::response::Response; +use session::context::QueryContext; + +use crate::hint_headers; + +pub async fn extract_hints(mut request: Request, next: Next) -> Response { + let hints = hint_headers::extract_hints(request.headers()); + if let Some(query_ctx) = request.extensions_mut().get_mut::() { + for (key, value) in hints { + query_ctx.set_extension(key, value); + } + } + next.run(request).await +} diff --git a/src/servers/src/lib.rs b/src/servers/src/lib.rs index ce6857c6d23f..92f2b8b9d0ba 100644 --- a/src/servers/src/lib.rs +++ b/src/servers/src/lib.rs @@ -27,6 +27,7 @@ pub mod error; pub mod export_metrics; pub mod grpc; pub mod heartbeat_options; +mod hint_headers; pub mod http; pub mod influxdb; pub mod interceptor;