druid_io/
client.rs

1use crate::connection::{BrokersPool, SelectionStategy, StaticPool};
2use crate::query::response::GroupByResponse;
3use crate::query::response::MetadataResponse;
4use crate::query::response::ScanResponse;
5use crate::query::response::SearchResponse;
6use crate::query::response::SegmentMetadataResponse;
7use crate::query::response::TimeBoundaryResponse;
8use crate::query::response::{TimeseriesResponse, TopNResponse};
9use crate::query::timeseries::Timeseries;
10use crate::query::{
11    group_by::GroupBy, scan::Scan, search::Search, segment_metadata::SegmentMetadata,
12    time_boundary::TimeBoundary, top_n::TopN, DataSource,
13};
14use crate::query::{DataSourceMetadata, Query};
15use reqwest::Client;
16use serde::de::DeserializeOwned;
17use serde::Serialize;
18use std::collections::HashMap;
19use thiserror::Error;
20
21#[derive(Error, Debug)]
22#[non_exhaustive]
23pub enum DruidClientError {
24    #[error("http connection error")]
25    HttpConnection { source: reqwest::Error },
26    #[error("the data for key `{0}` is not available")]
27    Redaction(String),
28    #[error("invalid header (expected {expected:?}, found {found:?})")]
29    InvalidHeader { expected: String, found: String },
30    #[error("couldn't serialize object to json")]
31    ParsingError { source: serde_json::Error },
32    #[error("couldn't deserialize json to object")]
33    ParsingResponseError { source: serde_json::Error }, // todo: original json but with manageable size
34    #[error("Server responded with an error")]
35    ServerError { response: String },
36    #[error("unknown data store error")]
37    Unknown,
38}
39type ClientResult<T> = Result<T, DruidClientError>;
40
41pub struct DruidClient {
42    http_client: Client,
43    brokers_pool: Box<dyn BrokersPool>,
44}
45
46impl DruidClient {
47    pub fn new(nodes: Vec<String>) -> Self {
48        let strategy = SelectionStategy::default_for(&nodes);
49        DruidClient {
50            http_client: Client::new(),
51            brokers_pool: Box::new(StaticPool::new(nodes, strategy)),
52        }
53    }
54
55    fn url(&self) -> String {
56        format!("http://{}/druid/v2/?pretty", self.brokers_pool.broker())
57    }
58
59    async fn http_query(&self, request: &str) -> Result<String, DruidClientError> {
60        let response_str = self
61            .http_client
62            .post(&self.url())
63            .body(request.to_string())
64            .header(reqwest::header::CONTENT_TYPE, "application/json")
65            .send()
66            .await
67            .map_err(|source| DruidClientError::HttpConnection { source: source })?
68            .text()
69            .await
70            .map_err(|source| DruidClientError::HttpConnection { source: source })?;
71
72        let json_value = serde_json::from_str::<serde_json::Value>(&response_str)
73            .map_err(|err| DruidClientError::ParsingError { source: err });
74        if let Some(_) = json_value?.get("error") {
75            return Err(DruidClientError::ServerError {
76                response: response_str,
77            });
78        }
79        Ok(response_str)
80    }
81
82    pub async fn query<'a, T: DeserializeOwned + std::fmt::Debug + Serialize>(
83        &self,
84        query: &Query,
85    ) -> ClientResult<Vec<T>> {
86        self._query(query).await
87    }
88    pub async fn top_n<'a, T: DeserializeOwned + std::fmt::Debug + Serialize>(
89        &self,
90        query: &TopN,
91    ) -> ClientResult<Vec<TopNResponse<T>>> {
92        self._query(query).await
93    }
94
95    pub async fn search<'a, T: DeserializeOwned + std::fmt::Debug + Serialize>(
96        &self,
97        query: &Search,
98    ) -> ClientResult<Vec<SearchResponse>> {
99        self._query(query).await
100    }
101
102    pub async fn group_by<'a, T: DeserializeOwned + std::fmt::Debug + Serialize>(
103        &self,
104        query: &GroupBy,
105    ) -> ClientResult<Vec<GroupByResponse<T>>> {
106        self._query(query).await
107    }
108    pub async fn scan<'a, T: DeserializeOwned + std::fmt::Debug + Serialize>(
109        &self,
110        query: &Scan,
111    ) -> ClientResult<Vec<ScanResponse<T>>> {
112        self._query(query).await
113    }
114    pub async fn time_boundary<'a, T: DeserializeOwned + std::fmt::Debug + Serialize>(
115        &self,
116        query: &TimeBoundary,
117    ) -> ClientResult<Vec<TimeBoundaryResponse>> {
118        self._query(query).await
119    }
120    pub async fn segment_metadata(
121        &self,
122        query: &SegmentMetadata,
123    ) -> ClientResult<Vec<SegmentMetadataResponse>> {
124        self._query(query).await
125    }
126
127    pub async fn timeseries<'a, T: DeserializeOwned + std::fmt::Debug + Serialize>(
128        &self,
129        query: &Timeseries,
130    ) -> ClientResult<Vec<TimeseriesResponse<T>>> {
131        self._query(query).await
132    }
133
134    async fn _query<Req, Resp>(&self, query: &Req) -> ClientResult<Resp>
135    where
136        Req: Serialize,
137        Resp: DeserializeOwned,
138    {
139        let request = serde_json::to_string(&query)
140            .map_err(|err| DruidClientError::ParsingError { source: err });
141
142        let response = match request {
143            Ok(str) => self.http_query(&str).await,
144            Err(e) => Err(e),
145        };
146
147        let response = response.and_then(|str| {
148            serde_json::from_str::<Resp>(&str)
149                .map_err(|source| DruidClientError::ParsingResponseError { source: source })
150        });
151
152        response
153    }
154
155    pub async fn datasource_metadata(
156        self,
157        data_source: DataSource,
158    ) -> ClientResult<Vec<MetadataResponse<HashMap<String, String>>>> {
159        let query = DataSourceMetadata {
160            data_source: data_source,
161            context: Default::default(),
162        };
163
164        self._query(&query).await
165    }
166}
druid_io/client.rs

druid_io/
client.rs
