diff --git a/README.md b/README.md index 6f02ec6a123c1e694db9212e47df8400fad9dccd..9d84c54be25accab1451a9784a6ec2d7201ad2b1 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 38c35655c94b7f5d167e5658520e018eedec8093..1bb88c7509f6509fcaa57c7df77d6f3040e7f4e3 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 0000000000000000000000000000000000000000..da1a63fd06224b5007a0cc70e8ff2ddbf5b35d51 --- /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 f256da7af064d14c13099aeec78679e5ffc9920e..f8d5e163578ec975935ac38c2ec28523bcb7b71f 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"); (