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
19pub(crate) fn extract_hints<T: ToHeaderMap>(headers: &T) -> Vec<(String, String)> {
20    let mut hints = Vec::new();
21    if let Some(value_str) = headers.get(HINTS_KEY) {
22        value_str.split(',').for_each(|hint| {
23            let mut parts = hint.splitn(2, '=');
24            if let (Some(key), Some(value)) = (parts.next(), parts.next()) {
25                hints.push((key.trim().to_string(), value.trim().to_string()));
26            }
27        });
28        // If hints are provided in the `x-greptime-hints` header, ignore the rest of the headers
29        return hints;
30    }
31    for key in HINT_KEYS.iter() {
32        if let Some(value) = headers.get(key) {
33            let new_key = key.replace(HINTS_KEY_PREFIX, "");
34            hints.push((new_key, value.trim().to_string()));
35        }
36    }
37    hints
38}
39
40pub(crate) trait ToHeaderMap {
41    fn get(&self, key: &str) -> Option<&str>;
42}
43
44impl ToHeaderMap for MetadataMap {
45    fn get(&self, key: &str) -> Option<&str> {
46        self.get(key).and_then(|v| v.to_str().ok())
47    }
48}
49
50impl ToHeaderMap for HeaderMap {
51    fn get(&self, key: &str) -> Option<&str> {
52        self.get(key).and_then(|v| v.to_str().ok())
53    }
54}
55#[cfg(test)]
56mod tests {
57    use http::header::{HeaderMap, HeaderValue};
58    use tonic::metadata::{MetadataMap, MetadataValue};
59
60    use super::*;
61
62    #[test]
63    fn test_extract_skip_wal_hint() {
64        use session::hints::INSERT_SKIP_WAL_HINT;
65
66        let mut headers = HeaderMap::new();
67        headers.insert(HINTS_KEY, HeaderValue::from_static("insert_skip_wal=true"));
68        let mut metadata = MetadataMap::new();
69        metadata.insert(
70            HINTS_KEY,
71            MetadataValue::from_static("insert_skip_wal=true"),
72        );
73        let expected = vec![(INSERT_SKIP_WAL_HINT.to_string(), "true".to_string())];
74        assert_eq!(extract_hints(&headers), expected);
75        assert_eq!(extract_hints(&metadata), expected);
76    }
77
78    #[test]
79    fn test_extract_hints_with_full_header_map() {
80        let mut headers = HeaderMap::new();
81        headers.insert(
82            "x-greptime-hint-auto_create_table",
83            HeaderValue::from_static("true"),
84        );
85        headers.insert("x-greptime-hint-ttl", HeaderValue::from_static("3600d"));
86        headers.insert(
87            "x-greptime-hint-append_mode",
88            HeaderValue::from_static("true"),
89        );
90        headers.insert(
91            "x-greptime-hint-merge_mode",
92            HeaderValue::from_static("false"),
93        );
94        headers.insert(
95            "x-greptime-hint-physical_table",
96            HeaderValue::from_static("table1"),
97        );
98        headers.insert(
99            "x-greptime-hint-read_preference",
100            HeaderValue::from_static("leader"),
101        );
102
103        let hints = extract_hints(&headers);
104
105        assert_eq!(hints.len(), 6);
106        assert_eq!(
107            hints[0],
108            ("auto_create_table".to_string(), "true".to_string())
109        );
110        assert_eq!(hints[1], ("ttl".to_string(), "3600d".to_string()));
111        assert_eq!(hints[2], ("append_mode".to_string(), "true".to_string()));
112        assert_eq!(hints[3], ("merge_mode".to_string(), "false".to_string()));
113        assert_eq!(
114            hints[4],
115            ("physical_table".to_string(), "table1".to_string())
116        );
117        assert_eq!(
118            hints[5],
119            ("read_preference".to_string(), "leader".to_string())
120        );
121    }
122
123    #[test]
124    fn test_extract_hints_with_missing_keys() {
125        let mut headers = HeaderMap::new();
126        headers.insert(
127            "x-greptime-hint-auto_create_table",
128            HeaderValue::from_static("true"),
129        );
130        headers.insert("x-greptime-hint-ttl", HeaderValue::from_static("3600d"));
131
132        let hints = extract_hints(&headers);
133
134        assert_eq!(hints.len(), 2);
135        assert_eq!(
136            hints[0],
137            ("auto_create_table".to_string(), "true".to_string())
138        );
139        assert_eq!(hints[1], ("ttl".to_string(), "3600d".to_string()));
140    }
141
142    #[test]
143    fn test_extract_hints_all_in_one() {
144        let mut headers = HeaderMap::new();
145        headers.insert(
146            "x-greptime-hints",
147            HeaderValue::from_static(" auto_create_table=true, ttl =3600d, append_mode=true , merge_mode=false , physical_table= table1,\
148            read_preference=leader, compaction.twcs.time_window=1h"),
149        );
150
151        let hints = extract_hints(&headers);
152
153        assert_eq!(hints.len(), 7);
154        assert_eq!(
155            hints[0],
156            ("auto_create_table".to_string(), "true".to_string())
157        );
158        assert_eq!(hints[1], ("ttl".to_string(), "3600d".to_string()));
159        assert_eq!(hints[2], ("append_mode".to_string(), "true".to_string()));
160        assert_eq!(hints[3], ("merge_mode".to_string(), "false".to_string()));
161        assert_eq!(
162            hints[4],
163            ("physical_table".to_string(), "table1".to_string())
164        );
165        assert_eq!(
166            hints[5],
167            ("read_preference".to_string(), "leader".to_string())
168        );
169        assert_eq!(
170            hints[6],
171            ("compaction.twcs.time_window".to_string(), "1h".to_string())
172        );
173    }
174
175    #[test]
176    fn test_extract_hints_with_metadata_map() {
177        let mut metadata = MetadataMap::new();
178        metadata.insert(
179            "x-greptime-hint-auto_create_table",
180            MetadataValue::from_static("true"),
181        );
182        metadata.insert("x-greptime-hint-ttl", MetadataValue::from_static("3600d"));
183        metadata.insert(
184            "x-greptime-hint-append_mode",
185            MetadataValue::from_static("true"),
186        );
187        metadata.insert(
188            "x-greptime-hint-merge_mode",
189            MetadataValue::from_static("false"),
190        );
191        metadata.insert(
192            "x-greptime-hint-physical_table",
193            MetadataValue::from_static("table1"),
194        );
195        metadata.insert(
196            "x-greptime-hint-read_preference",
197            MetadataValue::from_static("leader"),
198        );
199
200        let hints = extract_hints(&metadata);
201
202        assert_eq!(hints.len(), 6);
203        assert_eq!(
204            hints[0],
205            ("auto_create_table".to_string(), "true".to_string())
206        );
207        assert_eq!(hints[1], ("ttl".to_string(), "3600d".to_string()));
208        assert_eq!(hints[2], ("append_mode".to_string(), "true".to_string()));
209        assert_eq!(hints[3], ("merge_mode".to_string(), "false".to_string()));
210        assert_eq!(
211            hints[4],
212            ("physical_table".to_string(), "table1".to_string())
213        );
214        assert_eq!(
215            hints[5],
216            ("read_preference".to_string(), "leader".to_string())
217        );
218    }
219
220    #[test]
221    fn test_extract_hints_with_partial_metadata_map() {
222        let mut metadata = MetadataMap::new();
223        metadata.insert(
224            "x-greptime-hint-auto_create_table",
225            MetadataValue::from_static("true"),
226        );
227        metadata.insert("x-greptime-hint-ttl", MetadataValue::from_static("3600d"));
228
229        let hints = extract_hints(&metadata);
230
231        assert_eq!(hints.len(), 2);
232        assert_eq!(
233            hints[0],
234            ("auto_create_table".to_string(), "true".to_string())
235        );
236        assert_eq!(hints[1], ("ttl".to_string(), "3600d".to_string()));
237    }
238}
servers/hint_headers.rs

servers/
hint_headers.rs
