servers/http/
read_preference.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 std::str::FromStr;
16
17use axum::body::Body;
18use axum::http::Request;
19use axum::middleware::Next;
20use axum::response::Response;
21use session::context::QueryContext;
22use session::ReadPreference;
23
24use crate::http::header::GREPTIME_DB_HEADER_READ_PREFERENCE;
25
26/// Extract read preference from the request headers.
27pub async fn extract_read_preference(mut request: Request<Body>, next: Next) -> Response {
28    let read_preference = request
29        .headers()
30        .get(&GREPTIME_DB_HEADER_READ_PREFERENCE)
31        .and_then(|header| header.to_str().ok())
32        .and_then(|s| ReadPreference::from_str(s).ok())
33        .unwrap_or_default();
34
35    if let Some(query_ctx) = request.extensions_mut().get_mut::<QueryContext>() {
36        common_telemetry::debug!("Setting read preference to {}", read_preference);
37        query_ctx.set_read_preference(read_preference);
38    }
39    next.run(request).await
40}