rust/greptimedb
1
0
Fork 0
mirror of https://github.com/GreptimeTeam/greptimedb.git synced 2026-10-12 15:02:32 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat: accept validated query options in SQL SET and SHOW

Browse Source 
Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com>
This commit is contained in:
discord9
2026-10-09 15:47:21 +08:00
parent 1fc2bef02e
commit 08a49d0d06
6 changed files with 250 additions and 28 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/operator/src/statement.rs
+7 -5
View File
@@ -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<Output> {
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)?,
src/operator/src/statement/set.rs
+137 -21
View File
@@ -253,28 +253,57 @@ fn try_parse_datestyle(expr: &Expr) -> Result<(Option<PGDateTimeStyle>, Option<P
}
}
/// Set the allow query fallback configuration parameter to true or false based on the provided expressions.
///
pub fn set_allow_query_fallback(exprs: Vec<Expr>, 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<Expr>, ctx: QueryContextRef) -> Result<bool> {
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::<Vec<_>>()
.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();
src/operator/src/statement/show.rs
+9
View File
@@ -350,6 +350,15 @@ impl StatementExecutor {
#[tracing::instrument(skip_all)]
pub fn show_variable(&self, stmt: ShowVariables, query_ctx: QueryContextRef) -> Result<Output> {
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)
}
src/sql/src/parsers/set_var_parser.rs
+35 -1
View File
@@ -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(
src/sql/src/parsers/show_parser.rs
+29 -1
View File
@@ -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";
tests/cases/standalone/common/system/query_options.sql
+33
View File
@@ -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;