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::{HINTS_KEY, HINTS_KEY_PREFIX, HINT_KEYS};
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_hints_with_full_header_map() {
64        let mut headers = HeaderMap::new();
65        headers.insert(
66            "x-greptime-hint-auto_create_table",
67            HeaderValue::from_static("true"),
68        );
69        headers.insert("x-greptime-hint-ttl", HeaderValue::from_static("3600d"));
70        headers.insert(
71            "x-greptime-hint-append_mode",
72            HeaderValue::from_static("true"),
73        );
74        headers.insert(
75            "x-greptime-hint-merge_mode",
76            HeaderValue::from_static("false"),
77        );
78        headers.insert(
79            "x-greptime-hint-physical_table",
80            HeaderValue::from_static("table1"),
81        );
82        headers.insert(
83            "x-greptime-hint-read_preference",
84            HeaderValue::from_static("leader"),
85        );
86
87        let hints = extract_hints(&headers);
88
89        assert_eq!(hints.len(), 6);
90        assert_eq!(
91            hints[0],
92            ("auto_create_table".to_string(), "true".to_string())
93        );
94        assert_eq!(hints[1], ("ttl".to_string(), "3600d".to_string()));
95        assert_eq!(hints[2], ("append_mode".to_string(), "true".to_string()));
96        assert_eq!(hints[3], ("merge_mode".to_string(), "false".to_string()));
97        assert_eq!(
98            hints[4],
99            ("physical_table".to_string(), "table1".to_string())
100        );
101        assert_eq!(
102            hints[5],
103            ("read_preference".to_string(), "leader".to_string())
104        );
105    }
106
107    #[test]
108    fn test_extract_hints_with_missing_keys() {
109        let mut headers = HeaderMap::new();
110        headers.insert(
111            "x-greptime-hint-auto_create_table",
112            HeaderValue::from_static("true"),
113        );
114        headers.insert("x-greptime-hint-ttl", HeaderValue::from_static("3600d"));
115
116        let hints = extract_hints(&headers);
117
118        assert_eq!(hints.len(), 2);
119        assert_eq!(
120            hints[0],
121            ("auto_create_table".to_string(), "true".to_string())
122        );
123        assert_eq!(hints[1], ("ttl".to_string(), "3600d".to_string()));
124    }
125
126    #[test]
127    fn test_extract_hints_all_in_one() {
128        let mut headers = HeaderMap::new();
129        headers.insert(
130            "x-greptime-hints",
131            HeaderValue::from_static(" auto_create_table=true, ttl =3600d, append_mode=true , merge_mode=false , physical_table= table1,\
132            read_preference=leader"),
133        );
134
135        let hints = extract_hints(&headers);
136
137        assert_eq!(hints.len(), 6);
138        assert_eq!(
139            hints[0],
140            ("auto_create_table".to_string(), "true".to_string())
141        );
142        assert_eq!(hints[1], ("ttl".to_string(), "3600d".to_string()));
143        assert_eq!(hints[2], ("append_mode".to_string(), "true".to_string()));
144        assert_eq!(hints[3], ("merge_mode".to_string(), "false".to_string()));
145        assert_eq!(
146            hints[4],
147            ("physical_table".to_string(), "table1".to_string())
148        );
149        assert_eq!(
150            hints[5],
151            ("read_preference".to_string(), "leader".to_string())
152        );
153    }
154
155    #[test]
156    fn test_extract_hints_with_metadata_map() {
157        let mut metadata = MetadataMap::new();
158        metadata.insert(
159            "x-greptime-hint-auto_create_table",
160            MetadataValue::from_static("true"),
161        );
162        metadata.insert("x-greptime-hint-ttl", MetadataValue::from_static("3600d"));
163        metadata.insert(
164            "x-greptime-hint-append_mode",
165            MetadataValue::from_static("true"),
166        );
167        metadata.insert(
168            "x-greptime-hint-merge_mode",
169            MetadataValue::from_static("false"),
170        );
171        metadata.insert(
172            "x-greptime-hint-physical_table",
173            MetadataValue::from_static("table1"),
174        );
175        metadata.insert(
176            "x-greptime-hint-read_preference",
177            MetadataValue::from_static("leader"),
178        );
179
180        let hints = extract_hints(&metadata);
181
182        assert_eq!(hints.len(), 6);
183        assert_eq!(
184            hints[0],
185            ("auto_create_table".to_string(), "true".to_string())
186        );
187        assert_eq!(hints[1], ("ttl".to_string(), "3600d".to_string()));
188        assert_eq!(hints[2], ("append_mode".to_string(), "true".to_string()));
189        assert_eq!(hints[3], ("merge_mode".to_string(), "false".to_string()));
190        assert_eq!(
191            hints[4],
192            ("physical_table".to_string(), "table1".to_string())
193        );
194        assert_eq!(
195            hints[5],
196            ("read_preference".to_string(), "leader".to_string())
197        );
198    }
199
200    #[test]
201    fn test_extract_hints_with_partial_metadata_map() {
202        let mut metadata = MetadataMap::new();
203        metadata.insert(
204            "x-greptime-hint-auto_create_table",
205            MetadataValue::from_static("true"),
206        );
207        metadata.insert("x-greptime-hint-ttl", MetadataValue::from_static("3600d"));
208
209        let hints = extract_hints(&metadata);
210
211        assert_eq!(hints.len(), 2);
212        assert_eq!(
213            hints[0],
214            ("auto_create_table".to_string(), "true".to_string())
215        );
216        assert_eq!(hints[1], ("ttl".to_string(), "3600d".to_string()));
217    }
218}
servers/hint_headers.rs

servers/
hint_headers.rs
