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 row;
20mod stack_trace_debug;
21mod utils;
22
23use aggr_func::{impl_aggr_func_type_store, impl_as_aggr_func_creator};
24use print_caller::process_print_caller;
25use proc_macro::TokenStream;
26use quote::quote;
27use range_fn::process_range_fn;
28use syn::{Data, DeriveInput, Fields, parse_macro_input};
29
30use crate::admin_fn::process_admin_fn;
31use crate::row::into_row::derive_into_row_impl;
32use crate::row::schema::derive_schema_impl;
33use crate::row::to_row::derive_to_row_impl;
34
35/// Make struct implemented trait [AggrFuncTypeStore], which is necessary when writing UDAF.
36/// This derive macro is expect to be used along with attribute macro [macro@as_aggr_func_creator].
37#[proc_macro_derive(AggrFuncTypeStore)]
38pub fn aggr_func_type_store_derive(input: TokenStream) -> TokenStream {
39    let ast = parse_macro_input!(input as DeriveInput);
40    impl_aggr_func_type_store(&ast)
41}
42
43/// A struct can be used as a creator for aggregate function if it has been annotated with this
44/// attribute first.
45///
46/// This attribute add a necessary field which is intended to store the input
47/// data's types to the struct.
48/// This attribute is expected to be used along with derive macro [AggrFuncTypeStore].
49#[proc_macro_attribute]
50pub fn as_aggr_func_creator(args: TokenStream, input: TokenStream) -> TokenStream {
51    impl_as_aggr_func_creator(args, input)
52}
53
54/// Attribute macro to convert an arithimetic function to a range function. The annotated function
55/// should accept servaral arrays as input and return a single value as output.
56///
57/// This procedure macro can works on any number of input parameters. Return type can be either
58/// primitive type or wrapped in `Option`.
59///
60/// # Example
61/// Take `count_over_time()` in PromQL as an example:
62/// ```rust, ignore
63/// /// The count of all values in the specified interval.
64/// #[range_fn(
65///     name = "CountOverTime",
66///     ret = "Float64Array",
67///     display_name = "prom_count_over_time"
68/// )]
69/// pub fn count_over_time(_: &TimestampMillisecondArray, values: &Float64Array) -> f64 {
70///      values.len() as f64
71/// }
72/// ```
73///
74/// # Arguments
75/// - `name`: The name of the generated `ScalarUDF` struct.
76/// - `ret`: The return type of the generated UDF function.
77/// - `display_name`: The display name of the generated UDF function.
78/// - `evaluator`: Optional path to a specialized evaluator with the calling convention
79///   `fn(&[ColumnarValue], &str) -> Result<ColumnarValue, DataFusionError>`. When supplied,
80///   the generated UDF `calc` delegates directly to it; without it, the default expansion is
81///   unchanged.
82#[proc_macro_attribute]
83pub fn range_fn(args: TokenStream, input: TokenStream) -> TokenStream {
84    process_range_fn(args, input)
85}
86
87/// Attribute macro to convert a normal function to SQL administration function. The annotated function
88/// should accept:
89///    - `&ProcedureServiceHandlerRef` or `&TableMutationHandlerRef` or `FlowServiceHandlerRef` as the first argument,
90///    - `&QueryContextRef` as the second argument, and
91///    - `&[ValueRef<'_>]` as the third argument which is SQL function input values in each row.
92///
93/// Return type must be `common_query::error::Result<Value>`.
94///
95/// # Example see `common/function/src/system/procedure_state.rs`.
96///
97/// # Arguments
98/// - `name`: The name of the generated `Function` implementation.
99/// - `ret`: The return type of the generated SQL function, it will be transformed into `ConcreteDataType::{ret}_datatype()` result.
100/// - `display_name`: The display name of the generated SQL function.
101/// - `sig_fn`: the function to returns `Signature` of generated `Function`.
102/// - `user_path`: Optional path to the trait and context (e.g., `crate`);
103///   defaults to `crate` if not provided.
104#[proc_macro_attribute]
105pub fn admin_fn(args: TokenStream, input: TokenStream) -> TokenStream {
106    process_admin_fn(args, input)
107}
108
109/// Attribute macro to print the caller to the annotated function.
110/// The caller is printed as its filename and the call site line number.
111///
112/// This macro works like this: inject the tracking codes as the first statement to the annotated
113/// function body. The tracking codes use [backtrace-rs](https://crates.io/crates/backtrace) to get
114/// the callers. So you must dependent on the `backtrace-rs` crate.
115///
116/// # Arguments
117/// - `depth`: The max depth of call stack to print. Optional, defaults to 1.
118///
119/// # Example
120/// ```rust, ignore
121///
122/// #[print_caller(depth = 3)]
123/// fn foo() {}
124/// ```
125#[proc_macro_attribute]
126pub fn print_caller(args: TokenStream, input: TokenStream) -> TokenStream {
127    process_print_caller(args, input)
128}
129
130/// Attribute macro to derive [std::fmt::Debug] for the annotated `Error` type.
131///
132/// The generated `Debug` implementation will print the error in a stack trace style. E.g.:
133/// ```plaintext
134/// 0: Foo error, at src/common/catalog/src/error.rs:80:10
135/// 1: Bar error, at src/common/function/src/error.rs:90:10
136/// 2: Root cause, invalid table name, at src/common/catalog/src/error.rs:100:10
137/// ```
138///
139/// Notes on using this macro:
140/// - `#[snafu(display)]` must present on each enum variants,
141///   and should not include `location` and `source`.
142/// - Only our internal error can be named `source`.
143///   All external error should be `error` with an `#[snafu(source)]` annotation.
144/// - `common_error` crate must be accessible.
145#[proc_macro_attribute]
146pub fn stack_trace_debug(args: TokenStream, input: TokenStream) -> TokenStream {
147    stack_trace_debug::stack_trace_style_impl(args.into(), input.into()).into()
148}
149
150/// Generates implementation for `From<&TableMeta> for TableMetaBuilder`
151#[proc_macro_derive(ToMetaBuilder)]
152pub fn derive_meta_builder(input: TokenStream) -> TokenStream {
153    let input = parse_macro_input!(input as DeriveInput);
154
155    let Data::Struct(data_struct) = input.data else {
156        panic!("ToMetaBuilder can only be derived for structs");
157    };
158
159    let Fields::Named(fields) = data_struct.fields else {
160        panic!("ToMetaBuilder can only be derived for structs with named fields");
161    };
162
163    // Check that this is being applied to TableMeta struct
164    if input.ident != "TableMeta" {
165        panic!("ToMetaBuilder can only be derived for TableMeta struct");
166    }
167
168    let field_init = fields.named.iter().map(|field| {
169        let field_name = field.ident.as_ref().unwrap();
170        quote! {
171            #field_name: Default::default(),
172        }
173    });
174
175    let field_assignments = fields.named.iter().map(|field| {
176        let field_name = field.ident.as_ref().unwrap();
177        quote! {
178            builder.#field_name(meta.#field_name.clone());
179        }
180    });
181
182    let generated = quote! {
183        impl From<&TableMeta> for TableMetaBuilder {
184            fn from(meta: &TableMeta) -> Self {
185                let mut builder = Self {
186                    #(#field_init)*
187                };
188
189                #(#field_assignments)*
190                builder
191            }
192        }
193    };
194
195    generated.into()
196}
197
198/// Derive macro to convert a struct to a row.
199///
200/// # Example
201/// ```rust, ignore
202/// use api::v1::Row;
203/// use api::v1::value::ValueData;
204/// use api::v1::Value;
205///
206/// #[derive(ToRow)]
207/// struct ToRowTest {
208///     my_value: i32,
209///     #[col(name = "string_value", datatype = "string", semantic = "tag")]
210///     my_string: String,  
211///     my_bool: bool,
212///     my_float: f32,
213///     #[col(
214///         name = "timestamp_value",
215///         semantic = "Timestamp",
216///         datatype = "TimestampMillisecond"
217///     )]
218///     my_timestamp: i64,
219///     #[col(skip)]
220///     my_skip: i32,
221/// }
222///
223/// let row = ToRowTest {
224///     my_value: 1,
225///     my_string: "test".to_string(),
226///     my_bool: true,
227///     my_float: 1.0,
228///     my_timestamp: 1718563200000,
229///     my_skip: 1,
230/// }.to_row();
231/// ```
232#[proc_macro_derive(ToRow, attributes(col))]
233pub fn derive_to_row(input: TokenStream) -> TokenStream {
234    let input = parse_macro_input!(input as DeriveInput);
235    let output = derive_to_row_impl(input);
236    output.unwrap_or_else(|e| e.to_compile_error()).into()
237}
238
239/// Derive macro to convert a struct to a row with move semantics.
240///
241/// # Example
242/// ```rust, ignore
243/// use api::v1::Row;
244/// use api::v1::value::ValueData;
245/// use api::v1::Value;
246///
247/// #[derive(IntoRow)]
248/// struct IntoRowTest {
249///     my_value: i32,
250///     #[col(name = "string_value", datatype = "string", semantic = "tag")]
251///     my_string: String,  
252///     my_bool: bool,
253///     my_float: f32,
254///     #[col(
255///         name = "timestamp_value",
256///         semantic = "Timestamp",
257///         datatype = "TimestampMillisecond"
258///     )]
259///     my_timestamp: i64,
260///     #[col(skip)]
261///     my_skip: i32,
262/// }
263///
264/// let row = IntoRowTest {
265///     my_value: 1,
266///     my_string: "test".to_string(),
267///     my_bool: true,
268///     my_float: 1.0,
269///     my_timestamp: 1718563200000,
270///     my_skip: 1,
271/// }.into_row();
272/// ```
273#[proc_macro_derive(IntoRow, attributes(col))]
274pub fn derive_into_row(input: TokenStream) -> TokenStream {
275    let input = parse_macro_input!(input as DeriveInput);
276    let output = derive_into_row_impl(input);
277    output.unwrap_or_else(|e| e.to_compile_error()).into()
278}
279
280/// Derive macro to convert a struct to a schema.
281///
282/// # Example
283/// ```rust, ignore
284/// use api::v1::ColumnSchema;
285///
286/// #[derive(Schema)]
287/// struct SchemaTest {
288///     my_value: i32,
289///     #[col(name = "string_value", datatype = "string", semantic = "tag")]
290///     my_string: String,  
291///     my_bool: bool,
292///     my_float: f32,
293///     #[col(
294///         name = "timestamp_value",
295///         semantic = "Timestamp",
296///         datatype = "TimestampMillisecond"
297///     )]
298///     my_timestamp: i64,
299///     #[col(skip)]
300///     my_skip: i32,
301/// }
302///
303/// let schema = SchemaTest::schema();
304/// ```
305#[proc_macro_derive(Schema, attributes(col))]
306pub fn derive_schema(input: TokenStream) -> TokenStream {
307    let input = parse_macro_input!(input as DeriveInput);
308    let output = derive_schema_impl(input);
309    output.unwrap_or_else(|e| e.to_compile_error()).into()
310}
common_macro/lib.rs

common_macro/
lib.rs
