Skip to main content

mz_expr_derive_impl/
sqlfunc.rs

1// Copyright Materialize, Inc. and contributors. All rights reserved.
2//
3// Use of this software is governed by the Business Source License
4// included in the LICENSE file.
5//
6// As of the Change Date specified in that file, in accordance with
7// the Business Source License, use of this software will be governed
8// by the Apache License, Version 2.0.
9//! The `#[sqlfunc]` attribute macro's entry point.
10//!
11//! Parses the attribute, classifies the annotated function's arity, and hands the
12//! matching [`crate::shape::Shape`] to [`crate::generate::generate`].
13
14use darling::FromMeta;
15use proc_macro2::{Ident, TokenStream};
16use quote::quote;
17
18use crate::generate::generate;
19use crate::modifiers::Modifiers;
20use crate::shape::Shape;
21use crate::signature;
22
23/// Implementation for the `#[sqlfunc]` macro. The first parameter is the attribute
24/// arguments, the second is the function body. The third parameter indicates
25/// whether to include the test function in the output.
26///
27/// The feature `test` must be enabled to include the test function.
28pub fn sqlfunc(
29    attr: TokenStream,
30    item: TokenStream,
31    include_test: bool,
32) -> darling::Result<TokenStream> {
33    let mut attr_args = darling::ast::NestedMeta::parse_meta_list(attr.clone())?;
34
35    // Check if the first attribute arg is a bare Path (struct name).
36    let struct_ty = match attr_args.first() {
37        Some(darling::ast::NestedMeta::Meta(syn::Meta::Path(_))) => {
38            let darling::ast::NestedMeta::Meta(syn::Meta::Path(path)) = attr_args.remove(0) else {
39                unreachable!()
40            };
41            Some(path)
42        }
43        _ => None,
44    };
45
46    let modifiers = Modifiers::from_list(&attr_args)?;
47    let generate_tests = modifiers.generates_test();
48    let func = syn::parse2::<syn::ItemFn>(item.clone())?;
49
50    let Some(shape) = determine_shape(&func) else {
51        return Err(darling::Error::custom(
52            "nullary functions are not supported",
53        ));
54    };
55    let tokens = generate(shape, &func, modifiers, struct_ty, &attr)?;
56
57    let test = (generate_tests && include_test).then(|| generate_test(attr, item, &func.sig.ident));
58
59    Ok(quote! {
60        #tokens
61        #test
62    })
63}
64
65#[cfg(any(feature = "test", test))]
66fn generate_test(attr: TokenStream, item: TokenStream, name: &Ident) -> TokenStream {
67    let attr = attr.to_string();
68    let item = item.to_string();
69    let test_name = Ident::new(&format!("test_{}", name), name.span());
70    let fn_name = name.to_string();
71
72    quote! {
73        #[cfg(test)]
74        #[cfg_attr(miri, ignore)] // unsupported operation: extern static `pidfd_spawnp` is not supported by Miri
75        #[mz_ore::test]
76        fn #test_name() {
77            let (output, input) = mz_expr_derive_impl::test_sqlfunc_str(#attr, #item);
78            insta::assert_snapshot!(#fn_name, output, &input);
79        }
80    }
81}
82
83#[cfg(not(any(feature = "test", test)))]
84fn generate_test(_attr: TokenStream, _item: TokenStream, _name: &Ident) -> TokenStream {
85    quote! {}
86}
87
88/// The shape of a function annotated with `#[sqlfunc]`, or `None` for a function
89/// without SQL arguments.
90///
91/// Only the SQL arguments count, so a `&self` receiver and a trailing `&RowArena` are
92/// skipped. A `Variadic<T>` or `OptionalArg<T>` argument makes the function variadic
93/// regardless of how many arguments it has.
94fn determine_shape(func: &syn::ItemFn) -> Option<Shape> {
95    let args = signature::sql_args(func);
96    if args.iter().any(|arg| signature::is_variadic_arg(arg)) {
97        return Some(Shape::Variadic);
98    }
99    match args.len() {
100        0 => None,
101        1 => Some(Shape::Unary),
102        2 => Some(Shape::Binary),
103        _ => Some(Shape::Variadic),
104    }
105}