session/
query_options.rs

1// Copyright 2023 Greptime Team
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use crate::session_config::{Error, InvalidConfigValueSnafu};
16
17/// Returns a canonical query option alias or a DataFusion optimizer option name.
18pub fn canonical_query_option_name(key: &str) -> Option<String> {
19    let key = key.to_ascii_lowercase();
20    match key.as_str() {
21        "query_parallelism" | "query.parallelism" => Some("query.parallelism".to_string()),
22        "query_fallback" | "allow_query_fallback" | "query.allow_query_fallback" => {
23            Some("query.allow_query_fallback".to_string())
24        }
25        "query.enable_remote_dynamic_filter_pushdown" => {
26            Some("query.enable_remote_dynamic_filter_pushdown".to_string())
27        }
28        _ if key.starts_with("datafusion.optimizer.") => Some(key),
29        _ => None,
30    }
31}
32
33/// Maximum query parallelism accepted through public query-option inputs.
34pub const MAX_PUBLIC_PARALLELISM: u16 = 1024;
35
36/// Validate a query option, returning its canonical name and normalized value.
37pub fn parse_query_option(
38    key: &str,
39    value: &str,
40    max_parallelism: u16,
41) -> Result<Option<(String, String)>, Error> {
42    let key_lower = key.to_ascii_lowercase();
43    let canonical = match canonical_query_option_name(key) {
44        Some(name) => name,
45        None if key_lower.starts_with("query.") || key_lower.starts_with("datafusion.") => {
46            return InvalidConfigValueSnafu {
47                name: key,
48                value,
49                hint: "Unknown or unsupported query option",
50            }
51            .fail();
52        }
53        None => return Ok(None),
54    };
55
56    let normalized = if canonical == "query.parallelism" {
57        match value.parse::<u16>() {
58            Ok(n) if (1..=max_parallelism).contains(&n) => n.to_string(),
59            _ => {
60                return InvalidConfigValueSnafu {
61                    name: key,
62                    value,
63                    hint: format!("query.parallelism must be between 1 and {max_parallelism}"),
64                }
65                .fail();
66            }
67        }
68    } else if canonical.starts_with("query.") {
69        match value.to_ascii_lowercase().as_str() {
70            "true" => "true".to_string(),
71            "false" => "false".to_string(),
72            _ => {
73                return InvalidConfigValueSnafu {
74                    name: key,
75                    value,
76                    hint: "Expected true or false",
77                }
78                .fail();
79            }
80        }
81    } else {
82        let mut options = datafusion_common::config::ConfigOptions::default();
83        let set_result = options.set(&canonical, value).or_else(|error| {
84            if value.eq_ignore_ascii_case("true") || value.eq_ignore_ascii_case("false") {
85                options.set(&canonical, &value.to_ascii_lowercase())
86            } else {
87                Err(error)
88            }
89        });
90        if let Err(error) = set_result {
91            return InvalidConfigValueSnafu {
92                name: key,
93                value,
94                hint: error.to_string(),
95            }
96            .fail();
97        }
98        match options
99            .entries()
100            .into_iter()
101            .find(|entry| entry.key == canonical)
102            .and_then(|entry| entry.value)
103        {
104            Some(value) => value,
105            None => {
106                return InvalidConfigValueSnafu {
107                    name: key,
108                    value,
109                    hint: "No matching DataFusion configuration entry",
110                }
111                .fail();
112            }
113        }
114    };
115    Ok(Some((canonical, normalized)))
116}
session/query_options.rs

session/
query_options.rs
