rumeter_component/samplers/
gql.rs

1use std::collections::HashMap;
2
3use async_trait::async_trait;
4use serde::Serialize;
5use tracing::*;
6use crate::{Sampler, record::{RecordData, ResponseResult}};
7
8use super::http::HeaderMap;
9
10
11#[derive(Clone)]
12pub struct GraphQLSampler<T: Serialize + Clone + Send + Sync>{
13    label: String,
14    endpoint: String,
15    headers: HeaderMap,
16    body: RequestBody<T>,
17}
18
19#[derive(Serialize, Clone)]
20struct RequestBody<T: Serialize + Clone + Send> {
21    query: String,
22    variables: Option<T>,
23}
24
25impl <T: Serialize + Clone + Send + Sync> GraphQLSampler<T> {
26    pub fn new(label: &str, endpoint: &str, query: &str, headers: HeaderMap, vars: Option<T>) -> Self {
27        let body = RequestBody {
28            query: query.to_string(),
29            variables: vars,
30        };
31        Self { label: label.to_string(), endpoint: endpoint.to_string(), headers, body }
32    }
33
34    fn request_size(&self) -> u32 {
35        self.request_line_size() + self.request_headers_size() + self.request_body_size()
36    }
37
38    fn request_headers_size(&self) -> u32 {
39        let mut size = 0u32;
40        for (key, value) in self.headers.clone() {
41            match key {
42                Some(header_name) => {
43                    size = size + (header_name.to_string().len() + value.len() + ":\r\n".len()) as u32;
44                },
45                None => {},
46            }
47        }
48        size
49    }
50
51    fn request_line_size(&self) -> u32 {
52        ("POST".len() + self.endpoint.len() + "  HTTP/1.1\r\n".len()) as u32
53    }
54
55    fn request_body_size(&self) -> u32 {
56        let s = serde_json::to_string(&self.body);
57        (s.unwrap_or("".to_string()).len() + "\r\n".len()) as u32 
58    }
59}
60
61#[async_trait]
62impl <T: Serialize + Clone + Send + Sync> Sampler for GraphQLSampler<T> {
63    async fn run(&self) -> RecordData {
64        let client = reqwest::Client::new();
65        let start_send_timestamp = chrono::Local::now();
66        let resp = client.post(&self.endpoint).json(&self.body).headers(self.headers.clone()).send().await;
67        let finish_send_timestamp = chrono::Local::now();
68
69        match resp {
70            Ok(r) => {
71                let data_type = String::from("text");
72                let code = r.status().as_u16();
73                let resp_msg = r.status().canonical_reason().unwrap_or("Unknown");
74                let success = code < 400u16;
75                let fail_msg = if success {
76                    None
77                } else {
78                    Some(resp_msg.to_string())
79                };
80                let mut resp_headers: HashMap<String, String> = HashMap::new();
81                for (h_key, h_val) in r.headers() {
82                    resp_headers.insert(h_key.to_string(), h_val.to_str().unwrap().to_string());
83                }
84
85                let resp_body = r.text().await.unwrap_or("".to_string());
86
87                RecordData::new(
88                    start_send_timestamp.timestamp_millis() as u128,
89                    (finish_send_timestamp - start_send_timestamp).num_milliseconds() as u64,
90                    self.label.clone(),
91                    code,
92                    resp_msg.into(),
93                    "".to_string(),
94                    data_type,
95                    success,
96                    fail_msg,
97                    resp_body.len() as u64,
98                    self.request_size() as u64,
99                    0,
100                    0,
101                    self.endpoint.clone(),
102                    (finish_send_timestamp - start_send_timestamp).num_milliseconds() as u64,
103                    0,
104                    0,
105                    Some(ResponseResult::new(resp_headers, resp_body)),
106                )
107            },
108            Err(e) => {
109                error!("failed! --> {}", e.to_string());
110                RecordData::new(
111                    start_send_timestamp.timestamp_millis() as u128,
112                    (finish_send_timestamp - start_send_timestamp).num_milliseconds() as u64,
113                    self.label.clone(),
114                    0,
115                    "no data".to_string(),
116                    "".to_string(),
117                    "no data".to_string(),
118                    false,
119                    Some(e.to_string()),
120                    0u64,
121                    self.request_size() as u64,
122                    0,
123                    0,
124                    self.endpoint.clone(),
125                    (finish_send_timestamp - start_send_timestamp).num_milliseconds() as u64,
126                    0,
127                    0,
128                    None,
129                )
130            },
131        }
132
133    }
134}
135
rumeter_component/samplers/gql.rs

rumeter_component/samplers/
gql.rs
