diff --git a/derive-encode/src/lib.rs b/derive-encode/src/lib.rs index 381c5bf7..2ed3f062 100644 --- 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 index 446af0ba..8cd4df7a 100644 --- 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, ®istry).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); +}