From 3c276049140eda4c8cff6afb72d9ba6ba5bbbe95 Mon Sep 17 00:00:00 2001
From: Maaz Ahmed <maaz.a@subcom.tech>
Date: Tue, 26 Dec 2023 16:58:17 +0530
Subject: [PATCH] feat: initial support for handling multiple queries
---
Cargo.lock | 2 +-
Cargo.toml | 2 ++
src/lib.rs | 58 +++++++++++++++++++++++++++++++++++++++++++++++++++-
tests/api.rs | 17 ++++++++++++++-
4 files changed, 76 insertions(+), 3 deletions(-)
diff --git a/Cargo.lock b/Cargo.lock
index 4d3a98d..da88ed4 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 18d90aa..ad34590 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 42b74f6..1318329 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 f8d5e16..da2fe67 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");
(
--
GitLab