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::{IntoResponse, Response};
19use common_telemetry::debug;
20use session::context::QueryContext;
21use session::hints::is_reserved_extension_key;
22
23use crate::error::{InvalidParameterSnafu, Result};
24use crate::hint_headers;
25
26pub async fn extract_hints(mut request: Request<Body>, next: Next) -> Response {
27    let hints = hint_headers::extract_hints(request.headers());
28    let dedicated_read_preference = request
29        .headers()
30        .contains_key(&crate::http::header::GREPTIME_DB_HEADER_READ_PREFERENCE);
31    let hints = match hint_headers::validate_public_hints(hints) {
32        Ok(hints) => hints,
33        Err(error) => return error.into_response(),
34    };
35    if let Some(query_ctx) = request.extensions_mut().get_mut::<QueryContext>()
36        && let Err(error) = apply_hints(query_ctx, hints, dedicated_read_preference)
37    {
38        return error.into_response();
39    }
40    next.run(request).await
41}
42
43fn apply_hints(
44    query_ctx: &mut QueryContext,
45    hints: Vec<(String, String)>,
46    dedicated_read_preference: bool,
47) -> Result<()> {
48    for (key, value) in hints {
49        if is_reserved_extension_key(&key) {
50            debug!(
51                key = key.as_str(),
52                "Ignoring reserved external query context extension key"
53            );
54            continue;
55        }
56        if key.eq_ignore_ascii_case(session::hints::READ_PREFERENCE_HINT) {
57            if !dedicated_read_preference {
58                let preference = value.parse::<session::ReadPreference>().map_err(|_| {
59                    InvalidParameterSnafu {
60                        reason: format!("Invalid read_preference hint: {value}"),
61                    }
62                    .build()
63                })?;
64                query_ctx.set_read_preference(preference);
65            }
66        } else {
67            query_ctx.set_extension(key, value);
68        }
69    }
70    Ok(())
71}
72
73#[cfg(test)]
74mod tests {
75    use common_query::request::INITIAL_REMOTE_DYN_FILTER_REGISTRATIONS_EXTENSION_KEY as COMMON_INITIAL_REMOTE_DYN_FILTER_REGISTRATIONS_EXTENSION_KEY;
76    use query::options::FLOW_SCHEDULED_TIME_MILLIS;
77    use session::context::{QueryContextBuilder, generate_remote_query_id};
78    use session::hints::{
79        INITIAL_REMOTE_DYN_FILTER_REGISTRATIONS_EXTENSION_KEY, READ_PREFERENCE_EXTENSION_KEY,
80        REMOTE_QUERY_ID_EXTENSION_KEY,
81    };
82
83    use super::apply_hints;
84
85    #[test]
86    fn test_apply_hints_ignores_reserved_extension_keys_only() {
87        let original_query_id = generate_remote_query_id();
88        let mut query_ctx = QueryContextBuilder::default()
89            .set_extension(
90                REMOTE_QUERY_ID_EXTENSION_KEY.to_string(),
91                original_query_id.clone(),
92            )
93            .build();
94
95        apply_hints(
96            &mut query_ctx,
97            vec![
98                (
99                    REMOTE_QUERY_ID_EXTENSION_KEY.to_string(),
100                    "spoofed".to_string(),
101                ),
102                (
103                    INITIAL_REMOTE_DYN_FILTER_REGISTRATIONS_EXTENSION_KEY.to_string(),
104                    "spoofed-regs".to_string(),
105                ),
106                (
107                    READ_PREFERENCE_EXTENSION_KEY.to_string(),
108                    "follower".to_string(),
109                ),
110                (
111                    FLOW_SCHEDULED_TIME_MILLIS.to_string(),
112                    "1700000000000".to_string(),
113                ),
114                ("ttl".to_string(), "7d".to_string()),
115            ],
116            false,
117        )
118        .unwrap();
119
120        assert_eq!(
121            query_ctx.remote_query_id(),
122            Some(original_query_id.as_str())
123        );
124        assert!(
125            query_ctx
126                .extension(INITIAL_REMOTE_DYN_FILTER_REGISTRATIONS_EXTENSION_KEY)
127                .is_none()
128        );
129        assert_eq!(query_ctx.read_preference(), session::ReadPreference::Leader);
130        assert!(query_ctx.extension(READ_PREFERENCE_EXTENSION_KEY).is_none());
131        assert_eq!(
132            query_ctx.extension(FLOW_SCHEDULED_TIME_MILLIS),
133            Some("1700000000000")
134        );
135        assert_eq!(query_ctx.extension("ttl"), Some("7d"));
136    }
137
138    #[test]
139    fn test_initial_dyn_filter_registration_key_matches_common_query_constant() {
140        assert_eq!(
141            INITIAL_REMOTE_DYN_FILTER_REGISTRATIONS_EXTENSION_KEY,
142            COMMON_INITIAL_REMOTE_DYN_FILTER_REGISTRATIONS_EXTENSION_KEY
143        );
144    }
145}
servers/http/hints.rs

servers/http/
hints.rs
