From 766158ae5cd15a68b492ddcf963331cf92befa22 Mon Sep 17 00:00:00 2001
From: Maaz Ahmed <maaz.a@subcom.tech>
Date: Thu, 30 Nov 2023 15:34:07 +0530
Subject: [PATCH] feat: basic metric query selector/builder
---
README.md | 8 +-
src/lib.rs | 13 ++-
src/query.rs | 254 +++++++++++++++++++++++++++++++++++++++++++++++++++
tests/api.rs | 9 +-
4 files changed, 276 insertions(+), 8 deletions(-)
create mode 100644 src/query.rs
diff --git a/README.md b/README.md
index 6f02ec6..9d84c54 100644
--- a/README.md
+++ b/README.md
@@ -3,14 +3,16 @@ A Rust library for handling PROMQL (and possibly MetricsQL) queries.
## Usage Example
```rust
-use mquery::{Auth, QueryManager};
+use mquery::{Auth, QueryManager, query};
#[tokio::main]
async fn main() {
dotenv::dotenv().expect("No .env file found in working dir");
let url = std::env::var("VM_URL").expect("VM URL not found in env");
let token = std::env::var("VM_TOKEN").expect("VM URL not found in env");
- let query = "total_http_requests";
+
+ let mut query = query::Metric::new("total_http_requests");
+ query.label_eq("method", "get");
let _response = QueryManager::new(url.parse().unwrap())
.auth(Auth::Bearer(token))
@@ -23,7 +25,7 @@ async fn main() {
# Roadmap
- [x] Basic raw queries (Instant and Ranged)
- [ ] Query Builder
- - [ ] Basic queries
+ - [x] Basic queries
- [ ] Operators
- [ ] Functions
- [ ] Runtime syntax checking
diff --git a/src/lib.rs b/src/lib.rs
index 38c3565..1bb88c7 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -2,7 +2,10 @@
//!
//! All HTTP requests are made using the `reqwest` client
+pub mod query;
+
use prometheus_http_query::{response::PromqlResult, Client, Error};
+use query::IntoQuery;
use reqwest::header::HeaderValue;
use url::Url;
@@ -59,8 +62,8 @@ impl QueryManager {
}
/// Send a query to the Prometheus server and retreived deserialized data
- pub async fn query<Q: AsRef<str>>(&self, query: Q) -> Result<PromqlResult, Error> {
- let mut builder = self.client.query(query.as_ref());
+ pub async fn query<Q: IntoQuery>(&self, query: Q) -> Result<PromqlResult, Error> {
+ let mut builder = self.client.query(query.into_query().as_ref());
if let Some((name, val)) = self.auth.get_header() {
builder = builder.header(name, val);
}
@@ -74,14 +77,16 @@ impl QueryManager {
}
/// Send a ranged query to the Prometheus server and retreived deserialized data
- pub async fn query_range<Q: AsRef<str>>(
+ pub async fn query_range<Q: IntoQuery>(
&self,
query: Q,
start: i64,
end: i64,
step: f64,
) -> Result<PromqlResult, Error> {
- let mut builder = self.client.query_range(query.as_ref(), start, end, step);
+ let mut builder = self
+ .client
+ .query_range(query.into_query().as_ref(), start, end, step);
if let Some((name, val)) = self.auth.get_header() {
builder = builder.header(name, val);
}
diff --git a/src/query.rs b/src/query.rs
new file mode 100644
index 0000000..da1a63f
--- /dev/null
+++ b/src/query.rs
@@ -0,0 +1,254 @@
+use std::fmt::Display;
+
+// Add default implementation of IntoQuery where the type simply returns itself
+// Used for types that already implement AsRef<str>
+macro_rules! impl_into_query {
+ ($($t:ty),+) => {
+ $(
+ impl IntoQuery for $t {
+ type Target = Self;
+ fn into_query(self) -> Self::Target {
+ self
+ }
+ }
+ )+
+ };
+}
+
+/// Marker trait to signify that a type is queriable
+///
+/// Has a blanket implementation for all types that implement `AsRef<str>`
+pub trait Query: AsRef<str> {}
+
+impl<T> Query for T where T: AsRef<str> {}
+
+/// Used by non-string types provided by the crate to convert into a type implements [``Query``]
+///
+/// For allowing raw queries, `IntoQuery` has been implemented for `String`, `&String`, and `&str`.
+pub trait IntoQuery {
+ type Target: Query;
+ fn into_query(self) -> Self::Target;
+}
+
+impl_into_query!(String, &String, &str);
+
+/// Metric selector/builder. Used to building queries with the help of the type system.
+///
+/// Note: besides making use of the type system, this type doesn't do any kind of validation on the resulting queries.
+#[derive(Default, Debug, Clone)]
+pub struct Metric<'a> {
+ name: &'a str,
+ labels: Vec<Label<'a>>,
+ at: Option<i32>,
+ lookback: Option<Duration>,
+ offset: Option<Offset>,
+}
+
+impl<'a> Metric<'a> {
+ /// Create a new metric with a metric name (for example `total_http_requests`)
+ pub fn new(name: &'a str) -> Self {
+ Self {
+ name,
+ ..Default::default()
+ }
+ }
+
+ /// Add a label matcher with an equality operator i.e. name="value"
+ ///
+ /// For addding multiple labels at once, see [`Self::labels`]
+ pub fn label_eq(&mut self, name: &'a str, value: &'a str) -> &mut Self {
+ self.labels.push(Label::Eq(name, value));
+ self
+ }
+
+ /// Add a label matcher with an inequality operator i.e. name!="value"
+ ///
+ /// For addding multiple labels at once, see [`Self::labels`]
+ pub fn label_ne(&mut self, name: &'a str, value: &'a str) -> &mut Self {
+ self.labels.push(Label::Ne(name, value));
+ self
+ }
+
+ /// Add a regular expression label matcher with an equality operator i.e. name=~".+"
+ ///
+ /// For addding multiple labels at once, see [`Self::labels`]
+ pub fn label_rx_eq(&mut self, name: &'a str, value: &'a str) -> &mut Self {
+ self.labels.push(Label::RxEq(name, value));
+ self
+ }
+
+ /// Add a regular expression label matcher with an equality operator i.e. name!~".+"
+ ///
+ /// For addding multiple labels at once, see [`Self::labels`]
+ pub fn label_rx_ne(&mut self, name: &'a str, value: &'a str) -> &mut Self {
+ self.labels.push(Label::RxNe(name, value));
+ self
+ }
+
+ /// Add multiple label matchers at once
+ ///
+ /// # Example
+ /// ```rust
+ /// use mquery::query::{Metric, Label};
+ ///
+ /// let _ = Metric::new("metric_name")
+ /// .labels([
+ /// Label::Eq("name", "value"),
+ /// Label::Ne("name", "value"),
+ /// Label::RxEq("name", "value"),
+ /// ]);
+ /// ```
+ pub fn labels<I: IntoIterator<Item = Label<'a>>>(&mut self, labels: I) -> &mut Self {
+ self.labels.extend(labels);
+ self
+ }
+
+ /// Set query evaluation time to the specified Unix timestamp using the `@` modifier.
+ pub fn at(&mut self, timestamp: i32) -> &mut Self {
+ // TODO: Convert between chrono DataTime
+ self.at.replace(timestamp);
+ self
+ }
+
+ /// Set lookback duration / range for the query with the appropriate unit of time
+ ///
+ /// More info on range selectors can be found at [Range Vector Selectors](https://prometheus.io/docs/prometheus/latest/querying/basics/#range-vector-selectors)
+ pub fn lookback(&mut self, duration: Duration) -> &mut Self {
+ self.lookback.replace(duration);
+ self
+ }
+
+ /// Set the offset time fromt the query evaluation time using the offset modifier with the appropriate unit
+ ///
+ /// Note: Offsets can contain [negative values](https://prometheus.io/docs/prometheus/latest/querying/basics/#offset-modifier)
+ pub fn offset(&mut self, offset: Offset) -> &mut Self {
+ self.offset.replace(offset);
+ self
+ }
+}
+
+impl Display for Metric<'_> {
+ fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+ f.write_str(self.name)?;
+ if !self.labels.is_empty() {
+ f.write_str("{")?;
+ for (i, l) in self.labels.iter().enumerate() {
+ write!(f, "{l}")?;
+ if i != (self.labels.len() - 1) {
+ f.write_str(",")?;
+ }
+ }
+ f.write_str("}")?;
+ }
+ if let Some(lookback) = self.lookback {
+ write!(f, "{lookback}")?;
+ }
+ if let Some(at) = self.at {
+ write!(f, " @ {}", at)?;
+ }
+ if let Some(off) = self.offset {
+ write!(f, " {off}")?;
+ }
+
+ Ok(())
+ }
+}
+
+impl IntoQuery for Metric<'_> {
+ type Target = String;
+ fn into_query(self) -> Self::Target {
+ self.to_string()
+ }
+}
+
+/// Label matcher type with (Name, Value) pairs
+#[derive(Debug, Clone)]
+pub enum Label<'a> {
+ Eq(&'a str, &'a str),
+ Ne(&'a str, &'a str),
+ RxEq(&'a str, &'a str),
+ RxNe(&'a str, &'a str),
+}
+
+impl Display for Label<'_> {
+ fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+ match self {
+ Label::Eq(n, v) => write!(f, "{n}=\"{v}\"")?,
+ Label::Ne(n, v) => write!(f, "{n}!=\"{v}\"")?,
+ Label::RxEq(n, v) => write!(f, "{n}=~\"{v}\"")?,
+ Label::RxNe(n, v) => write!(f, "{n}!~\"{v}\"")?,
+ }
+ Ok(())
+ }
+}
+
+/// Units of time to be used with [``Duration``] or [``Offset``]
+#[derive(Debug, Clone, Copy)]
+pub enum Unit {
+ Ms,
+ Sec,
+ Min,
+ Hr,
+ Day,
+ Week,
+ Year,
+}
+
+impl Display for Unit {
+ fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+ match self {
+ Unit::Ms => write!(f, "ms"),
+ Unit::Sec => write!(f, "s"),
+ Unit::Min => write!(f, "m"),
+ Unit::Hr => write!(f, "h"),
+ Unit::Day => write!(f, "d"),
+ Unit::Week => write!(f, "w"),
+ Unit::Year => write!(f, "y"),
+ }
+ }
+}
+
+/// Specify lookback / range duration for the queries
+#[derive(Debug, Clone, Copy)]
+pub struct Duration(u64, Unit);
+
+impl Display for Duration {
+ fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+ write!(f, "[{}{}]", self.0, self.1)
+ }
+}
+
+/// Specify time offset for the queries
+#[derive(Debug, Clone, Copy)]
+pub struct Offset(i64, Unit);
+
+impl Display for Offset {
+ fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+ write!(f, "offset {}{}", self.0, self.1)
+ }
+}
+
+#[cfg(test)]
+mod tests {
+ use super::*;
+
+ #[test]
+ fn metric_selector() {
+ let metric = Metric::new("metric_name")
+ .labels([
+ Label::Eq("name", "value"),
+ Label::Ne("name", "value"),
+ Label::RxEq("name", "value"),
+ ])
+ .label_rx_ne("name", "value")
+ .at(123456)
+ .lookback(Duration(1, Unit::Year))
+ .offset(Offset(-1, Unit::Year))
+ .to_string();
+
+ assert_eq!(
+ r#"metric_name{name="value",name!="value",name=~"value",name!~"value"}[1y] @ 123456 offset -1y"#,
+ metric
+ )
+ }
+}
diff --git a/tests/api.rs b/tests/api.rs
index f256da7..f8d5e16 100644
--- a/tests/api.rs
+++ b/tests/api.rs
@@ -1,4 +1,4 @@
-use mquery::{Auth, QueryManager};
+use mquery::{query, Auth, QueryManager};
use tokio::runtime::Runtime;
// TODO: Create proper tests with local server
@@ -19,6 +19,13 @@ fn deserialize_response() {
.unwrap();
}
+#[test]
+fn query_with_metric_selector() {
+ let metric = query::Metric::new("some_metric");
+ #[allow(clippy::let_underscore_future)]
+ let _ = QueryManager::default().query(metric);
+}
+
fn get_url_token() -> (String, String) {
dotenv::dotenv().expect("No .env file found in working dir");
(
--
GitLab