feat(derive): Introduce flatten attribute (#118)

Introduce a `#[prometheus(flatten)]` attribute which can be used when deriving `EncodeLabelSet`, allowing a nested struct to be flattened during encoding. Signed-off-by: John Howard <[email protected]> Signed-off-by: Max Inden <[email protected]>
prometheus · Jan 4, 2023 · 9141c11 · 9141c11
1 parent c619ad5
commit 9141c11
Show file tree

Hide file tree

Showing 5 changed files with 99 additions and 18 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -11,7 +11,33 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Introduce `Collector` abstraction allowing users to provide additional metrics
   and their description on each scrape. See [PR 82].
 
+- Introduce a `#[prometheus(flatten)]` attribute which can be used when deriving `EncodeLabelSet`, allowing
+  a nested struct to be flattened during encoding. See [PR 118].
+
+  For example:
+
+  ```rust
+  #[derive(EncodeLabelSet, Hash, Clone, Eq, PartialEq, Debug)]
+  struct CommonLabels {
+      a: u64,
+      b: u64,
+  }
+  #[derive(EncodeLabelSet, Hash, Clone, Eq, PartialEq, Debug)]
+  struct Labels {
+      unique: u64,
+      #[prometheus(flatten)]
+      common: CommonLabels,
+  }
+  ```
+
+  Would be encoded as:
+
+  ```
+  my_metric{a="42",b="42",unique="42"} 42
+  ```
+
 [PR 82]: https://github.com/prometheus/client_rust/pull/82
+[PR 118]: https://github.com/prometheus/client_rust/pull/118
 
 ## [0.19.0]
 

diff --git a/Cargo.toml b/Cargo.toml
@@ -21,7 +21,7 @@ members = ["derive-encode"]
 dtoa = "1.0"
 itoa = "1.0"
 parking_lot = "0.12"
-prometheus-client-derive-encode = { version = "0.4.0", path = "derive-encode" }
+prometheus-client-derive-encode = { version = "0.4.1", path = "derive-encode" }
 prost = { version = "0.11.0", optional = true }
 prost-types = { version = "0.11.0", optional = true }
 

diff --git a/derive-encode/Cargo.toml b/derive-encode/Cargo.toml
@@ -1,6 +1,6 @@
 [package]
 name = "prometheus-client-derive-encode"
-version = "0.4.0"
+version = "0.4.1"
 authors = ["Max Inden <[email protected]>"]
 edition = "2021"
 description = "Auxiliary crate to derive Encode trait from prometheus-client."

diff --git a/derive-encode/src/lib.rs b/derive-encode/src/lib.rs
@@ -12,7 +12,7 @@ use quote::quote;
 use syn::DeriveInput;
 
 /// Derive `prometheus_client::encoding::EncodeLabelSet`.
-#[proc_macro_derive(EncodeLabelSet)]
+#[proc_macro_derive(EncodeLabelSet, attributes(prometheus))]
 pub fn derive_encode_label_set(input: TokenStream) -> TokenStream {
     let ast: DeriveInput = syn::parse(input).unwrap();
     let name = &ast.ident;
@@ -22,22 +22,40 @@ pub fn derive_encode_label_set(input: TokenStream) -> TokenStream {
             syn::Fields::Named(syn::FieldsNamed { named, .. }) => named
                 .into_iter()
                 .map(|f| {
-                    let ident = f.ident.unwrap();
-                    let ident_string = KEYWORD_IDENTIFIERS
+                    let attribute = f
+                        .attrs
                         .iter()
-                        .find(|pair| ident == pair.1)
-                        .map(|pair| pair.0.to_string())
-                        .unwrap_or_else(|| ident.to_string());
-
-                    quote! {
-                        let mut label_encoder = encoder.encode_label();
-                        let mut label_key_encoder = label_encoder.encode_label_key()?;
-                        EncodeLabelKey::encode(&#ident_string, &mut label_key_encoder)?;
-
-                        let mut label_value_encoder = label_key_encoder.encode_label_value()?;
-                        EncodeLabelValue::encode(&self.#ident, &mut label_value_encoder)?;
-
-                        label_value_encoder.finish()?;
+                        .find(|a| a.path.is_ident("prometheus"))
+                        .map(|a| a.parse_args::<syn::Ident>().unwrap().to_string());
+                    let flatten = match attribute.as_deref() {
+                        Some("flatten") => true,
+                        Some(other) => {
+                            panic!("Provided attribute '{other}', but only 'flatten' is supported")
+                        }
+                        None => false,
+                    };
+                    let ident = f.ident.unwrap();
+                    if flatten {
+                        quote! {
+                             EncodeLabelSet::encode(&self.#ident, encoder)?;
+                        }
+                    } else {
+                        let ident_string = KEYWORD_IDENTIFIERS
+                            .iter()
+                            .find(|pair| ident == pair.1)
+                            .map(|pair| pair.0.to_string())
+                            .unwrap_or_else(|| ident.to_string());
+
+                        quote! {
+                            let mut label_encoder = encoder.encode_label();
+                            let mut label_key_encoder = label_encoder.encode_label_key()?;
+                            EncodeLabelKey::encode(&#ident_string, &mut label_key_encoder)?;
+
+                            let mut label_value_encoder = label_key_encoder.encode_label_value()?;
+                            EncodeLabelValue::encode(&self.#ident, &mut label_value_encoder)?;
+
+                            label_value_encoder.finish()?;
+                        }
                     }
                 })
                 .collect(),

diff --git a/derive-encode/tests/lib.rs b/derive-encode/tests/lib.rs
@@ -136,3 +136,40 @@ fn remap_keyword_identifiers() {
         + "# EOF\n";
     assert_eq!(expected, buffer);
 }
+
+#[test]
+fn flatten() {
+    #[derive(EncodeLabelSet, Hash, Clone, Eq, PartialEq, Debug)]
+    struct CommonLabels {
+        a: u64,
+        b: u64,
+    }
+    #[derive(EncodeLabelSet, Hash, Clone, Eq, PartialEq, Debug)]
+    struct Labels {
+        unique: u64,
+        #[prometheus(flatten)]
+        common: CommonLabels,
+    }
+
+    let mut registry = Registry::default();
+    let family = Family::<Labels, Counter>::default();
+    registry.register("my_counter", "This is my counter", family.clone());
+
+    // Record a single HTTP GET request.
+    family
+        .get_or_create(&Labels {
+            unique: 1,
+            common: CommonLabels { a: 2, b: 3 },
+        })
+        .inc();
+
+    // Encode all metrics in the registry in the text format.
+    let mut buffer = String::new();
+    encode(&mut buffer, &registry).unwrap();
+
+    let expected = "# HELP my_counter This is my counter.\n".to_owned()
+        + "# TYPE my_counter counter\n"
+        + "my_counter_total{unique=\"1\",a=\"2\",b=\"3\"} 1\n"
+        + "# EOF\n";
+    assert_eq!(expected, buffer);
+}