From 1fc2bef02e1a7c938049922b32a95970c6986daa Mon Sep 17 00:00:00 2001 From: discord9 <55937128+discord9@users.noreply.github.com> Date: Fri, 9 Oct 2026 15:33:30 +0800 Subject: [PATCH] feat: validate and snapshot per-query option overrides Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> --- src/session/src/context.rs | 111 ++++++++++++++++++++++++++- src/session/src/lib.rs | 1 + src/session/src/query_options.rs | 128 +++++++++++++++++++++++++++++++ 3 files changed, 236 insertions(+), 4 deletions(-) create mode 100644 src/session/src/query_options.rs diff --git a/src/session/src/context.rs b/src/session/src/context.rs index 0bb3f0f6a31..a600980e130 100644 --- a/src/session/src/context.rs +++ b/src/session/src/context.rs @@ -443,6 +443,49 @@ impl QueryContext { self.extensions.clone() } + /// Returns normalized query options from the request over session defaults. + pub fn effective_query_options( + &self, + ) -> Result, crate::session_config::Error> { + let mut request_options = HashMap::new(); + for (key, value) in &self.extensions { + if crate::hints::is_reserved_extension_key(key) { + continue; + } + if let Some((canonical, normalized)) = + crate::query_options::parse_query_option(key, value)? + { + if request_options + .get(&canonical) + .is_some_and(|previous| previous != &normalized) + { + return crate::session_config::InvalidConfigValueSnafu { + name: key, + value, + hint: "Conflicting aliases for the same query option", + } + .fail(); + } + request_options.insert(canonical, normalized); + } + } + let mut effective = self.configuration_parameter.query_options(); + effective.extend(request_options); + Ok(effective) + } + + /// Creates a per-query immutable option snapshot while retaining shared mutable session data. + pub fn query_option_snapshot(&self) -> Result { + let mut snapshot = self.clone(); + snapshot.configuration_parameter = Arc::new(self.configuration_parameter.as_ref().clone()); + let options = snapshot.effective_query_options()?; + snapshot + .extensions + .retain(|key, _| crate::query_options::canonical_query_option_name(key).is_none()); + snapshot.extensions.extend(options); + Ok(snapshot) + } + /// Default to double quote and fallback to back quote pub fn quote_style(&self) -> char { if self.sql_dialect().is_delimited_identifier_start('"') { @@ -691,7 +734,7 @@ pub struct ConfigurationVariables { postgres_bytea_output: ArcSwap, pg_datestyle_format: ArcSwap<(PGDateTimeStyle, PGDateOrder)>, pg_intervalstyle_format: ArcSwap, - allow_query_fallback: ArcSwap, + query_options: RwLock>, } impl Clone for ConfigurationVariables { @@ -700,7 +743,7 @@ impl Clone for ConfigurationVariables { postgres_bytea_output: ArcSwap::new(self.postgres_bytea_output.load().clone()), pg_datestyle_format: ArcSwap::new(self.pg_datestyle_format.load().clone()), pg_intervalstyle_format: ArcSwap::new(self.pg_intervalstyle_format.load().clone()), - allow_query_fallback: ArcSwap::new(self.allow_query_fallback.load().clone()), + query_options: RwLock::new(self.query_options.read().unwrap().clone()), } } } @@ -735,11 +778,39 @@ impl ConfigurationVariables { } pub fn allow_query_fallback(&self) -> bool { - **self.allow_query_fallback.load() + self.query_options().get("query.allow_query_fallback") == Some(&"true".to_string()) } pub fn set_allow_query_fallback(&self, allow: bool) { - self.allow_query_fallback.swap(Arc::new(allow)); + self.query_options + .write() + .unwrap() + .insert("query.allow_query_fallback".to_string(), allow.to_string()); + } + + /// Validates and stores one canonicalized per-session query option. + pub fn set_query_option( + &self, + key: &str, + value: &str, + ) -> Result<(), crate::session_config::Error> { + let (canonical, normalized) = crate::query_options::parse_query_option(key, value)? + .ok_or_else(|| crate::session_config::Error::InvalidConfigValue { + name: key.to_string(), + value: value.to_string(), + hint: "Not a supported query option".to_string(), + location: snafu::Location::new(file!(), line!(), column!()), + })?; + self.query_options + .write() + .unwrap() + .insert(canonical, normalized); + Ok(()) + } + + /// Returns the session's canonical query option overrides. + pub fn query_options(&self) -> HashMap { + self.query_options.read().unwrap().clone() } } @@ -752,6 +823,38 @@ mod test { use crate::Session; use crate::context::{Channel, *}; + #[test] + fn query_options_snapshot_validates_and_overlays_request_options() { + let mut ctx = QueryContextBuilder::default().build(); + ctx.configuration_parameter() + .set_query_option("query_parallelism", "8") + .unwrap(); + ctx.set_extension("query_parallelism", "4"); + ctx.set_extension("flow.return_region_seq", "true"); + let snapshot = ctx.query_option_snapshot().unwrap(); + assert_eq!(snapshot.extension("query.parallelism"), Some("4")); + assert_eq!(snapshot.extension("query_parallelism"), None); + assert_eq!(snapshot.extension("flow.return_region_seq"), Some("true")); + ctx.configuration_parameter() + .set_query_option("query.parallelism", "16") + .unwrap(); + assert_eq!(snapshot.extension("query.parallelism"), Some("4")); + + let mut aliases = QueryContextBuilder::default().build(); + aliases.set_extension("query_fallback", "TRUE"); + aliases.set_extension("query.allow_query_fallback", "true"); + assert_eq!( + aliases + .effective_query_options() + .unwrap() + .get("query.allow_query_fallback") + .map(String::as_str), + Some("true") + ); + aliases.set_extension("allow_query_fallback", "false"); + assert!(aliases.query_option_snapshot().is_err()); + } + #[test] fn test_session() { let session = Session::new( diff --git a/src/session/src/lib.rs b/src/session/src/lib.rs index 40d2e3e67fe..478458180ef 100644 --- a/src/session/src/lib.rs +++ b/src/session/src/lib.rs @@ -16,6 +16,7 @@ pub mod context; pub mod hints; pub mod protocol_ctx; pub mod query_id; +pub mod query_options; pub mod session_config; pub mod table_name; diff --git a/src/session/src/query_options.rs b/src/session/src/query_options.rs new file mode 100644 index 00000000000..4e7d47b344a --- /dev/null +++ b/src/session/src/query_options.rs @@ -0,0 +1,128 @@ +// 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 crate::session_config::{Error, InvalidConfigValueSnafu}; + +/// Returns the canonical name for a supported option key. +pub fn canonical_query_option_name(key: &str) -> Option<&'static str> { + let key = key.to_ascii_lowercase(); + match key.as_str() { + "query_parallelism" | "query.parallelism" => Some("query.parallelism"), + "query_fallback" | "allow_query_fallback" | "query.allow_query_fallback" => { + Some("query.allow_query_fallback") + } + "query.enable_remote_dynamic_filter_pushdown" => { + Some("query.enable_remote_dynamic_filter_pushdown") + } + _ => DATAFUSION_OPTIONS + .iter() + .find(|name| **name == key) + .copied(), + } +} + +const DATAFUSION_OPTIONS: &[&str] = &[ + "datafusion.optimizer.repartition_joins", + "datafusion.optimizer.repartition_aggregations", + "datafusion.optimizer.repartition_sorts", + "datafusion.optimizer.repartition_windows", + "datafusion.optimizer.enable_round_robin_repartition", + "datafusion.optimizer.prefer_existing_sort", + "datafusion.optimizer.prefer_hash_join", + "datafusion.optimizer.join_reordering", + "datafusion.optimizer.enable_topk_aggregation", + "datafusion.optimizer.enable_dynamic_filter_pushdown", +]; + +/// Validate a query option, returning its canonical name and normalized value. +pub fn parse_query_option(key: &str, value: &str) -> Result, Error> { + let key_lower = key.to_ascii_lowercase(); + let canonical = match canonical_query_option_name(key) { + Some(name) => name, + None if key_lower.starts_with("query.") || key_lower.starts_with("datafusion.") => { + return InvalidConfigValueSnafu { + name: key, + value, + hint: "Unknown or unsupported query option", + } + .fail(); + } + None => return Ok(None), + }; + + let normalized = if canonical == "query.parallelism" { + match value.parse::() { + Ok(n @ 1..=1024) => n.to_string(), + _ => { + return InvalidConfigValueSnafu { + name: key, + value, + hint: "query.parallelism must be between 1 and 1024", + } + .fail(); + } + } + } else if canonical.starts_with("query.") { + match value.to_ascii_lowercase().as_str() { + "true" => "true".to_string(), + "false" => "false".to_string(), + _ => { + return InvalidConfigValueSnafu { + name: key, + value, + hint: "Expected true or false", + } + .fail(); + } + } + } else { + let mut options = datafusion_common::config::ConfigOptions::default(); + if let Err(error) = options.set(canonical, value) { + return InvalidConfigValueSnafu { + name: key, + value, + hint: error.to_string(), + } + .fail(); + } + value.to_ascii_lowercase() + }; + Ok(Some((canonical.to_string(), normalized))) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn canonicalizes_and_validates_query_options() { + assert_eq!( + parse_query_option("QUERY_PARALLELISM", "1024").unwrap(), + Some(("query.parallelism".into(), "1024".into())) + ); + assert!(parse_query_option("query.parallelism", "0").is_err()); + assert!(parse_query_option("query.parallelism", "1025").is_err()); + assert_eq!( + parse_query_option("query_fallback", "TRUE").unwrap(), + Some(("query.allow_query_fallback".into(), "true".into())) + ); + assert!(parse_query_option("query.allow_query_fallback", "yes").is_err()); + assert!(parse_query_option("datafusion.optimizer.repartition_joins", "nope").is_err()); + assert!(parse_query_option("datafusion.execution.batch_size", "12").is_err()); + assert_eq!( + parse_query_option("flow.return_region_seq", "true").unwrap(), + None + ); + } +}