Skip to main content

deltalake_derive/
lib.rs

1use convert_case::{Case, Casing};
2use itertools::Itertools;
3use proc_macro::TokenStream;
4use proc_macro2::Span;
5use quote::quote;
6use syn::parse::Parser;
7use syn::{Attribute, Error, Expr, Ident, Result};
8use syn::{
9    Data, DeriveInput, Field, Fields, Lit, Meta, MetaNameValue, Token, Type, parse_macro_input,
10    punctuated::Punctuated,
11};
12
13/// Derive macro for implementing the TryUpdateKey trait
14///
15/// This macro automatically implements TryUpdateKey for a struct,
16/// mapping field names to configuration keys and using appropriate parsers
17/// based on the field type.
18///
19/// Additional key aliases can be specified with the `#[delta(alias = "alias.name")]` attribute.
20/// Multiple aliases can be added by `#[delta(alias = "foo", alias = "bar")]`.
21///
22/// Reading configuration can be achieved by assigning environmene keys to a field
23/// `#[delta(env = "MY_ENV_KEY")]`.
24#[proc_macro_derive(DeltaConfig, attributes(delta))]
25pub fn derive_delta_config(input: TokenStream) -> TokenStream {
26    // Parse the input tokens into a syntax tree
27    let input = parse_macro_input!(input as DeriveInput);
28
29    // Get the name of the struct
30    let name = &input.ident;
31
32    // Extract the fields from the struct
33    let fields = match &input.data {
34        Data::Struct(data) => match &data.fields {
35            Fields::Named(fields) => &fields.named,
36            _ => panic!("TryUpdateKey can only be derived for structs with named fields"),
37        },
38        _ => panic!("TryUpdateKey can only be derived for structs"),
39    }
40    .into_iter()
41    .collect::<Vec<_>>();
42
43    // Generate the implementation for TryUpdateKey trait
44    let try_update_key = match generate_try_update_key(name, &fields) {
45        Ok(try_update_key) => try_update_key,
46        Err(err) => return syn::Error::into_compile_error(err).into(),
47    };
48
49    // generate an enum with all configuration keys
50    let config_keys = match generate_config_keys(name, &fields) {
51        Ok(config_keys) => config_keys,
52        Err(err) => return syn::Error::into_compile_error(err).into(),
53    };
54
55    // generate a FromIterator implementation
56    let from_iter = generate_from_iterator(name);
57
58    let expanded = quote! {
59        #try_update_key
60
61        #config_keys
62
63        #from_iter
64    };
65
66    TokenStream::from(expanded)
67}
68
69fn generate_config_keys(name: &Ident, fields: &[&Field]) -> Result<proc_macro2::TokenStream> {
70    let enum_name = Ident::new(&format!("{name}Key"), Span::call_site());
71    let variants: Vec<_> = fields
72        .iter()
73        .map(|field| {
74            let field_name = &field
75                .ident
76                .as_ref()
77                .ok_or_else(|| syn::Error::new_spanned(field, "expected name"))?
78                .to_string();
79            let pascal_case = Ident::new(&field_name.to_case(Case::Pascal), Span::call_site());
80            let attributes = extract_field_attributes(&field.attrs)?;
81
82            // Generate doc attribute if documentation exists
83            let doc_attr = if let Some(doc_string) = attributes.docs {
84                // Create a doc attribute for the enum variant
85                quote! { #[doc = #doc_string] }
86            } else {
87                // No documentation
88                quote! {}
89            };
90
91            // Return the variant with its documentation
92            Ok(quote! {
93                #doc_attr
94                #pascal_case
95            })
96        })
97        .collect::<Result<_>>()?;
98    Ok(quote! {
99        #[automatically_derived]
100        pub enum #enum_name {
101            #(#variants),*
102        }
103    })
104}
105
106fn generate_from_iterator(name: &Ident) -> proc_macro2::TokenStream {
107    quote! {
108        #[automatically_derived]
109        impl<K, V> FromIterator<(K, V)> for #name
110        where
111            K: AsRef<str> + Into<String>,
112            V: AsRef<str> + Into<String>,
113        {
114            fn from_iter<I: IntoIterator<Item = (K, V)>>(iter: I) -> Self {
115                crate::logstore::config::ParseResult::from_iter(iter).config
116            }
117        }
118    }
119}
120
121fn generate_try_update_key(name: &Ident, fields: &[&Field]) -> Result<proc_macro2::TokenStream> {
122    let match_arms: Vec<_> = fields
123        .iter()
124        .filter_map(|field| {
125            let field_name = &field.ident.as_ref().unwrap();
126            let field_name_str = field_name.to_string();
127
128            // Extract aliases from attributes
129            let attributes = match extract_field_attributes(&field.attrs) {
130                Ok(attributes) => attributes,
131                Err(e) => return Some(Err(e)),
132            };
133            if attributes.skip {
134                return None;
135            }
136
137            // Determine parser based on field type
138            let (parser, is_option) = match determine_parser(&field.ty) {
139                Ok((parser, is_option)) => (parser, is_option),
140                Err(e) => return Some(Err(e)),
141            };
142
143            // Build the match conditions: field name and all aliases
144            let mut match_conditions = vec![quote! { #field_name_str }];
145            for alias in attributes.aliases {
146                match_conditions.push(quote! { #alias });
147            }
148
149            let match_arm = if is_option {
150                quote! {
151                    #(#match_conditions)|* => self.#field_name = Some(#parser(v)?),
152                }
153            } else {
154                quote! {
155                    #(#match_conditions)|* => self.#field_name = #parser(v)?,
156                }
157            };
158
159            Some(Ok(match_arm))
160        })
161        .try_collect()?;
162
163    let env_setters = generate_load_from_env(fields)?;
164
165    Ok(quote! {
166        #[automatically_derived]
167        impl crate::logstore::config::TryUpdateKey for #name {
168            fn try_update_key(&mut self, key: &str, v: &str) -> crate::DeltaResult<Option<()>> {
169                match key {
170                    #(#match_arms)*
171                    _ => return Ok(None),
172                }
173                Ok(Some(()))
174            }
175
176            fn load_from_environment(&mut self) -> crate::DeltaResult<()> {
177                let default_values = Self::default();
178                #(#env_setters)*
179                Ok(())
180            }
181        }
182    })
183}
184
185fn generate_load_from_env(fields: &[&Field]) -> Result<Vec<proc_macro2::TokenStream>> {
186    fields.iter().filter_map(|field| {
187        let field_name = &field.ident.as_ref().unwrap();
188        let attributes = match extract_field_attributes(&field.attrs) {
189            Ok(attributes) => attributes,
190            Err(e) => return Some(Err(e)),
191        };
192        if attributes.skip || attributes.env_variable_names.is_empty() {
193            return None;
194        }
195
196        let (parser, is_option) = match determine_parser(&field.ty) {
197            Ok((parser, is_option)) => (parser, is_option),
198            Err(e) => return Some(Err(e))
199        };
200
201        let env_checks = attributes.env_variable_names.iter().map(|env_var| {
202            if is_option {
203                // For Option types, only set if None
204                quote! {
205                    if self.#field_name.is_none() {
206                        if let Ok(val) = std::env::var(#env_var) {
207                            match #parser(&val) {
208                                Ok(parsed) => self.#field_name = Some(parsed),
209                                Err(e) => ::tracing::warn!("Failed to parse environment variable {}: {}", #env_var, e),
210                            }
211                        }
212                    }
213                }
214            } else {
215                // For non-Option types, we override the default value
216                // but ignore it if the current value is not the default.
217                quote! {
218                    if self.#field_name == default_values.#field_name {
219                        if let Ok(val) = std::env::var(#env_var) {
220                            match #parser(&val) {
221                                Ok(parsed) => self.#field_name = parsed,
222                                Err(e) => ::tracing::warn!("Failed to parse environment variable {}: {}", #env_var, e),
223                            }
224                        }
225                    }
226                }
227            }
228        });
229
230        Some(Ok(quote! {
231            #(#env_checks)*
232        }))
233    }).try_collect()
234}
235
236// Helper function to determine the appropriate parser based on field type
237fn determine_parser(ty: &Type) -> Result<(proc_macro2::TokenStream, bool)> {
238    match ty {
239        Type::Path(type_path) => {
240            let type_str = quote! { #type_path }.to_string();
241            let is_option = type_str.starts_with("Option");
242
243            let caller = if type_str.contains("usize") {
244                quote! { crate::logstore::config::parse_usize }
245            } else if type_str.contains("f64") || type_str.contains("f32") {
246                quote! { crate::logstore::config::parse_f64 }
247            } else if type_str.contains("Duration") {
248                quote! { crate::logstore::config::parse_duration }
249            } else if type_str.contains("bool") {
250                quote! { crate::logstore::config::parse_bool }
251            } else if type_str.contains("String") {
252                quote! { crate::logstore::config::parse_string }
253            } else {
254                return Err(Error::new_spanned(
255                    ty,
256                    format!(
257                        "Unsupported field type: {type_str}. Consider implementing a custom parser."
258                    ),
259                ));
260            };
261
262            Ok((caller, is_option))
263        }
264        _ => panic!("Unsupported field type for TryUpdateKey"),
265    }
266}
267
268struct FieldAttributes {
269    aliases: Vec<String>,
270    env_variable_names: Vec<String>,
271    docs: Option<String>,
272    skip: bool,
273}
274
275/// Extract aliases from field attributes
276///
277/// Parse the annotations from individual fields into a convenient structure.
278/// The field is annotated via `#[delta(...)]`. The following attributes are supported:
279/// - `alias = "alias_name"`: Specifies an alias for the field.
280/// - `env = "env_name"`: Specifies an environment variable name for the field.
281/// - `skip`: Specifies whether the field should be skipped during parsing.
282fn extract_field_attributes(attrs: &[Attribute]) -> Result<FieldAttributes> {
283    let mut aliases = Vec::new();
284    let mut environments = Vec::new();
285    let mut docs = None;
286    let mut doc_strings = Vec::new();
287    let mut skip = false;
288
289    for attr in attrs {
290        if attr.path().is_ident("doc") {
291            // Handle doc comments
292            let meta = attr.meta.require_name_value()?;
293            if let Expr::Lit(expr_lit) = &meta.value
294                && let Lit::Str(lit_str) = &expr_lit.lit
295            {
296                // Collect all doc strings - they might span multiple lines
297                doc_strings.push(lit_str.value().trim().to_string());
298            }
299        }
300        if attr.path().is_ident("delta") {
301            match &attr.meta {
302                Meta::List(list) => {
303                    let parser = Punctuated::<Meta, Token![,]>::parse_terminated;
304                    let parsed = parser.parse(list.tokens.clone().into())?;
305
306                    for val in parsed {
307                        match val {
308                            Meta::NameValue(MetaNameValue { path, value, .. }) => {
309                                if path.is_ident("alias")
310                                    && let Expr::Lit(lit_expr) = &value
311                                    && let Lit::Str(lit_str) = &lit_expr.lit
312                                {
313                                    aliases.push(lit_str.value());
314                                }
315                                if (path.is_ident("environment") || path.is_ident("env"))
316                                    && let Expr::Lit(lit_expr) = &value
317                                    && let Lit::Str(lit_str) = &lit_expr.lit
318                                {
319                                    environments.push(lit_str.value());
320                                }
321                            }
322                            Meta::Path(path) => {
323                                if path.is_ident("skip") {
324                                    skip = true;
325                                }
326                            }
327                            _ => {
328                                return Err(Error::new_spanned(
329                                    &attr.meta,
330                                    "only NameValue and Path parameters are supported",
331                                ));
332                            }
333                        }
334                    }
335                }
336                _ => {
337                    return Err(Error::new_spanned(
338                        &attr.meta,
339                        "expected a list-style attribute",
340                    ));
341                }
342            }
343        }
344    }
345
346    // Combine all doc strings into a single documentation string
347    if !doc_strings.is_empty() {
348        docs = Some(doc_strings.join("\n"));
349    }
350
351    Ok(FieldAttributes {
352        aliases,
353        env_variable_names: environments,
354        docs,
355        skip,
356    })
357}