servers/
hint_headers.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 http::HeaderMap;
16use session::hints::{HINT_KEYS, HINTS_KEY, HINTS_KEY_PREFIX};
17use tonic::metadata::MetadataMap;
18
19use crate::error::{InvalidParameterSnafu, Result};
20
21pub(crate) fn extract_hints<T: ToHeaderMap>(headers: &T) -> Vec<(String, String)> {
22    let mut hints = Vec::new();
23    if let Some(value_str) = headers.get(HINTS_KEY) {
24        value_str.split(',').for_each(|hint| {
25            let mut parts = hint.splitn(2, '=');
26            if let (Some(key), Some(value)) = (parts.next(), parts.next()) {
27                hints.push((key.trim().to_string(), value.trim().to_string()));
28            } else {
29                let key = hint.trim();
30                let key_lower = key.to_ascii_lowercase();
31                let owned_prefix =
32                    key_lower.starts_with("query.") || key_lower.starts_with("datafusion.");
33                if owned_prefix
34                    || session::query_options::canonical_query_option_name(key).is_some()
35                {
36                    hints.push((key.to_string(), String::new()));
37                }
38            }
39        });
40        // If hints are provided in the `x-greptime-hints` header, ignore the rest of the headers
41        return hints;
42    }
43    for key in HINT_KEYS.iter() {
44        if let Some(value) = headers.get(key) {
45            let new_key = key.replace(HINTS_KEY_PREFIX, "");
46            hints.push((new_key, value.trim().to_string()));
47        }
48    }
49    hints
50}
51
52/// Validate query options supplied through public protocol hints.
53pub(crate) fn validate_public_hints(hints: Vec<(String, String)>) -> Result<Vec<(String, String)>> {
54    use std::collections::HashMap;
55
56    let mut normalized = Vec::with_capacity(hints.len());
57    let mut seen = HashMap::new();
58    for (key, value) in hints {
59        if session::hints::is_reserved_extension_key(&key) {
60            continue;
61        }
62        // Older Flownode clients send `query_parallelism=0` to request the
63        // server default; treat it as unset instead of rejecting the request.
64        if key.eq_ignore_ascii_case("query_parallelism") && value == "0" {
65            continue;
66        }
67        match session::query_options::parse_query_option(
68            &key,
69            &value,
70            session::query_options::MAX_PUBLIC_PARALLELISM,
71        ) {
72            Ok(Some((canonical, normalized_value))) => {
73                if let Some(existing) = seen.insert(canonical.clone(), normalized_value.clone())
74                    && existing != normalized_value
75                {
76                    return InvalidParameterSnafu {
77                        reason: format!("Conflicting values for hint `{canonical}`"),
78                    }
79                    .fail();
80                }
81                normalized.push((canonical, normalized_value));
82            }
83            Ok(None) => normalized.push((key, value)),
84            Err(error) => {
85                return InvalidParameterSnafu {
86                    reason: error.to_string(),
87                }
88                .fail();
89            }
90        }
91    }
92    Ok(normalized)
93}
94
95pub(crate) trait ToHeaderMap {
96    fn get(&self, key: &str) -> Option<&str>;
97}
98
99impl ToHeaderMap for MetadataMap {
100    fn get(&self, key: &str) -> Option<&str> {
101        self.get(key).and_then(|v| v.to_str().ok())
102    }
103}
104
105impl ToHeaderMap for HeaderMap {
106    fn get(&self, key: &str) -> Option<&str> {
107        self.get(key).and_then(|v| v.to_str().ok())
108    }
109}
110#[cfg(test)]
111mod tests {
112    use http::header::{HeaderMap, HeaderValue};
113    use tonic::metadata::{MetadataMap, MetadataValue};
114
115    use super::*;
116
117    #[test]
118    fn test_extract_skip_wal_hint() {
119        use session::hints::INSERT_SKIP_WAL_HINT;
120
121        let mut headers = HeaderMap::new();
122        headers.insert(HINTS_KEY, HeaderValue::from_static("insert_skip_wal=true"));
123        let mut metadata = MetadataMap::new();
124        metadata.insert(
125            HINTS_KEY,
126            MetadataValue::from_static("insert_skip_wal=true"),
127        );
128        let expected = vec![(INSERT_SKIP_WAL_HINT.to_string(), "true".to_string())];
129        assert_eq!(extract_hints(&headers), expected);
130        assert_eq!(extract_hints(&metadata), expected);
131    }
132
133    #[test]
134    fn test_extract_hints_with_full_header_map() {
135        let mut headers = HeaderMap::new();
136        headers.insert(
137            "x-greptime-hint-auto_create_table",
138            HeaderValue::from_static("true"),
139        );
140        headers.insert("x-greptime-hint-ttl", HeaderValue::from_static("3600d"));
141        headers.insert(
142            "x-greptime-hint-append_mode",
143            HeaderValue::from_static("true"),
144        );
145        headers.insert(
146            "x-greptime-hint-merge_mode",
147            HeaderValue::from_static("false"),
148        );
149        headers.insert(
150            "x-greptime-hint-physical_table",
151            HeaderValue::from_static("table1"),
152        );
153        headers.insert(
154            "x-greptime-hint-read_preference",
155            HeaderValue::from_static("leader"),
156        );
157
158        let hints = extract_hints(&headers);
159
160        assert_eq!(hints.len(), 6);
161        assert_eq!(
162            hints[0],
163            ("auto_create_table".to_string(), "true".to_string())
164        );
165        assert_eq!(hints[1], ("ttl".to_string(), "3600d".to_string()));
166        assert_eq!(hints[2], ("append_mode".to_string(), "true".to_string()));
167        assert_eq!(hints[3], ("merge_mode".to_string(), "false".to_string()));
168        assert_eq!(
169            hints[4],
170            ("physical_table".to_string(), "table1".to_string())
171        );
172        assert_eq!(
173            hints[5],
174            ("read_preference".to_string(), "leader".to_string())
175        );
176    }
177
178    #[test]
179    fn test_extract_hints_with_missing_keys() {
180        let mut headers = HeaderMap::new();
181        headers.insert(
182            "x-greptime-hint-auto_create_table",
183            HeaderValue::from_static("true"),
184        );
185        headers.insert("x-greptime-hint-ttl", HeaderValue::from_static("3600d"));
186
187        let hints = extract_hints(&headers);
188
189        assert_eq!(hints.len(), 2);
190        assert_eq!(
191            hints[0],
192            ("auto_create_table".to_string(), "true".to_string())
193        );
194        assert_eq!(hints[1], ("ttl".to_string(), "3600d".to_string()));
195    }
196
197    #[test]
198    fn test_extract_hints_all_in_one() {
199        let mut headers = HeaderMap::new();
200        headers.insert(
201            "x-greptime-hints",
202            HeaderValue::from_static(" auto_create_table=true, ttl =3600d, append_mode=true , merge_mode=false , physical_table= table1,\
203            read_preference=leader, compaction.twcs.time_window=1h"),
204        );
205
206        let hints = extract_hints(&headers);
207
208        assert_eq!(hints.len(), 7);
209        assert_eq!(
210            hints[0],
211            ("auto_create_table".to_string(), "true".to_string())
212        );
213        assert_eq!(hints[1], ("ttl".to_string(), "3600d".to_string()));
214        assert_eq!(hints[2], ("append_mode".to_string(), "true".to_string()));
215        assert_eq!(hints[3], ("merge_mode".to_string(), "false".to_string()));
216        assert_eq!(
217            hints[4],
218            ("physical_table".to_string(), "table1".to_string())
219        );
220        assert_eq!(
221            hints[5],
222            ("read_preference".to_string(), "leader".to_string())
223        );
224        assert_eq!(
225            hints[6],
226            ("compaction.twcs.time_window".to_string(), "1h".to_string())
227        );
228    }
229
230    #[test]
231    fn test_extract_hints_with_metadata_map() {
232        let mut metadata = MetadataMap::new();
233        metadata.insert(
234            "x-greptime-hint-auto_create_table",
235            MetadataValue::from_static("true"),
236        );
237        metadata.insert("x-greptime-hint-ttl", MetadataValue::from_static("3600d"));
238        metadata.insert(
239            "x-greptime-hint-append_mode",
240            MetadataValue::from_static("true"),
241        );
242        metadata.insert(
243            "x-greptime-hint-merge_mode",
244            MetadataValue::from_static("false"),
245        );
246        metadata.insert(
247            "x-greptime-hint-physical_table",
248            MetadataValue::from_static("table1"),
249        );
250        metadata.insert(
251            "x-greptime-hint-read_preference",
252            MetadataValue::from_static("leader"),
253        );
254
255        let hints = extract_hints(&metadata);
256
257        assert_eq!(hints.len(), 6);
258        assert_eq!(
259            hints[0],
260            ("auto_create_table".to_string(), "true".to_string())
261        );
262        assert_eq!(hints[1], ("ttl".to_string(), "3600d".to_string()));
263        assert_eq!(hints[2], ("append_mode".to_string(), "true".to_string()));
264        assert_eq!(hints[3], ("merge_mode".to_string(), "false".to_string()));
265        assert_eq!(
266            hints[4],
267            ("physical_table".to_string(), "table1".to_string())
268        );
269        assert_eq!(
270            hints[5],
271            ("read_preference".to_string(), "leader".to_string())
272        );
273    }
274
275    #[test]
276    fn test_extract_hints_with_partial_metadata_map() {
277        let mut metadata = MetadataMap::new();
278        metadata.insert(
279            "x-greptime-hint-auto_create_table",
280            MetadataValue::from_static("true"),
281        );
282        metadata.insert("x-greptime-hint-ttl", MetadataValue::from_static("3600d"));
283
284        let hints = extract_hints(&metadata);
285
286        assert_eq!(hints.len(), 2);
287        assert_eq!(
288            hints[0],
289            ("auto_create_table".to_string(), "true".to_string())
290        );
291        assert_eq!(hints[1], ("ttl".to_string(), "3600d".to_string()));
292    }
293}
servers/hint_headers.rs

servers/
hint_headers.rs
