servers/http/result/
greptime_manage_resp.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::response::IntoResponse;
16use axum::Json;
17use http::header::CONTENT_TYPE;
18use http::HeaderValue;
19use serde::{Deserialize, Serialize};
20
21use crate::http::header::{GREPTIME_DB_HEADER_EXECUTION_TIME, GREPTIME_DB_HEADER_FORMAT};
22
23/// Greptimedb Manage Api Response struct
24/// Currently we have `Pipelines` and `Scripts` as control panel api
25#[derive(Serialize, Deserialize, Debug)]
26pub struct GreptimedbManageResponse {
27    #[serde(flatten)]
28    pub(crate) manage_result: ManageResult,
29    pub(crate) execution_time_ms: u64,
30}
31
32impl GreptimedbManageResponse {
33    pub fn from_pipeline(
34        name: String,
35        version: String,
36        execution_time_ms: u64,
37        pipeline: Option<String>,
38    ) -> Self {
39        GreptimedbManageResponse {
40            manage_result: ManageResult::Pipelines {
41                pipelines: vec![PipelineOutput {
42                    name,
43                    version,
44                    pipeline,
45                }],
46            },
47            execution_time_ms,
48        }
49    }
50
51    pub fn from_pipelines(pipelines: Vec<PipelineOutput>, execution_time_ms: u64) -> Self {
52        GreptimedbManageResponse {
53            manage_result: ManageResult::Pipelines { pipelines },
54            execution_time_ms,
55        }
56    }
57
58    pub fn with_execution_time(mut self, execution_time: u64) -> Self {
59        self.execution_time_ms = execution_time;
60        self
61    }
62
63    pub fn execution_time_ms(&self) -> u64 {
64        self.execution_time_ms
65    }
66}
67
68#[derive(Serialize, Deserialize, Debug)]
69#[serde(untagged)]
70pub enum ManageResult {
71    Pipelines { pipelines: Vec<PipelineOutput> },
72    // todo(shuiyisong): refactor scripts api
73    Scripts(),
74}
75
76#[derive(Serialize, Deserialize, Debug)]
77pub struct PipelineOutput {
78    name: String,
79    version: String,
80    #[serde(skip_serializing_if = "Option::is_none")]
81    pipeline: Option<String>,
82}
83
84impl IntoResponse for GreptimedbManageResponse {
85    fn into_response(self) -> axum::response::Response {
86        let execution_time = self.execution_time_ms;
87
88        let mut resp = Json(self).into_response();
89
90        // We deliberately don't add this format into [`crate::http::ResponseFormat`]
91        // because this is a format for manage api other than the data query api
92        resp.headers_mut().insert(
93            &GREPTIME_DB_HEADER_FORMAT,
94            HeaderValue::from_static("greptimedb_manage"),
95        );
96        resp.headers_mut().insert(
97            &GREPTIME_DB_HEADER_EXECUTION_TIME,
98            HeaderValue::from(execution_time),
99        );
100        resp.headers_mut().insert(
101            CONTENT_TYPE,
102            HeaderValue::from_str(mime_guess::mime::APPLICATION_JSON.as_ref()).unwrap(),
103        );
104
105        resp
106    }
107}
108
109#[cfg(test)]
110mod tests {
111    use arrow::datatypes::ToByteSlice;
112    use axum::body::to_bytes;
113
114    use super::*;
115
116    #[tokio::test]
117    async fn test_into_response() {
118        let resp = GreptimedbManageResponse {
119            manage_result: ManageResult::Pipelines {
120                pipelines: vec![PipelineOutput {
121                    name: "test_name".to_string(),
122                    version: "test_version".to_string(),
123                    pipeline: None,
124                }],
125            },
126            execution_time_ms: 42,
127        };
128
129        let re = resp.into_response();
130        let data_str = format!("{:?}", re);
131        assert_eq!(
132            data_str,
133            r#"Response { status: 200, version: HTTP/1.1, headers: {"content-type": "application/json", "x-greptime-format": "greptimedb_manage", "x-greptime-execution-time": "42"}, body: Body(UnsyncBoxBody) }"#
134        );
135
136        let body_bytes = to_bytes(re.into_body(), usize::MAX).await.unwrap();
137        let body_str = String::from_utf8_lossy(body_bytes.to_byte_slice());
138        assert_eq!(
139            body_str,
140            r#"{"pipelines":[{"name":"test_name","version":"test_version"}],"execution_time_ms":42}"#
141        );
142    }
143}