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