servers/http/
opentsdb.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;
16use std::sync::Arc;
17
18use axum::body::Bytes;
19use axum::extract::{Query, State};
20use axum::http::StatusCode as HttpStatusCode;
21use axum::{Extension, Json};
22use common_error::ext::ErrorExt;
23use serde::{Deserialize, Serialize};
24use session::context::{Channel, QueryContext};
25use snafu::ResultExt;
26
27use crate::error::{self, Result};
28use crate::opentsdb::codec::DataPoint;
29use crate::query_handler::OpentsdbProtocolHandlerRef;
30
31#[derive(Serialize, Deserialize)]
32#[serde(untagged)]
33enum OneOrMany<T> {
34    One(T),
35    Vec(Vec<T>),
36}
37
38impl<T> From<OneOrMany<T>> for Vec<T> {
39    fn from(from: OneOrMany<T>) -> Self {
40        match from {
41            OneOrMany::One(val) => vec![val],
42            OneOrMany::Vec(vec) => vec,
43        }
44    }
45}
46
47#[derive(Serialize, Deserialize, Clone, Debug, PartialEq)]
48pub struct DataPointRequest {
49    metric: String,
50    timestamp: i64,
51    value: f64,
52    tags: HashMap<String, String>,
53}
54
55impl From<DataPointRequest> for DataPoint {
56    fn from(request: DataPointRequest) -> Self {
57        let ts_millis = DataPoint::timestamp_to_millis(request.timestamp);
58
59        let tags = request.tags.into_iter().collect::<Vec<(String, String)>>();
60
61        DataPoint::new(request.metric, ts_millis, request.value, tags)
62    }
63}
64
65#[derive(Serialize, Deserialize, Debug)]
66#[serde(untagged)]
67pub enum OpentsdbPutResponse {
68    Empty,
69    Debug(OpentsdbDebuggingResponse),
70}
71
72// Please refer to the OpenTSDB documents of ["api/put"](http://opentsdb.net/docs/build/html/api_http/put.html)
73// for more details.
74#[axum_macros::debug_handler]
75pub async fn put(
76    State(opentsdb_handler): State<OpentsdbProtocolHandlerRef>,
77    Query(params): Query<HashMap<String, String>>,
78    Extension(mut ctx): Extension<QueryContext>,
79    body: Bytes,
80) -> Result<(HttpStatusCode, Json<OpentsdbPutResponse>)> {
81    let summary = params.contains_key("summary");
82    let details = params.contains_key("details");
83
84    let data_point_requests = parse_data_points(body).await?;
85    let data_points = data_point_requests
86        .iter()
87        .map(|point| point.clone().into())
88        .collect::<Vec<_>>();
89
90    ctx.set_channel(Channel::Opentsdb);
91    let ctx = Arc::new(ctx);
92
93    if summary || details {
94        opentsdb_handler
95            .preflight(&data_points, ctx.clone())
96            .await?;
97    }
98
99    let response = if !summary && !details {
100        if let Err(e) = opentsdb_handler.exec(data_points, ctx.clone()).await {
101            // Not debugging purpose, failed fast.
102            return error::InternalSnafu {
103                err_msg: e.to_string(),
104            }
105            .fail();
106        }
107        (HttpStatusCode::NO_CONTENT, Json(OpentsdbPutResponse::Empty))
108    } else {
109        let mut response = OpentsdbDebuggingResponse {
110            success: 0,
111            failed: 0,
112            errors: if details {
113                Some(Vec::with_capacity(data_points.len()))
114            } else {
115                None
116            },
117        };
118
119        for (data_point, request) in data_points.into_iter().zip(data_point_requests) {
120            let result = opentsdb_handler.exec(vec![data_point], ctx.clone()).await;
121            match result {
122                Ok(affected_rows) => response.on_success(affected_rows),
123                Err(e) => response.on_failed(request, e),
124            }
125        }
126        (
127            HttpStatusCode::OK,
128            Json(OpentsdbPutResponse::Debug(response)),
129        )
130    };
131    Ok(response)
132}
133
134async fn parse_data_points(body: Bytes) -> Result<Vec<DataPointRequest>> {
135    let data_points = serde_json::from_slice::<OneOrMany<DataPointRequest>>(&body[..])
136        .context(error::InvalidOpentsdbJsonRequestSnafu)?;
137    Ok(data_points.into())
138}
139
140#[derive(Serialize, Deserialize, Debug)]
141struct OpentsdbDetailError {
142    datapoint: DataPointRequest,
143    error: String,
144}
145
146#[derive(Serialize, Deserialize, Debug)]
147pub struct OpentsdbDebuggingResponse {
148    success: i32,
149    failed: i32,
150    #[serde(skip_serializing_if = "Option::is_none")]
151    errors: Option<Vec<OpentsdbDetailError>>,
152}
153
154impl OpentsdbDebuggingResponse {
155    fn on_success(&mut self, affected_rows: usize) {
156        self.success += affected_rows as i32;
157    }
158
159    fn on_failed(&mut self, datapoint: DataPointRequest, error: impl ErrorExt) {
160        self.failed += 1;
161
162        if let Some(details) = self.errors.as_mut() {
163            let error = OpentsdbDetailError {
164                datapoint,
165                error: error.output_msg(),
166            };
167            details.push(error);
168        };
169    }
170}
171
172#[cfg(test)]
173mod test {
174
175    use super::*;
176
177    #[test]
178    fn test_into_opentsdb_data_point() {
179        let request = DataPointRequest {
180            metric: "hello".to_string(),
181            timestamp: 1234,
182            value: 1.0,
183            tags: HashMap::from([("foo".to_string(), "a".to_string())]),
184        };
185        let data_point: DataPoint = request.into();
186        assert_eq!(data_point.metric(), "hello");
187        assert_eq!(data_point.ts_millis(), 1234000);
188        assert_eq!(data_point.value(), 1.0);
189        assert_eq!(
190            data_point.tags(),
191            &vec![("foo".to_string(), "a".to_string())]
192        );
193    }
194
195    #[tokio::test]
196    async fn test_parse_data_points() {
197        let raw_data_point1 = r#"{
198                "metric": "sys.cpu.nice",
199                "timestamp": 1346846400,
200                "value": 18,
201                "tags": {
202                    "host": "web01",
203                    "dc": "lga"
204                }
205            }"#;
206        let data_point1 = serde_json::from_str::<DataPointRequest>(raw_data_point1).unwrap();
207
208        let raw_data_point2 = r#"{
209                "metric": "sys.cpu.nice",
210                "timestamp": 1346846400,
211                "value": 9,
212                "tags": {
213                    "host": "web02",
214                    "dc": "lga"
215                }
216            }"#;
217        let data_point2 = serde_json::from_str::<DataPointRequest>(raw_data_point2).unwrap();
218
219        let body = Bytes::from(raw_data_point1);
220        let data_points = parse_data_points(body).await.unwrap();
221        assert_eq!(data_points.len(), 1);
222        assert_eq!(data_points[0], data_point1);
223
224        let body = Bytes::from(format!("[{raw_data_point1},{raw_data_point2}]"));
225        let data_points = parse_data_points(body).await.unwrap();
226        assert_eq!(data_points.len(), 2);
227        assert_eq!(data_points[0], data_point1);
228        assert_eq!(data_points[1], data_point2);
229
230        let body = Bytes::from("");
231        let result = parse_data_points(body).await;
232        assert!(result.is_err());
233        let err = result.unwrap_err().output_msg();
234        assert!(err.contains("EOF while parsing a value at line 1 column 0"));
235
236        let body = Bytes::from("hello world");
237        let result = parse_data_points(body).await;
238        assert!(result.is_err());
239        let err = result.unwrap_err().output_msg();
240        assert!(err.contains("expected value at line 1 column 1"));
241    }
242}
servers/http/opentsdb.rs

servers/http/
opentsdb.rs
