rust/greptimedb
1
0
Fork 0
mirror of https://github.com/GreptimeTeam/greptimedb.git synced 2026-10-28 06:36:18 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat: unify query options across SQL, HTTP and gRPC (#9494)

Browse Source 
Isolated rebase verification of PR #9494. Preserve the reviewed net change; original PR branch remains untouched.

Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com>
This commit is contained in:
discord9
2026-10-16 04:48:32 +00:00
committed by GitHub
parent f663569cf0
commit 778f4556a8
30 changed files with 855 additions and 180 deletions
Download Patch File Download Diff File Expand all files Collapse all files
docs/how-to/query-options.md
+18
View File
@@ -0,0 +1,18 @@
# Query options
Query options adjust selected behavior for a connection or request; they are not server-wide settings or resource-pool controls.
| Canonical option | SQL aliases | Values |
| --- | --- | --- |
| `query.parallelism` | `query_parallelism`, `query.parallelism` | Integer `1`–`1024`; server default `0` means automatic CPU-based parallelism. This is not a resource-pool limit. |
| `query.allow_query_fallback` | `query_fallback`, `allow_query_fallback`, `query.allow_query_fallback` (legacy `ALLOW_QUERY_FALLBACK` syntax / `query_fallback` hint also work) | `true` or `false` |
| `query.enable_remote_dynamic_filter_pushdown` | canonical name | `true` or `false` |
| `datafusion.optimizer.*` | canonical name | Keys, types, defaults, and validation come from the linked DataFusion version (for example, numeric `max_passes` and Boolean `filter_null_join_keys`). |
Use `SET query_parallelism = 4` or `SET query.parallelism TO 4`; `SHOW VARIABLES query.parallelism` works in MySQL and PostgreSQL, which also accepts `SHOW query.parallelism`. `SHOW VARIABLES <option>` reports the effective value. Unspecified options inherit engine settings; defaults are automatic parallelism, fallback `false`, and remote dynamic-filter pushdown `true`. Invalid values and unsupported names are errors; unknown third-party MySQL/PostgreSQL initialization settings retain existing tolerance.
HTTP `X-Greptime-Hints` (for example, `query.parallelism=4`) and gRPC `x-greptime-hints` request metadata set request options; gRPC hints are not a protobuf `Hints` field. The Rust Flight client offers `Database::sql_with_hint` and `flight_request().with_hints`. Hints override session `SET`, including within a multi-statement HTTP request. SQL `SET` lasts for its connection only; `SET LOCAL` transaction-scoped restoration is not implemented.
`query.enable_remote_dynamic_filter_pushdown` controls propagation across frontend-to-datanode remote queries; it is distinct from DataFusion's native `datafusion.optimizer.enable_dynamic_filter_pushdown`.
Internal flow/query-reserved keys are not public query options. HTTP response timeout `x-greptime-timeout` defaults to `0` (no deadline); this is a response deadline, not a per-statement or stream/body-processing limit.
src/datanode/src/region_server.rs
+12 -6
View File
@@ -308,6 +308,17 @@ impl RegionServer {
};
let region_id = RegionId::from_u64(request.region_id);
let engine_ctx = self
.inner
.query_engine
.engine_context(query_ctx)
.map_err(|error| {
crate::error::InvalidQueryContextSnafu {
reason: error.to_string(),
}
.build()
})?;
let query_ctx = engine_ctx.query_ctx();
let catalog_list = Arc::new(NameAwareCatalogList::new(
self.clone(),
region_id,
@@ -318,12 +329,7 @@ impl RegionServer {
common_telemetry::info!("Handle remote read for region: {}", region_id);
}
let decoder = self
.inner
.query_engine
.engine_context(query_ctx.clone())
.new_plan_decoder()
.context(NewPlanDecoderSnafu)?;
let decoder = engine_ctx.new_plan_decoder().context(NewPlanDecoderSnafu)?;
let plan = decoder
.decode(Bytes::from(request.plan), catalog_list, false)
src/datanode/src/tests.rs
+4 -1
View File
@@ -90,7 +90,10 @@ impl QueryEngine for MockQueryEngine {
unimplemented!()
}
fn engine_context(&self, _query_ctx: QueryContextRef) -> QueryEngineContext {
fn engine_context(
&self,
_query_ctx: QueryContextRef,
) -> query::error::Result<QueryEngineContext> {
unimplemented!()
}
fn engine_state(&self) -> &QueryEngineState {
src/flow/src/batching_mode/frontend_client.rs
+33 -20
View File
@@ -590,10 +590,10 @@ impl FrontendClient {
query, batch_opts, ..
} => {
let query_parallelism = query.parallelism.to_string();
let hints = vec![
(QUERY_PARALLELISM_HINT, query_parallelism.as_str()),
(READ_PREFERENCE_HINT, batch_opts.read_preference.as_ref()),
];
let mut hints = vec![(READ_PREFERENCE_HINT, batch_opts.read_preference.as_ref())];
if query.parallelism > 0 {
hints.push((QUERY_PARALLELISM_HINT, query_parallelism.as_str()));
}
let db = self.get_random_active_frontend(catalog, schema).await?;
*peer_desc = Some(PeerDesc::Dist {
peer: db.peer.clone(),
@@ -615,10 +615,13 @@ impl FrontendClient {
..
} => {
*peer_desc = Some(PeerDesc::Standalone);
let mut extensions_map = HashMap::from([(
QUERY_PARALLELISM_HINT.to_string(),
query.parallelism.to_string(),
)]);
let mut extensions_map = HashMap::new();
if query.parallelism > 0 {
extensions_map.insert(
QUERY_PARALLELISM_HINT.to_string(),
query.parallelism.to_string(),
);
}
for (key, value) in extensions {
extensions_map.insert((*key).to_string(), (*value).to_string());
}
@@ -697,10 +700,14 @@ impl FrontendClient {
.current_catalog(catalog.to_string())
.current_schema(schema.to_string())
.channel(Channel::Internal)
.extensions(HashMap::from([(
QUERY_PARALLELISM_HINT.to_string(),
query.parallelism.to_string(),
)]))
.extensions(if query.parallelism > 0 {
HashMap::from([(
QUERY_PARALLELISM_HINT.to_string(),
query.parallelism.to_string(),
)])
} else {
HashMap::new()
})
.build();
let ctx = Arc::new(ctx);
let database_client = {
@@ -756,14 +763,16 @@ impl FrontendClient {
peer: db.peer.clone(),
});
let query_parallelism = query.parallelism.to_string();
let mut hints = vec![(READ_PREFERENCE_HINT, batch_opts.read_preference.as_ref())];
if query.parallelism > 0 {
hints.push((QUERY_PARALLELISM_HINT, query_parallelism.as_str()));
}
db.database
.handle_with_retry(
req.clone(),
batch_opts.experimental_grpc_max_retries,
&[
(QUERY_PARALLELISM_HINT, &query.parallelism.to_string()),
(READ_PREFERENCE_HINT, batch_opts.read_preference.as_ref()),
],
&hints,
)
.await
.with_context(|_| InvalidRequestSnafu {
@@ -778,10 +787,14 @@ impl FrontendClient {
let ctx = QueryContextBuilder::default()
.current_catalog(catalog.to_string())
.current_schema(schema.to_string())
.extensions(HashMap::from([(
QUERY_PARALLELISM_HINT.to_string(),
query.parallelism.to_string(),
)]))
.extensions(if query.parallelism > 0 {
HashMap::from([(
QUERY_PARALLELISM_HINT.to_string(),
query.parallelism.to_string(),
)])
} else {
HashMap::new()
})
.build();
let ctx = Arc::new(ctx);
{
src/frontend/src/instance.rs
+15
View File
@@ -307,6 +307,21 @@ impl Instance {
async fn query_statement(&self, stmt: Statement, query_ctx: QueryContextRef) -> Result<Output> {
check_permission(self.plugins.clone(), &stmt, &query_ctx)?;
let query_ctx = if matches!(&stmt, Statement::SetVariables(_)) {
query_ctx
} else {
Arc::new(
query_ctx
.query_option_snapshot()
.map_err(|error| {
query::error::InvalidQueryContextExtensionSnafu {
reason: error.to_string(),
}
.build()
})
.context(PlanStatementSnafu)?,
)
};
let query_interceptor = self.plugins.get::<SqlQueryInterceptorRef<Error>>();
let query_interceptor = query_interceptor.as_ref();
src/frontend/src/instance/grpc.rs
+12 -4
View File
@@ -141,9 +141,12 @@ impl GrpcQueryHandler for Instance {
let timer = GRPC_HANDLE_PLAN_ELAPSED.start_timer();
// use dummy catalog to provide table
let plan_decoder = self
let engine_ctx = self
.query_engine()
.engine_context(ctx.clone())
.context(PlanStatementSnafu)?;
let ctx = engine_ctx.query_ctx();
let plan_decoder = engine_ctx
.new_plan_decoder()
.context(PlanStatementSnafu)?;
@@ -502,11 +505,12 @@ impl Instance {
})?;
// use dummy catalog to provide table
let plan_decoder = self
let engine_ctx = self
.query_engine()
.engine_context(ctx.clone())
.new_plan_decoder()
.context(PlanStatementSnafu)?;
let ctx = engine_ctx.query_ctx();
let plan_decoder = engine_ctx.new_plan_decoder().context(PlanStatementSnafu)?;
let dummy_catalog_list = Arc::new(
catalog::table_source::dummy_catalog::DummyCatalogList::new_with_query_ctx(
@@ -549,7 +553,11 @@ impl Instance {
let insert_into = add_insert_to_logical_plan(table_name, table_source, logical_plan)
.context(SubstraitDecodeLogicalPlanSnafu)?;
let engine_ctx = self.query_engine().engine_context(ctx.clone());
let engine_ctx = self
.query_engine()
.engine_context(ctx.clone())
.context(PlanStatementSnafu)?;
let ctx = engine_ctx.query_ctx();
let state = engine_ctx.state();
// Analyze the plan
let analyzed_plan = state
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 = sql::ast::object_name_to_string_unquoted(&set_var.variable);
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
+43 -23
View File
@@ -22,7 +22,7 @@ use session::ReadPreference;
use session::context::Channel::Postgres;
use session::context::QueryContextRef;
use session::session_config::{PGByteaOutputValue, PGDateOrder, PGDateTimeStyle, PGIntervalStyle};
use snafu::{OptionExt, ResultExt, ensure};
use snafu::{IntoError, OptionExt, ResultExt, ensure};
use sql::ast::{Expr, Ident, Value};
use sql::statements::set_variables::SetVariables;
use sqlparser::ast::ValueWithSpan;
@@ -253,29 +253,49 @@ 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(())
}
expr => NotSupportedSnafu {
feat: format!(
"Unsupported allow query fallback expr {} in set variable statement",
expr
),
}
.fail(),
/// Applies supported query options and rejects unsupported owned namespaces.
pub fn set_query_option(variable: &str, exprs: Vec<Expr>, ctx: QueryContextRef) -> Result<bool> {
let lower = variable.to_ascii_lowercase();
if !lower.starts_with("query.")
&& !lower.starts_with("datafusion.")
&& session::query_options::canonical_query_option_name(variable).is_none()
{
return Ok(false);
}
let [Expr::Value(value)] = exprs.as_slice() else {
return Err(InvalidConfigValueSnafu.into_error(
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!()),
},
));
};
let value = match &value.value {
Value::Boolean(value) => value.to_string(),
Value::Number(value, _)
| Value::SingleQuotedString(value)
| Value::DoubleQuotedString(value) => value.clone(),
_ => {
return Err(InvalidConfigValueSnafu.into_error(
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!()),
},
));
}
};
ctx.configuration_parameter()
.set_query_option(variable, &value)
.context(InvalidConfigValueSnafu)?;
Ok(true)
}
pub fn set_intervalstyle(exprs: Vec<Expr>, ctx: QueryContextRef) -> Result<()> {
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/query/src/datafusion.rs
+8 -48
View File
@@ -55,9 +55,7 @@ use tracing::Span;
use crate::analyze::DistAnalyzeExec;
pub use crate::datafusion::planner::DfContextProviderAdapter;
use crate::dist_plan::{
DistPlannerOptions, MergeScanLogicalPlan, RemoteDynFilterReceiverInjectorRef,
};
use crate::dist_plan::{MergeScanLogicalPlan, RemoteDynFilterReceiverInjectorRef};
use crate::error::{
CatalogSnafu, CreateRecordBatchSnafu, MissingTableMutationHandlerSnafu,
MissingTimestampColumnSnafu, QueryExecutionSnafu, Result, TableMutationSnafu,
@@ -156,7 +154,8 @@ impl DatafusionQueryEngine {
plan: LogicalPlan,
query_ctx: QueryContextRef,
) -> Result<Output> {
let mut ctx = self.engine_context(query_ctx.clone());
let mut ctx = self.engine_context(query_ctx)?;
let query_ctx = ctx.query_ctx();
let plan = if let Some(receiver_injector) =
self.plugins.get::<RemoteDynFilterReceiverInjectorRef>()
{
@@ -552,54 +551,15 @@ impl QueryEngine for DatafusionQueryEngine {
self.state.read_table(table).map_err(Into::into)
}
fn engine_context(&self, query_ctx: QueryContextRef) -> QueryEngineContext {
fn engine_context(&self, query_ctx: QueryContextRef) -> Result<QueryEngineContext> {
let mut state = self.state.session_state();
let query_ctx = self.state.apply_query_options(&mut state, query_ctx)?;
state.config_mut().set_extension(query_ctx.clone());
state.config_mut().set_extension(self.state.clone());
// note that hints in "x-greptime-hints" is automatically parsed
// and set to query context's extension, so we can get it from query context.
if let Some(parallelism) = query_ctx.extension(QUERY_PARALLELISM_HINT) {
if let Ok(n) = parallelism.parse::<u64>() {
if n > 0 {
let new_cfg = state.config().clone().with_target_partitions(n as usize);
*state.config_mut() = new_cfg;
}
} else {
common_telemetry::warn!(
"Failed to parse query_parallelism: {}, using default value",
parallelism
);
}
}
// configure execution options
state.config_mut().options_mut().execution.time_zone =
Some(query_ctx.timezone().to_string());
// usually it's impossible to have both `set variable` set by sql client and
// hint in header by grpc client, so only need to deal with them separately
if query_ctx.configuration_parameter().allow_query_fallback() {
state
.config_mut()
.options_mut()
.extensions
.insert(DistPlannerOptions {
allow_query_fallback: true,
});
} else if let Some(fallback) = query_ctx.extension(QUERY_FALLBACK_HINT) {
// also check the query context for fallback hint
// if it is set, we will enable the fallback
if fallback.to_lowercase().parse::<bool>().unwrap_or(false) {
state
.config_mut()
.options_mut()
.extensions
.insert(DistPlannerOptions {
allow_query_fallback: true,
});
}
}
state
.config_mut()
.options_mut()
@@ -639,7 +599,7 @@ impl QueryEngine for DatafusionQueryEngine {
}
}
QueryEngineContext::new(state, query_ctx)
Ok(QueryEngineContext::new(state, query_ctx))
}
fn engine_state(&self) -> &QueryEngineState {
@@ -1194,7 +1154,7 @@ mod tests {
.as_any()
.downcast_ref::<DatafusionQueryEngine>()
.unwrap();
let engine_ctx = engine.engine_context(QueryContext::arc());
let engine_ctx = engine.engine_context(QueryContext::arc()).unwrap();
let state = engine_ctx.state();
let schema = Arc::new(datatypes::schema::Schema::new(vec![ColumnSchema::new(
@@ -1254,7 +1214,7 @@ mod tests {
.as_any()
.downcast_ref::<DatafusionQueryEngine>()
.unwrap();
let engine_ctx = engine.engine_context(QueryContext::arc());
let engine_ctx = engine.engine_context(QueryContext::arc()).unwrap();
let state = engine_ctx.state();
assert!(
src/query/src/dist_plan/analyzer.rs
+14 -2
View File
@@ -15,6 +15,7 @@
use std::collections::{BTreeMap, BTreeSet, HashSet};
use std::sync::Arc;
use common_function::function::FunctionContext;
use common_telemetry::debug;
use datafusion::config::{ConfigExtension, ExtensionOptions};
use datafusion::datasource::DefaultTableSource;
@@ -122,8 +123,19 @@ impl AnalyzerRule for DistPlannerAnalyzer {
) -> datafusion_common::Result<LogicalPlan> {
expr_depth::validate_plan(&plan)?;
let mut config = config.clone();
// Aligned with the behavior in `datafusion_optimizer::OptimizerContext::new()`.
config.optimizer.filter_null_join_keys = true;
// Preserve the historical pre-MergeScan default unless explicitly set.
let explicit = config
.extensions
.get::<FunctionContext>()
.is_some_and(|context| {
context
.query_ctx
.extension("datafusion.optimizer.filter_null_join_keys")
.is_some()
});
if !explicit {
config.optimizer.filter_null_join_keys = true;
}
let config = Arc::new(config);
let opt = config.extensions.get::<DistPlannerOptions>();
let allow_fallback = opt.map(|o| o.allow_query_fallback).unwrap_or(false);
src/query/src/optimizer/constant_term.rs
+1 -1
View File
@@ -470,7 +470,7 @@ mod tests {
.await
.unwrap();
let engine_ctx = engine.engine_context(query_ctx);
let engine_ctx = engine.engine_context(query_ctx).unwrap();
let state = engine_ctx.state();
let analyzed_plan = state
src/query/src/planner.rs
+46 -28
View File
@@ -56,6 +56,7 @@ use crate::error::{
};
use crate::log_query::planner::LogQueryPlanner;
use crate::optimizer::expr_depth;
use crate::options::ScheduledTimeExtension;
use crate::parser::{DEFAULT_LOOKBACK_STRING, PromQuery, QueryLanguageParser, QueryStatement};
use crate::promql::planner::PromPlanner;
use crate::query_engine::{DefaultPlanDecoder, QueryEngineState};
@@ -96,24 +97,40 @@ impl DfLogicalPlanner {
/// in the query context.
fn derive_session_state_with_scheduled_time(
&self,
query_ctx: &QueryContextRef,
) -> Result<SessionState> {
query_ctx: QueryContextRef,
) -> Result<QueryEngineContext> {
let extensions = query_ctx.extensions();
match crate::options::parse_scheduled_time_datetime(&extensions)? {
let mut state = match crate::options::parse_scheduled_time_datetime(&extensions)? {
Some(dt) => {
let execution_props = self
.session_state
.execution_props()
.clone()
.with_query_execution_start_time(dt);
Ok(
SessionStateBuilder::new_from_existing(self.session_state.clone())
.with_execution_props(execution_props)
.build(),
)
SessionStateBuilder::new_from_existing(self.session_state.clone())
.with_execution_props(execution_props)
.build()
}
None => Ok(self.session_state.clone()),
}
None => self.session_state.clone(),
};
let query_ctx = self
.engine_state
.apply_query_options(&mut state, query_ctx)?;
state.config_mut().set_extension(query_ctx.clone());
state.config_mut().set_extension(self.engine_state.clone());
state
.config_mut()
.options_mut()
.extensions
.insert(ScheduledTimeExtension {
scheduled_time: crate::options::scheduled_time_from_ctx(&query_ctx),
});
let config_options = state.config_options().clone();
let _ = state
.execution_props_mut()
.config_options
.insert(config_options);
Ok(QueryEngineContext::new(state, query_ctx))
}
/// Basically the same with `explain_to_plan` in DataFusion, but adapted to Greptime's
@@ -221,7 +238,9 @@ impl DfLogicalPlanner {
.fail()?;
}
let scheduled_state = self.derive_session_state_with_scheduled_time(&query_ctx)?;
let engine_context = self.derive_session_state_with_scheduled_time(query_ctx)?;
let scheduled_state = engine_context.state().clone();
let query_ctx = engine_context.query_ctx();
let table_provider = DfTableSourceProvider::new(
self.engine_state.catalog_manager().clone(),
self.engine_state.disallow_cross_catalog_query(),
@@ -244,7 +263,7 @@ impl DfLogicalPlanner {
)
.await?;
let config_options = self.session_state.config().options();
let config_options = scheduled_state.config().options();
let parser_options = &config_options.sql_parser;
let parser_options = ParserOptions {
map_string_types_to_utf8view: false,
@@ -275,10 +294,9 @@ impl DfLogicalPlanner {
.await?;
// Optimize logical plan by extension rules
let context = QueryEngineContext::new(scheduled_state, query_ctx);
let plan = self
.engine_state
.optimize_by_extension_rules(plan, &context)?;
.optimize_by_extension_rules(plan, &engine_context)?;
common_telemetry::debug!("Logical planner, optimize result: {plan}");
Ok(plan)
@@ -293,16 +311,17 @@ impl DfLogicalPlanner {
normalize_ident: bool,
query_ctx: QueryContextRef,
) -> Result<DfExpr> {
let scheduled_state = self.derive_session_state_with_scheduled_time(&query_ctx)?;
let engine_context = self.derive_session_state_with_scheduled_time(query_ctx)?;
let query_ctx = engine_context.query_ctx();
let context_provider = DfContextProviderAdapter::try_new(
self.engine_state.clone(),
scheduled_state,
engine_context.state().clone(),
None,
query_ctx,
)
.await?;
let config_options = self.session_state.config().options();
let config_options = engine_context.state().config().options();
let parser_options = &config_options.sql_parser;
let parser_options: ParserOptions = ParserOptions {
map_string_types_to_utf8view: false,
@@ -317,7 +336,9 @@ impl DfLogicalPlanner {
#[tracing::instrument(skip_all)]
async fn plan_pql(&self, stmt: &EvalStmt, query_ctx: QueryContextRef) -> Result<LogicalPlan> {
let mut scheduled_state = self.derive_session_state_with_scheduled_time(&query_ctx)?;
let engine_context = self.derive_session_state_with_scheduled_time(query_ctx)?;
let query_ctx = engine_context.query_ctx();
let mut scheduled_state = engine_context.state().clone();
let promql_annotations = query_ctx.remote_query_id().map(promql_annotation_collector);
if let Some(collector) = &promql_annotations {
scheduled_state
@@ -708,22 +729,19 @@ impl LogicalPlanner for DfLogicalPlanner {
query: LogQuery,
query_ctx: QueryContextRef,
) -> Result<LogicalPlan> {
let plan_decoder = Arc::new(DefaultPlanDecoder::new(
self.session_state.clone(),
&query_ctx,
)?);
let engine_context = self.derive_session_state_with_scheduled_time(query_ctx)?;
let query_ctx = engine_context.query_ctx();
let state = engine_context.state().clone();
let plan_decoder = Arc::new(DefaultPlanDecoder::new(state.clone(), &query_ctx)?);
let table_provider = DfTableSourceProvider::new(
self.engine_state.catalog_manager().clone(),
self.engine_state.disallow_cross_catalog_query(),
query_ctx,
plan_decoder,
self.session_state
.config_options()
.sql_parser
.enable_ident_normalization,
state.config_options().sql_parser.enable_ident_normalization,
);
let mut planner = LogQueryPlanner::new(table_provider, self.session_state.clone());
let mut planner = LogQueryPlanner::new(table_provider, state);
planner
.query_to_plan(query)
.await
@@ -964,7 +982,7 @@ mod tests {
.plan(&stmt, query_ctx.clone())
.await
.unwrap();
let context = engine.engine_context(query_ctx.clone());
let context = engine.engine_context(query_ctx.clone()).unwrap();
let state = context.state();
state
.analyzer()
src/query/src/query_engine.rs
+1 -1
View File
@@ -97,7 +97,7 @@ pub trait QueryEngine: Send + Sync {
fn read_table(&self, table: TableRef) -> Result<DataFrame>;
/// Create a [`QueryEngineContext`].
fn engine_context(&self, query_ctx: QueryContextRef) -> QueryEngineContext;
fn engine_context(&self, query_ctx: QueryContextRef) -> Result<QueryEngineContext>;
/// Retrieve the query engine state [`QueryEngineState`]
fn engine_state(&self) -> &QueryEngineState;
src/query/src/query_engine/default_serializer.rs
+15 -2
View File
@@ -433,6 +433,7 @@ mod tests {
let plan_decoder = engine
.engine_context(QueryContext::arc())
.unwrap()
.new_plan_decoder()
.unwrap();
let catalog_list = Arc::new(DummyCatalogList::with_table_provider(table_provider));
@@ -525,6 +526,7 @@ mod tests {
let table_provider = Arc::new(MemTable::try_new(schema, vec![vec![]]).unwrap());
let plan_decoder = engine
.engine_context(QueryContext::arc())
.unwrap()
.new_plan_decoder()
.unwrap();
@@ -831,6 +833,7 @@ mod tests {
.unwrap();
let plan_decoder = engine
.engine_context(query_ctx.clone())
.unwrap()
.new_plan_decoder()
.unwrap();
let decoded = plan_decoder
@@ -898,6 +901,7 @@ mod tests {
.unwrap();
let plan_decoder = engine
.engine_context(query_ctx.clone())
.unwrap()
.new_plan_decoder()
.unwrap();
let catalog_list = Arc::new(DummyCatalogList::with_table_provider(Arc::new(
@@ -964,6 +968,7 @@ mod tests {
fn plan_decoder(engine: &dyn crate::QueryEngine) -> SubstraitPlanDecoderRef {
engine
.engine_context(QueryContext::arc())
.unwrap()
.new_plan_decoder()
.unwrap()
}
@@ -997,7 +1002,11 @@ mod tests {
QueryOptions::default(),
)
.query_engine();
let mut state = engine.engine_context(resolver_ctx.clone()).state().clone();
let mut state = engine
.engine_context(resolver_ctx.clone())
.unwrap()
.state()
.clone();
state.config_mut().set_extension(resolver_ctx);
let extensions = Extensions::default();
let consumer = MergeScanSubstraitConsumer {
@@ -1470,7 +1479,11 @@ mod tests {
)
.query_engine();
let mut state = SessionStateBuilder::new_from_existing(
engine.engine_context(QueryContext::arc()).state().clone(),
engine
.engine_context(QueryContext::arc())
.unwrap()
.state()
.clone(),
)
.with_serializer_registry(Arc::new(DefaultSerializer))
.build();
src/query/src/query_engine/state.rs
+57 -3
View File
@@ -59,6 +59,7 @@ use crate::dist_plan::{
MergeSortExtensionPlanner, RemoteDynFilterReceiverExtensionPlanner,
RemoteDynFilterRegistryLease,
};
use crate::error::InvalidQueryContextExtensionSnafu;
use crate::metrics::{QUERY_MEMORY_POOL_REJECTED_TOTAL, QUERY_MEMORY_POOL_USAGE_BYTES};
use crate::optimizer::ExtensionAnalyzerRule;
use crate::optimizer::const_normalization::ConstNormalizationRule;
@@ -321,7 +322,61 @@ impl QueryEngineState {
rules.retain(|rule| rule.name() != name);
}
/// Optimize the logical plan by the extension analyzer rules.
/// Applies validated query options to one per-query DataFusion configuration.
pub fn apply_query_options(
&self,
state: &mut SessionState,
query_ctx: QueryContextRef,
) -> crate::error::Result<QueryContextRef> {
let query_ctx = Arc::new(query_ctx.query_option_snapshot().map_err(|error| {
InvalidQueryContextExtensionSnafu {
reason: error.to_string(),
}
.build()
})?);
let options = query_ctx.extensions();
const MASTER: &str = "datafusion.optimizer.enable_dynamic_filter_pushdown";
// The master switch cascades to children, so apply explicit child overrides afterward.
for (key, value) in options
.get_key_value(MASTER)
.into_iter()
.chain(options.iter().filter(|(key, _)| key.as_str() != MASTER))
{
if key.starts_with("datafusion.") {
state
.config_mut()
.options_mut()
.set(key, value)
.map_err(|error| {
InvalidQueryContextExtensionSnafu {
reason: error.to_string(),
}
.build()
})?;
}
}
if let Some(partitions) = options.get("query.parallelism") {
let partitions = partitions.parse::<usize>().map_err(|error| {
InvalidQueryContextExtensionSnafu {
reason: error.to_string(),
}
.build()
})?;
*state.config_mut() = state.config().clone().with_target_partitions(partitions);
}
if let Some(fallback) = options.get("query.allow_query_fallback") {
state
.config_mut()
.options_mut()
.extensions
.insert(DistPlannerOptions {
allow_query_fallback: fallback == "true",
});
}
Ok(query_ctx)
}
/// Optimize the logical plan using registered extension analyzer rules.
pub fn optimize_by_extension_rules(
&self,
plan: DfLogicalPlan,
@@ -331,8 +386,7 @@ impl QueryEngineState {
self.extension_rules
.iter()
.try_fold(plan, |acc_plan, rule| {
let plan =
rule.analyze(acc_plan, context, self.session_state().config_options())?;
let plan = rule.analyze(acc_plan, context, context.state().config_options())?;
expr_depth::validate_plan(&plan)?;
Ok(plan)
})
src/query/src/sql.rs
+86 -1
View File
@@ -376,7 +376,7 @@ async fn query_from_information_schema_dataframe(
let view = dataframe.into_view();
let dataframe = SessionContext::new_with_state(
query_engine
.engine_context(query_ctx.clone())
.engine_context(query_ctx.clone())?
.state()
.clone(),
)
@@ -844,6 +844,91 @@ pub async fn show_charsets_dataframe(
.await
}
/// Shows one supported query optimizer option, returning `None` for legacy SHOW handling.
pub fn show_query_option(
stmt: ShowVariables,
query_ctx: QueryContextRef,
engine: &dyn crate::QueryEngine,
) -> Result<Option<Output>> {
let name = sql::ast::object_name_to_string_unquoted(&stmt.variable);
let Some(canonical) = session::query_options::canonical_query_option_name(&name) else {
if name.to_ascii_lowercase().starts_with("query.")
|| name.to_ascii_lowercase().starts_with("datafusion.")
{
session::query_options::parse_query_option(
&name,
"",
session::query_options::MAX_PUBLIC_PARALLELISM,
)
.map_err(|error| {
error::InvalidQueryContextExtensionSnafu {
reason: error.to_string(),
}
.build()
})?;
}
return Ok(None);
};
let context = engine.engine_context(query_ctx.clone())?;
let value = match canonical.as_str() {
"query.allow_query_fallback" => context
.state()
.config_options()
.extensions
.get::<crate::dist_plan::DistPlannerOptions>()
.map(|options| options.allow_query_fallback)
.unwrap_or(false)
.to_string(),
"query.enable_remote_dynamic_filter_pushdown" => {
crate::options::remote_dyn_filter_pushdown_enabled_from_extensions(
&context.query_ctx().extensions(),
)?
.to_string()
}
"query.parallelism" => context.state().config().target_partitions().to_string(),
"datafusion.optimizer.filter_null_join_keys"
if context
.state()
.analyzer()
.rules
.iter()
.any(|rule| rule.name() == "DistPlannerAnalyzer")
&& context
.query_ctx()
.extension("datafusion.optimizer.filter_null_join_keys")
.is_none() =>
{
"true".to_string()
}
_ => context
.state()
.config_options()
.entries()
.into_iter()
.find(|entry| entry.key == canonical)
.and_then(|entry| entry.value)
.ok_or_else(|| {
error::InvalidQueryContextExtensionSnafu {
reason: format!(
"No DataFusion configuration entry for supported option {canonical}"
),
}
.build()
})?,
};
let schema = Arc::new(Schema::new(vec![ColumnSchema::new(
name.clone(),
ConcreteDataType::string_datatype(),
false,
)]));
let records = RecordBatches::try_from_columns(
schema,
vec![Arc::new(StringVector::from(vec![value])) as _],
)
.context(error::CreateRecordBatchSnafu)?;
Ok(Some(Output::new_with_record_batches(records)))
}
pub fn show_variable(stmt: ShowVariables, query_ctx: QueryContextRef) -> Result<Output> {
let variable = stmt.variable.to_string().to_uppercase();
let value = match variable.as_str() {
src/servers/src/grpc/flight.rs
+1
View File
@@ -198,6 +198,7 @@ impl FlightCraft for GreptimeRequestHandler {
request: Request<Ticket>,
) -> TonicResult<Response<TonicStream<FlightData>>> {
let mut hints = hint_headers::extract_hints(request.metadata());
hints = hint_headers::validate_public_hints(hints)?;
hints.extend(extract_flow_extensions(request.metadata())?);
let snapshot_seqs = extract_snapshot_seqs(request.metadata())?;
let channel = request
src/servers/src/grpc/greptime_handler.rs
+1
View File
@@ -81,6 +81,7 @@ impl GreptimeRequestHandler {
channel: Channel,
) -> Result<Output> {
let header = request.header.as_ref();
let hints = crate::hint_headers::validate_public_hints(hints)?;
let query_ctx = create_query_context(channel, header, hints, HashMap::new())?;
let query = request.request.context(InvalidQuerySnafu {
reason: "Expecting non-empty GreptimeRequest.",
src/servers/src/hint_headers.rs
+55
View File
@@ -16,6 +16,8 @@ use http::HeaderMap;
use session::hints::{HINT_KEYS, HINTS_KEY, HINTS_KEY_PREFIX};
use tonic::metadata::MetadataMap;
use crate::error::{InvalidParameterSnafu, Result};
pub(crate) fn extract_hints<T: ToHeaderMap>(headers: &T) -> Vec<(String, String)> {
let mut hints = Vec::new();
if let Some(value_str) = headers.get(HINTS_KEY) {
@@ -23,6 +25,16 @@ pub(crate) fn extract_hints<T: ToHeaderMap>(headers: &T) -> Vec<(String, String)
let mut parts = hint.splitn(2, '=');
if let (Some(key), Some(value)) = (parts.next(), parts.next()) {
hints.push((key.trim().to_string(), value.trim().to_string()));
} else {
let key = hint.trim();
let key_lower = key.to_ascii_lowercase();
let owned_prefix =
key_lower.starts_with("query.") || key_lower.starts_with("datafusion.");
if owned_prefix
|| session::query_options::canonical_query_option_name(key).is_some()
{
hints.push((key.to_string(), String::new()));
}
}
});
// If hints are provided in the `x-greptime-hints` header, ignore the rest of the headers
@@ -37,6 +49,49 @@ pub(crate) fn extract_hints<T: ToHeaderMap>(headers: &T) -> Vec<(String, String)
hints
}
/// Validate query options supplied through public protocol hints.
pub(crate) fn validate_public_hints(hints: Vec<(String, String)>) -> Result<Vec<(String, String)>> {
use std::collections::HashMap;
let mut normalized = Vec::with_capacity(hints.len());
let mut seen = HashMap::new();
for (key, value) in hints {
if session::hints::is_reserved_extension_key(&key) {
continue;
}
// Older Flownode clients send `query_parallelism=0` to request the
// server default; treat it as unset instead of rejecting the request.
if key.eq_ignore_ascii_case("query_parallelism") && value == "0" {
continue;
}
match session::query_options::parse_query_option(
&key,
&value,
session::query_options::MAX_PUBLIC_PARALLELISM,
) {
Ok(Some((canonical, normalized_value))) => {
if let Some(existing) = seen.insert(canonical.clone(), normalized_value.clone())
&& existing != normalized_value
{
return InvalidParameterSnafu {
reason: format!("Conflicting values for hint `{canonical}`"),
}
.fail();
}
normalized.push((canonical, normalized_value));
}
Ok(None) => normalized.push((key, value)),
Err(error) => {
return InvalidParameterSnafu {
reason: error.to_string(),
}
.fail();
}
}
}
Ok(normalized)
}
pub(crate) trait ToHeaderMap {
fn get(&self, key: &str) -> Option<&str>;
}
src/servers/src/http.rs
+5 -12
View File
@@ -1042,17 +1042,10 @@ impl HttpServer {
/// Attaches middlewares and debug routes to the router.
/// Callers should call this method after [HttpServer::make_app()].
pub fn build(&self, router: Router) -> Result<Router> {
let timeout_layer = if self.options.timeout != Duration::default() {
Some(
ServiceBuilder::new().layer(
DynamicTimeoutLayer::new(self.options.timeout)
.with_status_code_fn(Self::request_timeout_status_code),
),
)
} else {
info!("HTTP server timeout is disabled");
None
};
let timeout_layer = ServiceBuilder::new().layer(
DynamicTimeoutLayer::new(self.options.timeout)
.with_status_code_fn(Self::request_timeout_status_code),
);
let body_limit_layer = if self.options.body_limit != ReadableSize(0) {
Some(
ServiceBuilder::new()
@@ -1088,7 +1081,7 @@ impl HttpServer {
// and we have impl IntoResponse for Error. It will print out more detailed error messages
.layer(TraceLayer::new_for_http().on_failure(()))
.option_layer(cors_layer)
.option_layer(timeout_layer)
.layer(timeout_layer)
.option_layer(body_limit_layer)
// memory limit layer - must be before body is consumed
.layer(middleware::from_fn_with_state(
src/servers/src/http/hints.rs
+35 -6
View File
@@ -15,22 +15,36 @@
use axum::body::Body;
use axum::http::Request;
use axum::middleware::Next;
use axum::response::Response;
use axum::response::{IntoResponse, Response};
use common_telemetry::debug;
use session::context::QueryContext;
use session::hints::is_reserved_extension_key;
use crate::error::{InvalidParameterSnafu, Result};
use crate::hint_headers;
pub async fn extract_hints(mut request: Request<Body>, next: Next) -> Response {
let hints = hint_headers::extract_hints(request.headers());
if let Some(query_ctx) = request.extensions_mut().get_mut::<QueryContext>() {
apply_hints(query_ctx, hints);
let dedicated_read_preference = request
.headers()
.contains_key(&crate::http::header::GREPTIME_DB_HEADER_READ_PREFERENCE);
let hints = match hint_headers::validate_public_hints(hints) {
Ok(hints) => hints,
Err(error) => return error.into_response(),
};
if let Some(query_ctx) = request.extensions_mut().get_mut::<QueryContext>()
&& let Err(error) = apply_hints(query_ctx, hints, dedicated_read_preference)
{
return error.into_response();
}
next.run(request).await
}
fn apply_hints(query_ctx: &mut QueryContext, hints: Vec<(String, String)>) {
fn apply_hints(
query_ctx: &mut QueryContext,
hints: Vec<(String, String)>,
dedicated_read_preference: bool,
) -> Result<()> {
for (key, value) in hints {
if is_reserved_extension_key(&key) {
debug!(
@@ -39,8 +53,21 @@ fn apply_hints(query_ctx: &mut QueryContext, hints: Vec<(String, String)>) {
);
continue;
}
query_ctx.set_extension(key, value);
if key.eq_ignore_ascii_case(session::hints::READ_PREFERENCE_HINT) {
if !dedicated_read_preference {
let preference = value.parse::<session::ReadPreference>().map_err(|_| {
InvalidParameterSnafu {
reason: format!("Invalid read_preference hint: {value}"),
}
.build()
})?;
query_ctx.set_read_preference(preference);
}
} else {
query_ctx.set_extension(key, value);
}
}
Ok(())
}
#[cfg(test)]
@@ -86,7 +113,9 @@ mod tests {
),
("ttl".to_string(), "7d".to_string()),
],
);
false,
)
.unwrap();
assert_eq!(
query_ctx.remote_query_id(),
src/servers/src/http/read_preference.rs
+3 -4
View File
@@ -25,14 +25,13 @@ use crate::http::header::GREPTIME_DB_HEADER_READ_PREFERENCE;
/// Extract read preference from the request headers.
pub async fn extract_read_preference(mut request: Request<Body>, next: Next) -> Response {
let read_preference = request
if let Some(read_preference) = request
.headers()
.get(&GREPTIME_DB_HEADER_READ_PREFERENCE)
.and_then(|header| header.to_str().ok())
.and_then(|s| ReadPreference::from_str(s).ok())
.unwrap_or_default();
if let Some(query_ctx) = request.extensions_mut().get_mut::<QueryContext>() {
&& let Some(query_ctx) = request.extensions_mut().get_mut::<QueryContext>()
{
common_telemetry::debug!("Setting read preference to {}", read_preference);
query_ctx.set_read_preference(read_preference);
}
src/servers/src/mysql/federated.rs
+63 -2
View File
@@ -280,8 +280,31 @@ fn check_show_variables(query: &str) -> Option<Output> {
))
} else if SHOW_LOWER_CASE_PATTERN.is_match(query) {
Some(show_variables("lower_case_table_names", "0"))
} else if SHOW_VARIABLES_LIKE_PATTERN.is_match(query) {
Some(show_variables("", ""))
} else if let Some(matched) = SHOW_VARIABLES_LIKE_PATTERN.find(query) {
let mut remainder = strip_leading_comments(&query[matched.end()..]);
while let Some(rest) = remainder.strip_prefix(';') {
remainder = strip_leading_comments(rest);
}
if remainder.is_empty() {
Some(show_variables("", ""))
} else {
let remainder = remainder.trim_start();
// Filter keywords need an identifier boundary; named variables do not.
let looks_like_filter = ["WHERE", "LIKE"].iter().any(|keyword| {
remainder
.get(..keyword.len())
.is_some_and(|prefix| prefix.eq_ignore_ascii_case(keyword))
&& remainder[keyword.len()..]
.chars()
.next()
.is_none_or(|c| !(c == '_' || c == '$' || c.is_alphanumeric()))
});
if looks_like_filter {
Some(show_variables("", ""))
} else {
None
}
}
} else {
None
};
@@ -554,6 +577,44 @@ mod test {
use super::*;
#[test]
fn test_show_variables_where_filter_is_absorbed() {
// Old Connector/J `loadServerVariables` issues `SHOW VARIABLES WHERE
// Variable_name = '...'`; the federated path answers with an empty table
// instead of falling through to our parser, which would reject it.
let session = Arc::new(Session::new(None, Channel::Mysql, Default::default(), 0));
let output = check(
"show variables where Variable_name = 'sql_mode'",
QueryContext::arc(),
session,
)
.expect("WHERE-filtered SHOW VARIABLES should be absorbed");
assert!(matches!(output.data, OutputData::RecordBatches(_)));
}
#[test]
fn test_show_variables_named_variable_falls_back() {
// A named variable (e.g. `SHOW VARIABLES MAX_EXECUTION_TIME`,
// `SHOW VARIABLES query.parallelism`) reaches the real handler.
let session = Arc::new(Session::new(None, Channel::Mysql, Default::default(), 0));
assert!(
check(
"show variables max_execution_time",
QueryContext::arc(),
session.clone()
)
.is_none()
);
assert!(
check(
"show variables query.parallelism",
QueryContext::arc(),
session
)
.is_none()
);
}
#[test]
fn test_check_abnormal() {
let session = Arc::new(Session::new(None, Channel::Mysql, Default::default(), 0));
src/servers/src/postgres/handler.rs
+20 -2
View File
@@ -96,7 +96,7 @@ impl SimpleQueryHandler for PostgresServerHandlerInner {
}
statements
.iter()
.map(|s| s.to_string())
.map(statement_to_query)
.collect::<Vec<_>>()
.join(";")
} else {
@@ -131,6 +131,24 @@ impl SimpleQueryHandler for PostgresServerHandlerInner {
}
}
fn statement_to_query(statement: &SqlParserStatement) -> String {
if let SqlParserStatement::ShowVariable { variable } = statement
&& variable.len() > 1
&& (variable[0].value.eq_ignore_ascii_case("query")
|| variable[0].value.eq_ignore_ascii_case("datafusion"))
{
return format!(
"SHOW {}",
variable
.iter()
.map(ToString::to_string)
.collect::<Vec<_>>()
.join(".")
);
}
statement.to_string()
}
async fn send_warning_opt<C>(client: &mut C, query_context: QueryContextRef) -> PgWireResult<()>
where
C: Sink<PgWireBackendMessage> + Unpin + Send + Sync,
@@ -352,7 +370,7 @@ impl QueryParser for DefaultQueryParser {
}
let first_stmt = statements.remove(0);
let format = check_copy_to_stdout(&first_stmt);
(first_stmt.to_string(), format)
(statement_to_query(&first_stmt), format)
} else {
// bypass the error: it can run into error because of different
// versions of sqlparser
src/session/src/context.rs
+135 -4
View File
@@ -78,6 +78,8 @@ pub struct QueryContext {
sql_dialect: Arc<dyn Dialect + Send + Sync>,
#[builder(default)]
extensions: HashMap<String, String>,
#[builder(setter(skip))]
query_options_snapshotted: bool,
/// The configuration parameter are used to store the parameters that are set by the user
#[builder(default)]
configuration_parameter: Arc<ConfigurationVariables>,
@@ -481,6 +483,71 @@ 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;
}
// Older Frontends and Flownodes forward `query_parallelism=0` (their
// "use the server default" encoding) untouched inside the wire
// query context, so accept it as unset instead of failing the
// query on the Datanode side. The canonical
// `query.parallelism = 0` stays rejected.
if key.eq_ignore_ascii_case("query_parallelism") && value == "0" {
continue;
}
// Wire decoding preserves legacy frontend nonzero values; public entry points enforce the limit.
if let Some((canonical, normalized)) =
crate::query_options::parse_query_option(key, value, u16::MAX)?
{
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 = if self.query_options_snapshotted {
HashMap::new()
} else {
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();
let options = snapshot.effective_query_options()?;
snapshot
.extensions
.retain(|key, _| crate::query_options::canonical_query_option_name(key).is_none());
// Older Datanodes only understand the legacy `query_parallelism` hint key, so
// mirror the canonical parallelism override for the wire format while keeping
// the canonical key for in-process consumers.
let legacy_parallelism = options
.get("query.parallelism")
.map(|parallelism| ("query_parallelism".to_string(), parallelism.clone()));
snapshot.extensions.extend(options);
if let Some((key, value)) = legacy_parallelism {
snapshot.extensions.entry(key).or_insert(value);
}
snapshot.query_options_snapshotted = true;
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('"') {
@@ -676,6 +743,7 @@ impl QueryContextBuilder {
.sql_dialect
.unwrap_or_else(|| Arc::new(GreptimeDbDialect {})),
extensions,
query_options_snapshotted: false,
configuration_parameter: self
.configuration_parameter
.unwrap_or_else(|| Arc::new(ConfigurationVariables::default())),
@@ -740,7 +808,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 {
@@ -749,7 +817,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()),
}
}
}
@@ -784,11 +852,43 @@ 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,
crate::query_options::MAX_PUBLIC_PARALLELISM,
)?
.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()
}
}
@@ -801,6 +901,37 @@ mod test {
use crate::Session;
use crate::context::{Channel, *};
#[test]
fn test_query_parallelism_wire_accepts_legacy_value() {
let mut context = QueryContext::with(DEFAULT_CATALOG_NAME, "public");
context.set_extension("query_parallelism", "2048");
assert_eq!(
context
.effective_query_options()
.unwrap()
.get("query.parallelism"),
Some(&"2048".to_string())
);
}
#[test]
fn test_set_query_parallelism_public_limit() {
let context = QueryContext::with(DEFAULT_CATALOG_NAME, "public");
assert!(
context
.configuration_parameter()
.set_query_option("query.parallelism", "1025")
.is_err()
);
assert!(
context
.configuration_parameter()
.set_query_option("query.parallelism", "1024")
.is_ok()
);
}
#[test]
fn test_session() {
let session = Session::new(
src/session/src/lib.rs
+1
View File
@@ -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
+116
View File
@@ -0,0 +1,116 @@
// 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 a canonical query option alias or a DataFusion optimizer option name.
pub fn canonical_query_option_name(key: &str) -> Option<String> {
let key = key.to_ascii_lowercase();
match key.as_str() {
"query_parallelism" | "query.parallelism" => Some("query.parallelism".to_string()),
"query_fallback" | "allow_query_fallback" | "query.allow_query_fallback" => {
Some("query.allow_query_fallback".to_string())
}
"query.enable_remote_dynamic_filter_pushdown" => {
Some("query.enable_remote_dynamic_filter_pushdown".to_string())
}
_ if key.starts_with("datafusion.optimizer.") => Some(key),
_ => None,
}
}
/// Maximum query parallelism accepted through public query-option inputs.
pub const MAX_PUBLIC_PARALLELISM: u16 = 1024;
/// Validate a query option, returning its canonical name and normalized value.
pub fn parse_query_option(
key: &str,
value: &str,
max_parallelism: u16,
) -> 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) if (1..=max_parallelism).contains(&n) => n.to_string(),
_ => {
return InvalidConfigValueSnafu {
name: key,
value,
hint: format!("query.parallelism must be between 1 and {max_parallelism}"),
}
.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();
let set_result = options.set(&canonical, value).or_else(|error| {
if value.eq_ignore_ascii_case("true") || value.eq_ignore_ascii_case("false") {
options.set(&canonical, &value.to_ascii_lowercase())
} else {
Err(error)
}
});
if let Err(error) = set_result {
return InvalidConfigValueSnafu {
name: key,
value,
hint: error.to_string(),
}
.fail();
}
match options
.entries()
.into_iter()
.find(|entry| entry.key == canonical)
.and_then(|entry| entry.value)
{
Some(value) => value,
None => {
return InvalidConfigValueSnafu {
name: key,
value,
hint: "No matching DataFusion configuration entry",
}
.fail();
}
}
};
Ok(Some((canonical, normalized)))
}
src/sql/src/ast.rs
+34
View File
@@ -32,3 +32,37 @@ impl ObjectNamePartExt for ObjectNamePart {
ident.value.clone()
}
}
/// Formats an object name by joining its identifier values without quotes.
pub fn object_name_to_string_unquoted(name: &ObjectName) -> String {
name.0
.iter()
.map(ObjectNamePartExt::to_string_unquoted)
.collect::<Vec<_>>()
.join(".")
}
#[cfg(test)]
mod tests {
use sqlparser::ast::{Ident, ObjectName, ObjectNamePart};
use super::object_name_to_string_unquoted;
#[test]
fn test_object_name_to_string_unquoted() {
let quoted = ObjectName(vec![ObjectNamePart::Identifier(Ident::with_quote(
'"',
"query.parallelism",
))]);
let unquoted = ObjectName(vec![
ObjectNamePart::Identifier(Ident::new("query")),
ObjectNamePart::Identifier(Ident::new("parallelism")),
]);
assert_eq!("query.parallelism", object_name_to_string_unquoted(&quoted));
assert_eq!(
object_name_to_string_unquoted(&unquoted),
object_name_to_string_unquoted(&quoted)
);
}
}
tests/cases/standalone/common/basic.result
+5 -5
View File
@@ -196,11 +196,11 @@ affected_rows: 0
-- SQLNESS PROTOCOL MYSQL
SHOW VARIABLES MAX_EXECUTION_TIME;
+---------------+-------+
| Variable_name | Value |
+---------------+-------+
| | |
+---------------+-------+
+--------------------+
| MAX_EXECUTION_TIME |
+--------------------+
| 2000 |
+--------------------+
-- SQLNESS PROTOCOL MYSQL
SET MAX_EXECUTION_TIME = 0;