pinot_client_rust/
request.rs

1use std::fmt;
2
3use serde::ser::SerializeMap;
4
5const QUERY_OPTIONS: &str = "queryOptions";
6const QUERY_OPTIONS_SQL_VALUE: &str = "groupByMode=sql;responseFormat=sql";
7
8#[derive(Clone, Debug, Eq, PartialEq)]
9pub enum QueryFormat {
10    PQL,
11    SQL,
12}
13
14impl fmt::Display for QueryFormat {
15    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
16        match self {
17            QueryFormat::PQL => write!(f, "pql"),
18            QueryFormat::SQL => write!(f, "sql"),
19        }
20    }
21}
22
23/// Request is used in server request to host multiple pinot query types, like PQL, SQL.
24#[derive(Clone, Debug, Eq, PartialEq)]
25pub struct Request {
26    pub query_format: QueryFormat,
27    pub query: String,
28}
29
30impl Request {
31    pub fn new(query_format: QueryFormat, query: &str) -> Self {
32        Self { query_format, query: query.to_string() }
33    }
34}
35
36impl serde::Serialize for Request {
37    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
38        where
39            S: serde::Serializer,
40    {
41        let n = match self.query_format {
42            QueryFormat::PQL => 1,
43            QueryFormat::SQL => 2,
44        };
45        let mut map = serializer.serialize_map(Some(n))?;
46        map.serialize_entry(
47            self.query_format.to_string().as_str(),
48            &self.query,
49        )?;
50        if self.query_format == QueryFormat::SQL {
51            map.serialize_entry(
52                QUERY_OPTIONS,
53                QUERY_OPTIONS_SQL_VALUE,
54            )?;
55        }
56        map.end()
57    }
58}
59
60pub(crate) fn encode_query_address(broker_address: &str, query_format: &QueryFormat) -> String {
61    let query_address = match query_format {
62        QueryFormat::PQL => format!("{}/query", broker_address),
63        QueryFormat::SQL => format!("{}/query/sql", broker_address),
64    };
65    if !broker_address.starts_with("http://") && !broker_address.starts_with("https://") {
66        format!("http://{}", query_address)
67    } else {
68        query_address
69    }
70}
71
72#[cfg(test)]
73mod test {
74    use serde_json::json;
75
76    use super::*;
77
78    #[test]
79    fn query_format_displays_correctly() {
80        assert_eq!(QueryFormat::PQL.to_string(), "pql".to_string());
81        assert_eq!(QueryFormat::SQL.to_string(), "sql".to_string());
82    }
83
84    #[test]
85    fn request_serialises_pql() {
86        let request = Request::new(QueryFormat::PQL, "SELECT * FROM test");
87        let result = json!(request);
88        let expected = json!({
89            "pql": "SELECT * FROM test",
90        });
91        assert_eq!(result, expected);
92    }
93
94    #[test]
95    fn request_serialises_sql() {
96        let request = Request::new(QueryFormat::SQL, "SELECT * FROM test");
97        let result = json!(request);
98        let expected = json!({
99            "sql": "SELECT * FROM test",
100            QUERY_OPTIONS: QUERY_OPTIONS_SQL_VALUE,
101        });
102        assert_eq!(result, expected);
103    }
104
105    #[test]
106    fn encode_query_address_encodes_pql_without_http_prefix() {
107        let broker_address = "localhost:8000";
108        assert_eq!(
109            encode_query_address(broker_address, &QueryFormat::PQL),
110            "http://localhost:8000/query"
111        );
112    }
113
114    #[test]
115    fn encode_query_address_encodes_sql_without_http_prefix() {
116        let broker_address = "localhost:8000";
117        assert_eq!(
118            encode_query_address(broker_address, &QueryFormat::SQL),
119            "http://localhost:8000/query/sql"
120        );
121    }
122
123    #[test]
124    fn encode_query_address_encodes_pql_with_http_prefix() {
125        let broker_address = "http://localhost:8000";
126        assert_eq!(
127            encode_query_address(broker_address, &QueryFormat::PQL),
128            "http://localhost:8000/query"
129        );
130    }
131
132    #[test]
133    fn encode_query_address_encodes_sql_with_http_prefix() {
134        let broker_address = "http://localhost:8000";
135        assert_eq!(
136            encode_query_address(broker_address, &QueryFormat::SQL),
137            "http://localhost:8000/query/sql"
138        );
139    }
140
141    #[test]
142    fn encode_query_address_encodes_pql_with_https_prefix() {
143        let broker_address = "https://localhost:8000";
144        assert_eq!(
145            encode_query_address(broker_address, &QueryFormat::PQL),
146            "https://localhost:8000/query"
147        );
148    }
149
150    #[test]
151    fn encode_query_address_encodes_sql_with_https_prefix() {
152        let broker_address = "https://localhost:8000";
153        assert_eq!(
154            encode_query_address(broker_address, &QueryFormat::SQL),
155            "https://localhost:8000/query/sql"
156        );
157    }
158}
pinot_client_rust/request.rs

pinot_client_rust/
request.rs
