common_macro/
lib.rs

1// Copyright 2023 Greptime Team
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15mod admin_fn;
16mod aggr_func;
17mod print_caller;
18mod range_fn;
19mod stack_trace_debug;
20mod utils;
21
22use aggr_func::{impl_aggr_func_type_store, impl_as_aggr_func_creator};
23use print_caller::process_print_caller;
24use proc_macro::TokenStream;
25use quote::quote;
26use range_fn::process_range_fn;
27use syn::{parse_macro_input, Data, DeriveInput, Fields};
28
29use crate::admin_fn::process_admin_fn;
30
31/// Make struct implemented trait [AggrFuncTypeStore], which is necessary when writing UDAF.
32/// This derive macro is expect to be used along with attribute macro [macro@as_aggr_func_creator].
33#[proc_macro_derive(AggrFuncTypeStore)]
34pub fn aggr_func_type_store_derive(input: TokenStream) -> TokenStream {
35    let ast = parse_macro_input!(input as DeriveInput);
36    impl_aggr_func_type_store(&ast)
37}
38
39/// A struct can be used as a creator for aggregate function if it has been annotated with this
40/// attribute first.
41///
42/// This attribute add a necessary field which is intended to store the input
43/// data's types to the struct.
44/// This attribute is expected to be used along with derive macro [AggrFuncTypeStore].
45#[proc_macro_attribute]
46pub fn as_aggr_func_creator(args: TokenStream, input: TokenStream) -> TokenStream {
47    impl_as_aggr_func_creator(args, input)
48}
49
50/// Attribute macro to convert an arithimetic function to a range function. The annotated function
51/// should accept servaral arrays as input and return a single value as output.
52///
53/// This procedure macro can works on any number of input parameters. Return type can be either
54/// primitive type or wrapped in `Option`.
55///
56/// # Example
57/// Take `count_over_time()` in PromQL as an example:
58/// ```rust, ignore
59/// /// The count of all values in the specified interval.
60/// #[range_fn(
61///     name = "CountOverTime",
62///     ret = "Float64Array",
63///     display_name = "prom_count_over_time"
64/// )]
65/// pub fn count_over_time(_: &TimestampMillisecondArray, values: &Float64Array) -> f64 {
66///      values.len() as f64
67/// }
68/// ```
69///
70/// # Arguments
71/// - `name`: The name of the generated `ScalarUDF` struct.
72/// - `ret`: The return type of the generated UDF function.
73/// - `display_name`: The display name of the generated UDF function.
74#[proc_macro_attribute]
75pub fn range_fn(args: TokenStream, input: TokenStream) -> TokenStream {
76    process_range_fn(args, input)
77}
78
79/// Attribute macro to convert a normal function to SQL administration function. The annotated function
80/// should accept:
81///    - `&ProcedureServiceHandlerRef` or `&TableMutationHandlerRef` or `FlowServiceHandlerRef` as the first argument,
82///    - `&QueryContextRef` as the second argument, and
83///    - `&[ValueRef<'_>]` as the third argument which is SQL function input values in each row.
84///
85/// Return type must be `common_query::error::Result<Value>`.
86///
87/// # Example see `common/function/src/system/procedure_state.rs`.
88///
89/// # Arguments
90/// - `name`: The name of the generated `Function` implementation.
91/// - `ret`: The return type of the generated SQL function, it will be transformed into `ConcreteDataType::{ret}_datatype()` result.
92/// - `display_name`: The display name of the generated SQL function.
93/// - `sig_fn`: the function to returns `Signature` of generated `Function`.
94/// - `user_path`: Optional path to the trait and context (e.g., `crate`);
95///   defaults to `crate` if not provided.
96#[proc_macro_attribute]
97pub fn admin_fn(args: TokenStream, input: TokenStream) -> TokenStream {
98    process_admin_fn(args, input)
99}
100
101/// Attribute macro to print the caller to the annotated function.
102/// The caller is printed as its filename and the call site line number.
103///
104/// This macro works like this: inject the tracking codes as the first statement to the annotated
105/// function body. The tracking codes use [backtrace-rs](https://crates.io/crates/backtrace) to get
106/// the callers. So you must dependent on the `backtrace-rs` crate.
107///
108/// # Arguments
109/// - `depth`: The max depth of call stack to print. Optional, defaults to 1.
110///
111/// # Example
112/// ```rust, ignore
113///
114/// #[print_caller(depth = 3)]
115/// fn foo() {}
116/// ```
117#[proc_macro_attribute]
118pub fn print_caller(args: TokenStream, input: TokenStream) -> TokenStream {
119    process_print_caller(args, input)
120}
121
122/// Attribute macro to derive [std::fmt::Debug] for the annotated `Error` type.
123///
124/// The generated `Debug` implementation will print the error in a stack trace style. E.g.:
125/// ```plaintext
126/// 0: Foo error, at src/common/catalog/src/error.rs:80:10
127/// 1: Bar error, at src/common/function/src/error.rs:90:10
128/// 2: Root cause, invalid table name, at src/common/catalog/src/error.rs:100:10
129/// ```
130///
131/// Notes on using this macro:
132/// - `#[snafu(display)]` must present on each enum variants,
133///   and should not include `location` and `source`.
134/// - Only our internal error can be named `source`.
135///   All external error should be `error` with an `#[snafu(source)]` annotation.
136/// - `common_error` crate must be accessible.
137#[proc_macro_attribute]
138pub fn stack_trace_debug(args: TokenStream, input: TokenStream) -> TokenStream {
139    stack_trace_debug::stack_trace_style_impl(args.into(), input.into()).into()
140}
141
142/// Generates implementation for `From<&TableMeta> for TableMetaBuilder`
143#[proc_macro_derive(ToMetaBuilder)]
144pub fn derive_meta_builder(input: TokenStream) -> TokenStream {
145    let input = parse_macro_input!(input as DeriveInput);
146
147    let Data::Struct(data_struct) = input.data else {
148        panic!("ToMetaBuilder can only be derived for structs");
149    };
150
151    let Fields::Named(fields) = data_struct.fields else {
152        panic!("ToMetaBuilder can only be derived for structs with named fields");
153    };
154
155    // Check that this is being applied to TableMeta struct
156    if input.ident != "TableMeta" {
157        panic!("ToMetaBuilder can only be derived for TableMeta struct");
158    }
159
160    let field_init = fields.named.iter().map(|field| {
161        let field_name = field.ident.as_ref().unwrap();
162        quote! {
163            #field_name: Default::default(),
164        }
165    });
166
167    let field_assignments = fields.named.iter().map(|field| {
168        let field_name = field.ident.as_ref().unwrap();
169        quote! {
170            builder.#field_name(meta.#field_name.clone());
171        }
172    });
173
174    let gen = quote! {
175        impl From<&TableMeta> for TableMetaBuilder {
176            fn from(meta: &TableMeta) -> Self {
177                let mut builder = Self {
178                    #(#field_init)*
179                };
180
181                #(#field_assignments)*
182                builder
183            }
184        }
185    };
186
187    gen.into()
188}
common_macro/lib.rs

common_macro/
lib.rs
