feat: validate and snapshot per-query option overrides

Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com>

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<HashMap<String, String>, 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<Self, crate::session_config::Error> {
+        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<PGByteaOutputValue>,
     pg_datestyle_format: ArcSwap<(PGDateTimeStyle, PGDateOrder)>,
     pg_intervalstyle_format: ArcSwap<PGIntervalStyle>,
-    allow_query_fallback: ArcSwap<bool>,
+    query_options: RwLock<HashMap<String, String>>,
 }
 
 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<String, String> {
+        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(

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;
 

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<Option<(String, String)>, 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::<u16>() {
+            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
+        );
+    }
+}