servers/http/result/
greptime_result_v1.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::collections::HashMap;
16
17use axum::response::{IntoResponse, Response};
18use axum::Json;
19use common_query::Output;
20use headers::HeaderValue;
21use serde::{Deserialize, Serialize};
22use serde_json::Value;
23
24use crate::http::header::{
25    GREPTIME_DB_HEADER_EXECUTION_TIME, GREPTIME_DB_HEADER_FORMAT, GREPTIME_DB_HEADER_METRICS,
26};
27use crate::http::{handler, process_with_limit, GreptimeQueryOutput, HttpResponse, ResponseFormat};
28
29#[derive(Serialize, Deserialize, Debug)]
30pub struct GreptimedbV1Response {
31    #[serde(skip_serializing_if = "Vec::is_empty", default)]
32    pub(crate) output: Vec<GreptimeQueryOutput>,
33    pub(crate) execution_time_ms: u64,
34
35    // placeholder for header value
36    #[serde(skip)]
37    #[serde(default)]
38    pub(crate) resp_metrics: HashMap<String, Value>,
39}
40
41impl GreptimedbV1Response {
42    pub async fn from_output(outputs: Vec<crate::error::Result<Output>>) -> HttpResponse {
43        match handler::from_output(outputs).await {
44            Ok((output, resp_metrics)) => HttpResponse::GreptimedbV1(Self {
45                output,
46                execution_time_ms: 0,
47                resp_metrics,
48            }),
49            Err(err) => HttpResponse::Error(err),
50        }
51    }
52
53    pub fn output(&self) -> &[GreptimeQueryOutput] {
54        &self.output
55    }
56
57    pub fn with_execution_time(mut self, execution_time: u64) -> Self {
58        self.execution_time_ms = execution_time;
59        self
60    }
61
62    pub fn execution_time_ms(&self) -> u64 {
63        self.execution_time_ms
64    }
65
66    pub fn with_limit(mut self, limit: usize) -> Self {
67        self.output = process_with_limit(self.output, limit);
68        self
69    }
70}
71
72impl IntoResponse for GreptimedbV1Response {
73    fn into_response(self) -> Response {
74        let execution_time = self.execution_time_ms;
75        let metrics = if self.resp_metrics.is_empty() {
76            None
77        } else {
78            serde_json::to_string(&self.resp_metrics).ok()
79        };
80
81        let mut resp = Json(self).into_response();
82
83        resp.headers_mut().insert(
84            &GREPTIME_DB_HEADER_FORMAT,
85            HeaderValue::from_static(ResponseFormat::GreptimedbV1.as_str()),
86        );
87        resp.headers_mut().insert(
88            &GREPTIME_DB_HEADER_EXECUTION_TIME,
89            HeaderValue::from(execution_time),
90        );
91        if let Some(m) = metrics.and_then(|m| HeaderValue::from_str(&m).ok()) {
92            resp.headers_mut().insert(&GREPTIME_DB_HEADER_METRICS, m);
93        }
94
95        resp
96    }
97}