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"),
149        );
150
151        let hints = extract_hints(&headers);
152
153        assert_eq!(hints.len(), 6);
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    }
170
171    #[test]
172    fn test_extract_hints_with_metadata_map() {
173        let mut metadata = MetadataMap::new();
174        metadata.insert(
175            "x-greptime-hint-auto_create_table",
176            MetadataValue::from_static("true"),
177        );
178        metadata.insert("x-greptime-hint-ttl", MetadataValue::from_static("3600d"));
179        metadata.insert(
180            "x-greptime-hint-append_mode",
181            MetadataValue::from_static("true"),
182        );
183        metadata.insert(
184            "x-greptime-hint-merge_mode",
185            MetadataValue::from_static("false"),
186        );
187        metadata.insert(
188            "x-greptime-hint-physical_table",
189            MetadataValue::from_static("table1"),
190        );
191        metadata.insert(
192            "x-greptime-hint-read_preference",
193            MetadataValue::from_static("leader"),
194        );
195
196        let hints = extract_hints(&metadata);
197
198        assert_eq!(hints.len(), 6);
199        assert_eq!(
200            hints[0],
201            ("auto_create_table".to_string(), "true".to_string())
202        );
203        assert_eq!(hints[1], ("ttl".to_string(), "3600d".to_string()));
204        assert_eq!(hints[2], ("append_mode".to_string(), "true".to_string()));
205        assert_eq!(hints[3], ("merge_mode".to_string(), "false".to_string()));
206        assert_eq!(
207            hints[4],
208            ("physical_table".to_string(), "table1".to_string())
209        );
210        assert_eq!(
211            hints[5],
212            ("read_preference".to_string(), "leader".to_string())
213        );
214    }
215
216    #[test]
217    fn test_extract_hints_with_partial_metadata_map() {
218        let mut metadata = MetadataMap::new();
219        metadata.insert(
220            "x-greptime-hint-auto_create_table",
221            MetadataValue::from_static("true"),
222        );
223        metadata.insert("x-greptime-hint-ttl", MetadataValue::from_static("3600d"));
224
225        let hints = extract_hints(&metadata);
226
227        assert_eq!(hints.len(), 2);
228        assert_eq!(
229            hints[0],
230            ("auto_create_table".to_string(), "true".to_string())
231        );
232        assert_eq!(hints[1], ("ttl".to_string(), "3600d".to_string()));
233    }
234}
servers/hint_headers.rs

servers/
hint_headers.rs
