From 08a49d0d06a2b35b718817ead8dff7a9b72653a7 Mon Sep 17 00:00:00 2001 From: discord9 <55937128+discord9@users.noreply.github.com> Date: Fri, 9 Oct 2026 15:45:12 +0800 Subject: [PATCH] feat: accept validated query options in SQL SET and SHOW Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> --- src/operator/src/statement.rs | 12 +- src/operator/src/statement/set.rs | 158 +++++++++++++++--- src/operator/src/statement/show.rs | 9 + src/sql/src/parsers/set_var_parser.rs | 36 +++- src/sql/src/parsers/show_parser.rs | 30 +++- .../common/system/query_options.sql | 33 ++++ 6 files changed, 250 insertions(+), 28 deletions(-) create mode 100644 tests/cases/standalone/common/system/query_options.sql diff --git a/src/operator/src/statement.rs b/src/operator/src/statement.rs index 17d29ab4259..bcf167c91ff 100644 --- a/src/operator/src/statement.rs +++ b/src/operator/src/statement.rs @@ -100,7 +100,7 @@ use crate::error::{ }; use crate::insert::InserterRef; use crate::statement::copy_database::{COPY_DATABASE_TIME_END_KEY, COPY_DATABASE_TIME_START_KEY}; -use crate::statement::set::set_allow_query_fallback; +use crate::statement::set::set_query_option; /// A configurator that customizes or enhances a [`StatementExecutor`]. #[async_trait::async_trait] @@ -527,7 +527,8 @@ impl StatementExecutor { } fn set_variables(&self, set_var: SetVariables, query_ctx: QueryContextRef) -> Result { - let var_name = set_var.variable.to_string().to_uppercase(); + let variable = set_var.variable.to_string(); + let var_name = variable.to_uppercase(); debug!( "Trying to set {}={} for session: {} ", @@ -536,6 +537,10 @@ impl StatementExecutor { query_ctx.conn_info() ); + if set_query_option(&variable, set_var.value.clone(), query_ctx.clone())? { + return Ok(Output::new_with_affected_rows(0)); + } + match var_name.as_str() { "READ_PREFERENCE" => set_read_preference(set_var.value, query_ctx)?, "SKIP_WAL" => set_skip_wal(set_var.value, query_ctx)?, @@ -552,9 +557,6 @@ impl StatementExecutor { "DATESTYLE" => set_datestyle(set_var.value, query_ctx)?, "INTERVALSTYLE" => set_intervalstyle(set_var.value, query_ctx)?, - // Allow query to fallback when failed to push down. - "ALLOW_QUERY_FALLBACK" => set_allow_query_fallback(set_var.value, query_ctx)?, - "CLIENT_ENCODING" => validate_client_encoding(set_var)?, "@@SESSION.MAX_EXECUTION_TIME" | "MAX_EXECUTION_TIME" => match query_ctx.channel() { Channel::Mysql => set_query_timeout(set_var.value, query_ctx)?, diff --git a/src/operator/src/statement/set.rs b/src/operator/src/statement/set.rs index 83924953e75..c5246f065c2 100644 --- a/src/operator/src/statement/set.rs +++ b/src/operator/src/statement/set.rs @@ -253,28 +253,57 @@ fn try_parse_datestyle(expr: &Expr) -> Result<(Option, Option

, ctx: QueryContextRef) -> Result<()> { - let allow_fallback_expr = exprs.first().context(NotSupportedSnafu { - feat: "No allow query fallback value find in set variable statement", - })?; - match allow_fallback_expr { - Expr::Value(ValueWithSpan { - value: Value::Boolean(allow), - span: _, - }) => { - ctx.configuration_parameter() - .set_allow_query_fallback(*allow); - Ok(()) +/// Parses and stores a supported query option, returning false for unrelated SET variables. +pub fn set_query_option(variable: &str, exprs: Vec, ctx: QueryContextRef) -> Result { + let owned = session::query_options::canonical_query_option_name(variable).is_some(); + if !owned + && !variable.to_ascii_lowercase().starts_with("query.") + && !variable.to_ascii_lowercase().starts_with("datafusion.") + { + return Ok(false); + } + + let [Expr::Value(value)] = exprs.as_slice() else { + return InvalidConfigValueSnafu { + source: session::session_config::Error::InvalidConfigValue { + name: variable.to_string(), + value: exprs + .iter() + .map(ToString::to_string) + .collect::>() + .join(", "), + hint: "Expected exactly one literal value".to_string(), + location: snafu::Location::new(file!(), line!(), column!()), + }, } - expr => NotSupportedSnafu { - feat: format!( - "Unsupported allow query fallback expr {} in set variable statement", - expr - ), + .fail(); + }; + let value = match &value.value { + Value::Boolean(value) => value.to_string(), + Value::Number(value, _) + | Value::SingleQuotedString(value) + | Value::DoubleQuotedString(value) => value.clone(), + _ => { + return InvalidConfigValueSnafu { + source: session::session_config::Error::InvalidConfigValue { + name: variable.to_string(), + value: value.value.to_string(), + hint: "Expected a boolean, number, or quoted string literal".to_string(), + location: snafu::Location::new(file!(), line!(), column!()), + }, + } + .fail(); } - .fail(), + }; + if let Some((key, value)) = session::query_options::parse_query_option(variable, &value) + .context(InvalidConfigValueSnafu)? + { + ctx.configuration_parameter() + .set_query_option(&key, &value) + .context(InvalidConfigValueSnafu)?; + Ok(true) + } else { + Ok(false) } } @@ -414,9 +443,96 @@ mod test { use session::context::{Channel, QueryContext}; use sql::ast::{Expr, Value}; - use super::set_skip_wal; + use super::{set_query_option, set_skip_wal}; use crate::statement::set::parse_pg_query_timeout_input; + #[test] + fn test_set_query_option_storage_validation_and_isolation() { + let session = Session::new(None, Channel::Mysql, Default::default(), 0); + let ctx = session.new_query_context(); + let other = Session::new(None, Channel::Mysql, Default::default(), 1); + + assert!( + set_query_option( + "query_parallelism", + vec![Expr::Value(Value::Number("8".into(), false).into())], + ctx.clone(), + ) + .unwrap() + ); + assert_eq!( + ctx.configuration_parameter() + .query_options() + .get("query.parallelism"), + Some(&"8".to_string()) + ); + assert!( + set_query_option( + "QUERY.ALLOW_QUERY_FALLBACK", + vec![Expr::Value(Value::Boolean(false).into())], + ctx.clone(), + ) + .unwrap() + ); + assert_eq!( + ctx.configuration_parameter() + .query_options() + .get("query.allow_query_fallback"), + Some(&"false".to_string()) + ); + assert!( + other + .new_query_context() + .configuration_parameter() + .query_options() + .is_empty() + ); + let mut hinted = ctx.as_ref().clone(); + hinted.set_extension("query_parallelism", "16"); + assert_eq!( + hinted + .effective_query_options() + .unwrap() + .get("query.parallelism"), + Some(&"16".to_string()) + ); + assert_eq!( + ctx.configuration_parameter() + .query_options() + .get("query.parallelism"), + Some(&"8".to_string()) + ); + assert!(!set_query_option("unrelated_client_option", vec![], ctx.clone()).unwrap()); + assert!( + !set_query_option( + "unrelated_client_option", + vec![Expr::Identifier(sql::ast::Ident::new("true"))], + ctx.clone(), + ) + .unwrap() + ); + for (key, values) in [ + ( + "query.unknown", + vec![Expr::Value(Value::Boolean(true).into())], + ), + ( + "datafusion.execution.batch_size", + vec![Expr::Value(Value::Number("1".into(), false).into())], + ), + ( + "query.parallelism", + vec![Expr::Identifier(sql::ast::Ident::new("true"))], + ), + ( + "query.allow_query_fallback", + vec![Expr::Value(Value::Boolean(true).into()); 2], + ), + ] { + assert!(set_query_option(key, values, ctx.clone()).is_err(), "{key}"); + } + } + #[test] fn test_set_skip_wal() { let ctx = QueryContext::arc(); diff --git a/src/operator/src/statement/show.rs b/src/operator/src/statement/show.rs index 4a3fa5ae9b5..b721abc4999 100644 --- a/src/operator/src/statement/show.rs +++ b/src/operator/src/statement/show.rs @@ -350,6 +350,15 @@ impl StatementExecutor { #[tracing::instrument(skip_all)] pub fn show_variable(&self, stmt: ShowVariables, query_ctx: QueryContextRef) -> Result { + if let Some(output) = query::sql::show_query_option( + stmt.clone(), + query_ctx.clone(), + self.query_engine.as_ref(), + ) + .context(error::ExecuteStatementSnafu)? + { + return Ok(output); + } query::sql::show_variable(stmt, query_ctx).context(error::ExecuteStatementSnafu) } diff --git a/src/sql/src/parsers/set_var_parser.rs b/src/sql/src/parsers/set_var_parser.rs index 19f0c48278d..b80e30ad3f7 100644 --- a/src/sql/src/parsers/set_var_parser.rs +++ b/src/sql/src/parsers/set_var_parser.rs @@ -60,7 +60,7 @@ mod tests { use sqlparser::ast::{Expr, Ident, ObjectName, Value}; use super::*; - use crate::dialect::GreptimeDbDialect; + use crate::dialect::{GenericDialect, GreptimeDbDialect, MySqlDialect, PostgreSqlDialect}; use crate::parser::ParseOptions; fn assert_mysql_parse_result(sql: &str, indent_str: &str, expr: Expr) { @@ -144,6 +144,40 @@ mod tests { assert_pg_parse_result(sql, "STATEMENT_TIMEOUT", expected_query_timeout_expr); } + #[test] + fn test_set_query_option_dialects_and_case() { + let cases = [ + ( + "SET query_parallelism = 4", + &MySqlDialect {} as &dyn sqlparser::dialect::Dialect, + "query_parallelism", + "4", + ), + ( + "SET QUERY.ALLOW_QUERY_FALLBACK TO false", + &PostgreSqlDialect {}, + "QUERY.ALLOW_QUERY_FALLBACK", + "false", + ), + ( + "SET datafusion.optimizer.repartition_joins = true", + &GenericDialect {}, + "datafusion.optimizer.repartition_joins", + "true", + ), + ]; + for (sql, dialect, variable, value) in cases { + let result = + ParserContext::create_with_dialect(sql, dialect, ParseOptions::default()).unwrap(); + let [Statement::SetVariables(set)] = result.as_slice() else { + panic!("unexpected SET parse result for {sql}"); + }; + assert_eq!(set.variable.to_string(), variable); + assert_eq!(set.value.len(), 1); + assert_eq!(set.value[0].to_string(), value, "{sql}"); + } + } + #[test] fn test_unsupported_set_variant_remains_rejected() { let result = ParserContext::create_with_dialect( diff --git a/src/sql/src/parsers/show_parser.rs b/src/sql/src/parsers/show_parser.rs index a039dceea63..8ccf92887e6 100644 --- a/src/sql/src/parsers/show_parser.rs +++ b/src/sql/src/parsers/show_parser.rs @@ -632,7 +632,7 @@ mod tests { use sqlparser::ast::{Ident, ObjectName}; use super::*; - use crate::dialect::GreptimeDbDialect; + use crate::dialect::{GenericDialect, GreptimeDbDialect, MySqlDialect, PostgreSqlDialect}; use crate::parser::ParseOptions; use crate::statements::show::ShowDatabases; #[cfg(feature = "enterprise")] @@ -899,6 +899,34 @@ mod tests { ); } + #[test] + fn test_show_dotted_query_options() { + for (sql, dialect, expected) in [ + ( + "SHOW VARIABLES query.parallelism", + &MySqlDialect {} as &dyn sqlparser::dialect::Dialect, + "query.parallelism", + ), + ( + "SHOW QUERY.ALLOW_QUERY_FALLBACK", + &PostgreSqlDialect {}, + "QUERY.ALLOW_QUERY_FALLBACK", + ), + ( + "SHOW VARIABLES datafusion.optimizer.repartition_joins", + &GenericDialect {}, + "datafusion.optimizer.repartition_joins", + ), + ] { + let statements = + ParserContext::create_with_dialect(sql, dialect, ParseOptions::default()).unwrap(); + let [Statement::ShowVariables(show)] = statements.as_slice() else { + panic!("unexpected SHOW parse result for {sql}"); + }; + assert_eq!(show.variable.to_string(), expected); + } + } + #[test] pub fn test_show_variables() { let sql = "SHOW VARIABLES system_time_zone"; diff --git a/tests/cases/standalone/common/system/query_options.sql b/tests/cases/standalone/common/system/query_options.sql new file mode 100644 index 00000000000..f24a85f3ad9 --- /dev/null +++ b/tests/cases/standalone/common/system/query_options.sql @@ -0,0 +1,33 @@ +-- Query options are session-scoped; request hints take precedence per query. +-- SQLNESS PROTOCOL MYSQL +SET query_parallelism = 8; + +-- SQLNESS PROTOCOL MYSQL +SHOW VARIABLES query.parallelism; + +-- SQLNESS PROTOCOL MYSQL +SET query.allow_query_fallback = false; + +-- SQLNESS PROTOCOL MYSQL +SHOW VARIABLES query.allow_query_fallback; + +-- Query namespace typos and unapproved DataFusion options are errors. +-- SQLNESS PROTOCOL MYSQL +SET query.unknown_option = true; + +-- SQLNESS PROTOCOL MYSQL +SET datafusion.execution.batch_size = 128; + +-- Computed expressions are not accepted for query options. +-- SQLNESS PROTOCOL MYSQL +SET query.parallelism = 4 + 4; + +-- Third-party startup variables remain compatible. +-- SQLNESS PROTOCOL MYSQL +SET autocommit = 1; + +-- SQLNESS PROTOCOL POSTGRES +SET QUERY.PARALLELISM TO 16; + +-- SQLNESS PROTOCOL POSTGRES +SHOW QUERY.PARALLELISM;