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::Json;
16use axum::response::IntoResponse;
17use http::HeaderValue;
18use http::header::CONTENT_TYPE;
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 from_sql(sql: SqlOutput, execution_time_ms: u64) -> Self {
59        GreptimedbManageResponse {
60            manage_result: ManageResult::Sql { sql },
61            execution_time_ms,
62        }
63    }
64
65    pub fn from_dashboard(name: String, execution_time_ms: u64) -> Self {
66        GreptimedbManageResponse {
67            manage_result: ManageResult::Dashboards {
68                dashboards: vec![DashboardOutput {
69                    name,
70                    definition: String::new(),
71                }],
72            },
73            execution_time_ms,
74        }
75    }
76
77    pub fn from_dashboards(dashboards: Vec<DashboardOutput>, execution_time_ms: u64) -> Self {
78        GreptimedbManageResponse {
79            manage_result: ManageResult::Dashboards { dashboards },
80            execution_time_ms,
81        }
82    }
83
84    pub fn with_execution_time(mut self, execution_time: u64) -> Self {
85        self.execution_time_ms = execution_time;
86        self
87    }
88
89    pub fn execution_time_ms(&self) -> u64 {
90        self.execution_time_ms
91    }
92}
93
94#[derive(Serialize, Deserialize, Debug)]
95#[serde(untagged)]
96pub enum ManageResult {
97    Pipelines { pipelines: Vec<PipelineOutput> },
98    Sql { sql: SqlOutput },
99    Dashboards { dashboards: Vec<DashboardOutput> },
100}
101
102#[derive(Serialize, Deserialize, Debug)]
103pub struct PipelineOutput {
104    name: String,
105    version: String,
106    #[serde(skip_serializing_if = "Option::is_none")]
107    pipeline: Option<String>,
108}
109
110#[derive(Serialize, Deserialize, Debug)]
111pub struct DashboardOutput {
112    pub name: String,
113    #[serde(skip_serializing_if = "String::is_empty")]
114    pub definition: String,
115}
116
117#[derive(Serialize, Deserialize, Debug)]
118pub struct SqlOutput {
119    pub(crate) sql: String,
120    #[serde(skip_serializing_if = "Option::is_none")]
121    pub(crate) message: Option<String>,
122}
123
124impl IntoResponse for GreptimedbManageResponse {
125    fn into_response(self) -> axum::response::Response {
126        let execution_time = self.execution_time_ms;
127
128        let mut resp = Json(self).into_response();
129
130        // We deliberately don't add this format into [`crate::http::ResponseFormat`]
131        // because this is a format for manage api other than the data query api
132        resp.headers_mut().insert(
133            &GREPTIME_DB_HEADER_FORMAT,
134            HeaderValue::from_static("greptimedb_manage"),
135        );
136        resp.headers_mut().insert(
137            &GREPTIME_DB_HEADER_EXECUTION_TIME,
138            HeaderValue::from(execution_time),
139        );
140        resp.headers_mut().insert(
141            CONTENT_TYPE,
142            HeaderValue::from_str(mime_guess::mime::APPLICATION_JSON.as_ref()).unwrap(),
143        );
144
145        resp
146    }
147}
148
149#[cfg(test)]
150mod tests {
151    use arrow::datatypes::ToByteSlice;
152    use axum::body::to_bytes;
153
154    use super::*;
155
156    #[tokio::test]
157    async fn test_into_response() {
158        let resp = GreptimedbManageResponse {
159            manage_result: ManageResult::Pipelines {
160                pipelines: vec![PipelineOutput {
161                    name: "test_name".to_string(),
162                    version: "test_version".to_string(),
163                    pipeline: None,
164                }],
165            },
166            execution_time_ms: 42,
167        };
168
169        let re = resp.into_response();
170        let data_str = format!("{:?}", re);
171        assert_eq!(
172            data_str,
173            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) }"#
174        );
175
176        let body_bytes = to_bytes(re.into_body(), usize::MAX).await.unwrap();
177        let body_str = String::from_utf8_lossy(body_bytes.to_byte_slice());
178        assert_eq!(
179            body_str,
180            r#"{"pipelines":[{"name":"test_name","version":"test_version"}],"execution_time_ms":42}"#
181        );
182    }
183}