servers/http/
hints.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 axum::body::Body;
16use axum::http::Request;
17use axum::middleware::Next;
18use axum::response::Response;
19use common_telemetry::debug;
20use session::context::QueryContext;
21use session::hints::is_reserved_extension_key;
22
23use crate::hint_headers;
24
25pub async fn extract_hints(mut request: Request<Body>, next: Next) -> Response {
26    let hints = hint_headers::extract_hints(request.headers());
27    if let Some(query_ctx) = request.extensions_mut().get_mut::<QueryContext>() {
28        apply_hints(query_ctx, hints);
29    }
30    next.run(request).await
31}
32
33fn apply_hints(query_ctx: &mut QueryContext, hints: Vec<(String, String)>) {
34    for (key, value) in hints {
35        if is_reserved_extension_key(&key) {
36            debug!(
37                key = key.as_str(),
38                "Ignoring reserved external query context extension key"
39            );
40            continue;
41        }
42        query_ctx.set_extension(key, value);
43    }
44}
45
46#[cfg(test)]
47mod tests {
48    use common_query::request::INITIAL_REMOTE_DYN_FILTER_REGISTRATIONS_EXTENSION_KEY as COMMON_INITIAL_REMOTE_DYN_FILTER_REGISTRATIONS_EXTENSION_KEY;
49    use query::options::FLOW_SCHEDULED_TIME_MILLIS;
50    use session::context::{QueryContextBuilder, generate_remote_query_id};
51    use session::hints::{
52        INITIAL_REMOTE_DYN_FILTER_REGISTRATIONS_EXTENSION_KEY, READ_PREFERENCE_EXTENSION_KEY,
53        REMOTE_QUERY_ID_EXTENSION_KEY,
54    };
55
56    use super::apply_hints;
57
58    #[test]
59    fn test_apply_hints_ignores_reserved_extension_keys_only() {
60        let original_query_id = generate_remote_query_id();
61        let mut query_ctx = QueryContextBuilder::default()
62            .set_extension(
63                REMOTE_QUERY_ID_EXTENSION_KEY.to_string(),
64                original_query_id.clone(),
65            )
66            .build();
67
68        apply_hints(
69            &mut query_ctx,
70            vec![
71                (
72                    REMOTE_QUERY_ID_EXTENSION_KEY.to_string(),
73                    "spoofed".to_string(),
74                ),
75                (
76                    INITIAL_REMOTE_DYN_FILTER_REGISTRATIONS_EXTENSION_KEY.to_string(),
77                    "spoofed-regs".to_string(),
78                ),
79                (
80                    READ_PREFERENCE_EXTENSION_KEY.to_string(),
81                    "follower".to_string(),
82                ),
83                (
84                    FLOW_SCHEDULED_TIME_MILLIS.to_string(),
85                    "1700000000000".to_string(),
86                ),
87                ("ttl".to_string(), "7d".to_string()),
88            ],
89        );
90
91        assert_eq!(
92            query_ctx.remote_query_id(),
93            Some(original_query_id.as_str())
94        );
95        assert!(
96            query_ctx
97                .extension(INITIAL_REMOTE_DYN_FILTER_REGISTRATIONS_EXTENSION_KEY)
98                .is_none()
99        );
100        assert_eq!(query_ctx.read_preference(), session::ReadPreference::Leader);
101        assert!(query_ctx.extension(READ_PREFERENCE_EXTENSION_KEY).is_none());
102        assert_eq!(
103            query_ctx.extension(FLOW_SCHEDULED_TIME_MILLIS),
104            Some("1700000000000")
105        );
106        assert_eq!(query_ctx.extension("ttl"), Some("7d"));
107    }
108
109    #[test]
110    fn test_initial_dyn_filter_registration_key_matches_common_query_constant() {
111        assert_eq!(
112            INITIAL_REMOTE_DYN_FILTER_REGISTRATIONS_EXTENSION_KEY,
113            COMMON_INITIAL_REMOTE_DYN_FILTER_REGISTRATIONS_EXTENSION_KEY
114        );
115    }
116}
servers/http/hints.rs

servers/http/
hints.rs
