diff --git a/Cargo.lock b/Cargo.lock index 4d3a98d55159230a3cfc4ad7df3bff4023d15385..da88ed47b18f6c97c809624af8f58ccee403b136 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -649,7 +649,7 @@ dependencies = [ [[package]] name = "mquery" -version = "0.3.1" +version = "0.4.0" dependencies = [ "dotenv", "enum-as-inner", diff --git a/Cargo.toml b/Cargo.toml index 18d90aa13522d4f25b3560caca8a5bb096089d44..ad3459057dfc52d332674de317aa9a74944b843e 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -14,12 +14,14 @@ serde = { version = "1.0.193", optional = true } serde_json = { version = "1.0.108", optional = true } rayon = { version = "1.8.0", optional = true } mquery-macros = { path = "mquery-macros", optional = true } +tokio = { version = "1.34.0", default-features = false, optional = true } [features] default = [] utils = ["dep:serde", "dep:serde_json", "dep:rayon"] metricsql = [] macros = ["dep:mquery-macros"] +multi = ["dep:tokio"] [dev-dependencies] tokio = { version = "1.34.0", default-features = false, features = ["rt", "rt-multi-thread", "macros"]} diff --git a/src/lib.rs b/src/lib.rs index 42b74f64436b49c55312464d89d77ec573146877..131832986d10d52c667a3bc012f17acb3d7f9b71 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -16,6 +16,9 @@ use result::{IntoQryResult, QryResult}; use seal::Sealed; use url::Url; +#[cfg(feature = "multi")] +use tokio::task::{JoinError, JoinSet}; + /// The primary way to send queries and get deserialized responses from the metrics server /// /// The default instance uses no authentication, initializes it's own reqwest client and the @@ -68,7 +71,7 @@ impl QueryManager { self } - /// Send a query to the Prometheus server and retreived deserialized data + /// Send a query to the Prometheus server and retreive deserialized data /// /// This method accepts any type that implements the `IntoQuery` trait. /// The trait is implemented for `&str`, `String`, `&String`, and __mquery's__ @@ -88,6 +91,38 @@ impl QueryManager { } } + /// Send multiple queries to the Prometheus server and get a Vec of deserialized results + /// + /// This method accepts an iterator of any type that implements the `IntoQuery` trait. + /// The trait is implemented for `&str`, `String`, `&String`, and __mquery's__ + /// internal types such as [`query::Metric`], [`query::Scalar`] and others that result + /// from the query builder API. + /// + /// If there are different types of queries being used, you can turn them all into Strings + /// using the `to_string` method before passing them to the `queries` parameter. + #[cfg(feature = "multi")] + pub async fn query_multi<Q: IntoQuery + Send + Sync + 'static>( + &self, + queries: impl IntoIterator<Item = Q>, + ) -> Result<Vec<QryResult>, JoinError> { + let Self { + auth, + client, + timeout, + method, + } = self; + + let mut set = JoinSet::new(); + for q in queries { + set.spawn(query(q, auth.clone(), *timeout, client.clone(), *method)); + } + let mut res = Vec::with_capacity(set.len()); + while let Some(qr) = set.join_next().await { + res.push(qr?); + } + Ok(res) + } + /// Send a ranged query to the Prometheus server and retreived deserialized data /// /// This method accepts any type that implements the `IntoQuery` trait. @@ -182,3 +217,24 @@ impl<T: AsRef<str>> Display for ValidatedQuery<T> { impl<T: AsRef<str>> Operable for ValidatedQuery<T> {} impl<T: AsRef<str>> Sealed for ValidatedQuery<T> {} + +#[cfg(feature = "multi")] +async fn query<Q: IntoQuery>( + query: Q, + auth: Auth, + timeout: Option<i64>, + client: Client, + method: Method, +) -> QryResult { + let mut builder = client.query(query.into_query().as_ref()); + if let Some((name, val)) = auth.get_header() { + builder = builder.header(name, val); + } + if let Some(timeout) = timeout.as_ref() { + builder = builder.timeout(*timeout); + } + match method { + Method::Get => builder.get().await.into_qry_result(), + Method::Post => builder.post().await.into_qry_result(), + } +} diff --git a/tests/api.rs b/tests/api.rs index f8d5e163578ec975935ac38c2ec28523bcb7b71f..da2fe672114b890ae10ce6ca84c607be4f4473bc 100644 --- a/tests/api.rs +++ b/tests/api.rs @@ -1,6 +1,11 @@ -use mquery::{query, Auth, QueryManager}; +use mquery::{ + query::{self, Metric}, + Auth, QueryManager, +}; use tokio::runtime::Runtime; +mod utils; + // TODO: Create proper tests with local server #[test] fn deserialize_response() { @@ -26,6 +31,16 @@ fn query_with_metric_selector() { let _ = QueryManager::default().query(metric); } +#[cfg(feature = "multi")] +#[tokio::test] +async fn multi_queries() { + let res = QueryManager::new(utils::URL.parse().unwrap()) + .query_multi([Metric::new("metric"), Metric::new("metric_two")]) + .await + .expect("failed to get multi query result"); + assert!(res.len() == 2); +} + fn get_url_token() -> (String, String) { dotenv::dotenv().expect("No .env file found in working dir"); (