Introduce attribute for derive

prometheus · Dec 12, 2022 · 87f0b7e · 87f0b7e
1 parent 68e908c
commit 87f0b7e
Show file tree

Hide file tree

Showing 2 changed files with 67 additions and 16 deletions.
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,35 @@ 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 flatten = 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() == "flatten")
+                        .unwrap_or(false);
+                    if flatten {
+                        let ident = f.ident.unwrap();
+                        quote! {
+                            self.#ident.encode(encoder)?;
+                        }
+                    } else {
+                        let ident = f.ident.unwrap();
+                        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,41 @@ 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);
+}