Skip to main content

kubus_derive/
lib.rs

1use proc_macro::TokenStream;
2use quote::{ToTokens, quote};
3use syn::meta::ParseNestedMeta;
4use syn::{FnArg, Ident, ItemFn, LitInt, LitStr, PatType, ReturnType, Type};
5use syn::{parse_macro_input, parse_quote};
6
7#[derive(PartialEq, Eq)]
8enum EventType {
9    Apply,
10    Delete,
11}
12
13impl TryFrom<String> for EventType {
14    type Error = String;
15
16    fn try_from(value: String) -> Result<Self, Self::Error> {
17        match value.as_str() {
18            "Apply" => Ok(EventType::Apply),
19            "Delete" => Ok(EventType::Delete),
20            _ => Err(format!("Invalid event type {value}")),
21        }
22    }
23}
24
25impl ToTokens for EventType {
26    fn to_tokens(&self, tokens: &mut proc_macro2::TokenStream) {
27        let ty: Type = match self {
28            EventType::Apply => parse_quote!(::kubus::EventType::Apply),
29            EventType::Delete => parse_quote!(::kubus::EventType::Delete),
30        };
31        ty.to_tokens(tokens);
32    }
33}
34
35fn extract_generic_arg(ty: &Type, pos: usize) -> Option<&Type> {
36    if let syn::Type::Path(type_path) = ty
37        && let Some(last_segment) = type_path.path.segments.last()
38        && let syn::PathArguments::AngleBracketed(args) = &last_segment.arguments
39        && let Some(syn::GenericArgument::Type(inner_type)) = args.args.get(pos)
40    {
41        return Some(inner_type);
42    }
43    None
44}
45
46fn extract_return_type(func: &ItemFn) -> Option<&Type> {
47    match func.sig.output {
48        ReturnType::Type(_, ref ty) => Some(ty),
49        ReturnType::Default => None,
50    }
51}
52
53fn extract_func_arg_type(arg: &FnArg) -> Option<&PatType> {
54    match arg {
55        FnArg::Typed(pat) => Some(pat),
56        FnArg::Receiver(_) => None,
57    }
58}
59
60fn extract_resource_type(func: &ItemFn) -> Option<&Type> {
61    func.sig
62        .inputs
63        .first()
64        .and_then(|arg| extract_func_arg_type(arg))
65        .and_then(|pat| extract_generic_arg(&pat.ty, 0))
66}
67
68fn extract_context_type(func: &ItemFn) -> Option<&Type> {
69    func.sig
70        .inputs
71        .iter()
72        .nth(1)
73        .and_then(|arg| extract_func_arg_type(arg))
74        .and_then(|arc| extract_generic_arg(&arc.ty, 0))
75        .and_then(|ctx| extract_generic_arg(ctx, 0))
76}
77
78fn extract_function_return_error_type(func: &ItemFn) -> Option<&Type> {
79    extract_return_type(func).and_then(|ty| extract_generic_arg(ty, 1))
80}
81
82fn internal_prefix(ident: Ident) -> Ident {
83    let value = ident.to_string();
84    let prefixed = format!("__kubus_{value}");
85    Ident::new(&prefixed, ident.span())
86}
87
88fn quote_option<T>(value: Option<T>) -> proc_macro2::TokenStream
89where
90    T: ToTokens,
91{
92    match value {
93        Some(inner) => quote! { Some(#inner) },
94        None => quote! { None },
95    }
96}
97
98#[derive(Default)]
99struct EventHandlerAttrs {
100    event: Option<EventType>,
101    finalizer: Option<LitStr>,
102    label_selector: Option<LitStr>,
103    field_selector: Option<LitStr>,
104    requeue_interval: Option<LitInt>,
105}
106
107impl EventHandlerAttrs {
108    fn parse(&mut self, meta: ParseNestedMeta) -> syn::parse::Result<()> {
109        if meta.path.is_ident("event") {
110            let str: Ident = meta.value()?.parse()?;
111            self.event = Some(
112                str.to_string()
113                    .try_into()
114                    .map_err(|err| syn::parse::Error::new(str.span(), err))?,
115            );
116            Ok(())
117        } else if meta.path.is_ident("finalizer") {
118            self.finalizer = Some(meta.value()?.parse()?);
119            Ok(())
120        } else if meta.path.is_ident("label_selector") {
121            self.label_selector = Some(meta.value()?.parse()?);
122            Ok(())
123        } else if meta.path.is_ident("field_selector") {
124            self.field_selector = Some(meta.value()?.parse()?);
125            Ok(())
126        } else if meta.path.is_ident("requeue_interval") {
127            self.requeue_interval = Some(meta.value()?.parse()?);
128            Ok(())
129        } else {
130            Err(meta.error("unsupported kubus property"))
131        }
132    }
133}
134
135/// A procedural macro for defining Kubernetes event handlers in the Kubus framework.
136///
137/// # Attributes
138///
139/// * `event` (required) - The event type to handle: `Apply` or `Delete`
140/// * `finalizer` (optional) - The finalizer name to add (for Apply) or remove (for Delete)
141/// * `label_selector` (optional) - Label selector to filter which resources trigger this handler
142/// * `field_selector` (optional) - Field selector to filter which resources trigger this handler
143/// * `requeue_interval` (optional) - Requeue interval in seconds (default: 30)
144///
145/// # Function Signature
146///
147/// The annotated function must have the following signature:
148/// ```ignore
149/// async fn handler_name(
150///     resource: Arc<ResourceType>,
151///     context: Arc<Context<ContextType>>,
152/// ) -> Result<(), ErrorType>
153/// ```
154///
155/// # Examples
156///
157/// ```ignore
158/// #[kubus(event = Apply, finalizer = "my-app.example.com/cleanup")]
159/// async fn handle_pod_create(
160///     pod: Arc<Pod>,
161///     ctx: Arc<Context<MyContext>>,
162/// ) -> Result<(), MyError> {
163///     // Handle pod creation
164///     Ok(())
165/// }
166///
167/// #[kubus(
168///     event = Delete,
169///     finalizer = "my-app.example.com/cleanup",
170///     requeue_interval = 60
171/// )]
172/// async fn handle_pod_delete(
173///     pod: Arc<Pod>,
174///     ctx: Arc<Context<MyContext>>,
175/// ) -> Result<(), MyError> {
176///     // Handle pod deletion
177///     Ok(())
178/// }
179///
180/// #[kubus(
181///     event = Apply,
182///     label_selector = "app=my-app",
183///     field_selector = "status.phase=Running"
184/// )]
185/// async fn handle_running_pods(
186///     pod: Arc<Pod>,
187///     ctx: Arc<Context<MyContext>>,
188/// ) -> Result<(), MyError> {
189///     // Only handle running pods with app=my-app label
190///     Ok(())
191/// }
192/// ```
193#[proc_macro_attribute]
194pub fn kubus(args: TokenStream, input: TokenStream) -> TokenStream {
195    let mut attrs = EventHandlerAttrs::default();
196    let attr_parser = syn::meta::parser(|meta| attrs.parse(meta));
197    parse_macro_input!(args with attr_parser);
198
199    let event = attrs.event.expect("event kubus attribute missing");
200
201    let func = parse_macro_input!(input as ItemFn);
202    let resource_ty = extract_resource_type(&func)
203        .expect("unable to extract resource type from handler function");
204    let context_ty =
205        extract_context_type(&func).expect("unable to extract resource type from handler function");
206    let error_ty = extract_function_return_error_type(&func)
207        .cloned()
208        .unwrap_or_else(|| parse_quote! { ::kubus::HandlerError });
209
210    let internal_func = {
211        let mut func = func.clone();
212        func.sig.ident = internal_prefix(func.sig.ident);
213        func
214    };
215
216    let struct_name = func.sig.ident.clone();
217    let internal_func_name = internal_func.sig.ident.clone();
218
219    let field_selector = quote_option(attrs.field_selector);
220    let label_selector = quote_option(attrs.label_selector);
221    let requeue_interval = attrs.requeue_interval.unwrap_or_else(|| parse_quote!(30));
222
223    let update_finalizer = attrs
224        .finalizer
225        .map(|finalizer| {
226            let update_func = match event {
227                EventType::Apply => quote! { ::kubus::apply_finalizer },
228                EventType::Delete => quote! { ::kubus::remove_finalizer },
229            };
230            quote! {
231                let namespace = resource.namespace();
232                let client = context.client.clone();
233                let api: ::kube::Api<#resource_ty> =
234                    <#resource_ty as ::kube::Resource>::Scope::api(client, namespace);
235
236                #update_func(&api, #finalizer, resource).await?;
237            }
238        })
239        .unwrap_or_default();
240
241    let handler_name = LitStr::new(&struct_name.to_string(), struct_name.span());
242
243    quote! {
244        #[allow(non_snake_case)]
245        #internal_func
246
247        #[allow(non_camel_case_types)]
248        #[doc(hidden)]
249        pub struct #struct_name;
250
251        #[::async_trait::async_trait]
252        impl ::kubus::Handler<#resource_ty, #context_ty, #error_ty> for #struct_name {
253            const NAME: &'static str = #handler_name;
254            const FIELD_SELECTOR: Option<&'static str> = #field_selector;
255            const LABEL_SELECTOR: Option<&'static str> = #label_selector;
256
257            async fn handle(
258                resource: ::std::sync::Arc<#resource_ty>,
259                context: ::std::sync::Arc<::kubus::Context<#context_ty>>,
260            ) -> ::std::result::Result<::kube::runtime::controller::Action, #error_ty> {
261                use ::kubus::ScopeExt;
262                use ::kube::{Resource, ResourceExt};
263                let requeue = ::kube::runtime::controller::Action::requeue(
264                    ::std::time::Duration::from_secs(#requeue_interval)
265                );
266
267                if let (::kubus::EventType::Apply, Some(_)) | (::kubus::EventType::Delete, None) =
268                    (#event, resource.meta().deletion_timestamp.as_ref())
269                {
270                    return Ok(requeue);
271                }
272
273                #internal_func_name(resource.clone(), context.clone())
274                    .await
275                    .map_err(|err| ::kubus::Error::Handler(Box::new(err)))?;
276
277                #update_finalizer
278
279                Ok(requeue)
280            }
281        }
282
283        #[::async_trait::async_trait]
284        impl ::kubus::Runnable<#context_ty, #error_ty> for #struct_name {
285            fn name(&self) -> &'static str {
286                <#struct_name as ::kubus::Handler<#resource_ty, #context_ty, #error_ty>>::NAME
287            }
288
289            async fn run(
290                &self,
291                client: ::kube::Client,
292                context: ::std::sync::Arc<::kubus::Context<#context_ty>>,
293            ) -> ::std::result::Result<(), #error_ty> {
294                <#struct_name as ::kubus::Handler<#resource_ty, #context_ty, #error_ty>>::run(self, client, context).await
295            }
296        }
297    }
298    .into()
299}
300
301enum AdmissionKind {
302    Validating,
303    Mutating,
304}
305
306#[derive(Default)]
307struct AdmissionHandlerAttrs {
308    kind: Option<AdmissionKind>,
309}
310
311impl AdmissionHandlerAttrs {
312    fn parse(&mut self, meta: ParseNestedMeta) -> syn::parse::Result<()> {
313        if meta.path.is_ident("validating") && self.kind.is_none() {
314            self.kind = Some(AdmissionKind::Validating);
315            Ok(())
316        } else if meta.path.is_ident("mutating") && self.kind.is_none() {
317            self.kind = Some(AdmissionKind::Mutating);
318            Ok(())
319        } else {
320            Err(meta.error("unsupported kubus property"))
321        }
322    }
323}
324
325#[proc_macro_attribute]
326pub fn admission(args: TokenStream, input: TokenStream) -> TokenStream {
327    let mut attrs = AdmissionHandlerAttrs::default();
328    let attr_parser = syn::meta::parser(|meta| attrs.parse(meta));
329    parse_macro_input!(args with attr_parser);
330
331    let kind = attrs
332        .kind
333        .expect("admission attribute must specify either 'mutating' or 'validating'");
334
335    let func = parse_macro_input!(input as ItemFn);
336    let func_name = &func.sig.ident;
337    let name_string = LitStr::new(&func_name.to_string(), func_name.span());
338
339    // Extract error type from return type
340    let error_ty = extract_function_return_error_type(&func)
341        .cloned()
342        .unwrap_or_else(|| parse_quote! { ::kubus::HandlerError });
343
344    let internal_func = {
345        let mut func = func.clone();
346        func.sig.ident = internal_prefix(func.sig.ident);
347        func
348    };
349    let internal_func_name = &internal_func.sig.ident;
350
351    let (trait_name, method_name): (Type, Ident) = match kind {
352        AdmissionKind::Mutating => (
353            parse_quote! { ::kubus::admission::MutatingAdmissionHandler },
354            Ident::new("mutate", func_name.span()),
355        ),
356        AdmissionKind::Validating => (
357            parse_quote! { ::kubus::admission::ValidatingAdmissionHandler },
358            Ident::new("validate", func_name.span()),
359        ),
360    };
361
362    quote! {
363        #[allow(non_snake_case)]
364        #internal_func
365
366        #[allow(non_camel_case_types)]
367        #[doc(hidden)]
368        pub struct #func_name;
369
370        #[::async_trait::async_trait]
371        impl #trait_name for #func_name {
372            type Err = #error_ty;
373
374            fn name(&self) -> &'static str {
375                #name_string
376            }
377
378            async fn #method_name(
379                &self,
380                req: &::kube::core::admission::AdmissionRequest<::kube::api::DynamicObject>,
381            ) -> ::std::result::Result<::kube::core::admission::AdmissionResponse, Self::Err> {
382                #internal_func_name(req).await
383            }
384        }
385    }
386    .into()
387}