common_macro/
admin_fn.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
15use proc_macro::TokenStream;
16use quote::quote;
17use syn::spanned::Spanned;
18use syn::{
19    Attribute, Ident, ItemFn, Path, Signature, Type, TypePath, TypeReference, Visibility,
20    parse_macro_input,
21};
22
23use crate::utils::extract_input_types;
24
25/// Internal util macro to early return on error.
26macro_rules! ok {
27    ($item:expr) => {
28        match $item {
29            Ok(item) => item,
30            Err(e) => return e.into_compile_error().into(),
31        }
32    };
33}
34
35/// Internal util macro to create an error.
36macro_rules! error {
37    ($span:expr, $msg: expr) => {
38        Err(syn::Error::new($span, $msg))
39    };
40}
41
42pub(crate) fn process_admin_fn(args: TokenStream, input: TokenStream) -> TokenStream {
43    let mut name: Option<Ident> = None;
44    let mut display_name: Option<Ident> = None;
45    let mut sig_fn: Option<Ident> = None;
46    let mut ret: Option<Ident> = None;
47    let mut user_path: Option<Path> = None;
48    let mut single_row = false;
49
50    let parser = syn::meta::parser(|meta| {
51        if meta.path.is_ident("name") {
52            name = Some(meta.value()?.parse()?);
53            Ok(())
54        } else if meta.path.is_ident("display_name") {
55            display_name = Some(meta.value()?.parse()?);
56            Ok(())
57        } else if meta.path.is_ident("sig_fn") {
58            sig_fn = Some(meta.value()?.parse()?);
59            Ok(())
60        } else if meta.path.is_ident("ret") {
61            ret = Some(meta.value()?.parse()?);
62            Ok(())
63        } else if meta.path.is_ident("user_path") {
64            user_path = Some(meta.value()?.parse()?);
65            Ok(())
66        } else if meta.path.is_ident("single_row") {
67            single_row = true;
68            Ok(())
69        } else {
70            Err(meta.error("unsupported property"))
71        }
72    });
73
74    // extract arg map
75    parse_macro_input!(args with parser);
76
77    if user_path.is_none() {
78        user_path = Some(syn::parse_str("crate").expect("failed to parse user path"));
79    }
80
81    // decompose the fn block
82    let compute_fn = parse_macro_input!(input as ItemFn);
83    let ItemFn {
84        attrs,
85        vis,
86        sig,
87        block,
88    } = compute_fn;
89
90    // extract fn arg list
91    let Signature {
92        inputs,
93        ident: fn_name,
94        ..
95    } = &sig;
96
97    let arg_types = ok!(extract_input_types(inputs));
98    if arg_types.len() < 2 {
99        ok!(error!(
100            sig.span(),
101            "Expect at least two argument for admin fn: (handler, query_ctx)"
102        ));
103    }
104    let handler_type = ok!(extract_handler_type(&arg_types));
105
106    let mut result = TokenStream::new();
107    // build the struct and its impl block
108    // only do this when `display_name` is specified
109    if let Some(display_name) = display_name {
110        let struct_code = build_struct(
111            attrs,
112            vis,
113            fn_name,
114            name.expect("name required"),
115            sig_fn.expect("sig_fn required"),
116            ret.expect("ret required"),
117            handler_type,
118            display_name,
119            user_path.expect("user_path required"),
120            single_row,
121        );
122        result.extend(struct_code);
123    }
124
125    // preserve this fn
126    let input_fn_code: TokenStream = quote! {
127        #sig { #block }
128    }
129    .into();
130
131    result.extend(input_fn_code);
132    result
133}
134
135/// Retrieve the handler type, `ProcedureServiceHandlerRef` or `TableMutationHandlerRef`.
136fn extract_handler_type(arg_types: &[Type]) -> Result<&Ident, syn::Error> {
137    match &arg_types[0] {
138        Type::Reference(TypeReference { elem, .. }) => match &**elem {
139            Type::Path(TypePath { path, .. }) => Ok(&path
140                .segments
141                .first()
142                .expect("Expected a reference of handler")
143                .ident),
144            other => {
145                error!(other.span(), "Expected a reference of handler")
146            }
147        },
148        other => {
149            error!(other.span(), "Expected a reference of handler")
150        }
151    }
152}
153
154/// Build the function struct
155#[allow(clippy::too_many_arguments)]
156fn build_struct(
157    attrs: Vec<Attribute>,
158    vis: Visibility,
159    fn_name: &Ident,
160    name: Ident,
161    sig_fn: Ident,
162    ret: Ident,
163    handler_type: &Ident,
164    display_name_ident: Ident,
165    user_path: Path,
166    single_row: bool,
167) -> TokenStream {
168    let display_name = display_name_ident.to_string();
169    let ret = Ident::new(&format!("{ret}_datatype"), ret.span());
170    let uppcase_display_name = display_name.to_uppercase();
171    let validate_rows = single_row.then(|| {
172        quote! {
173            if args.number_rows != 1 {
174                return Err(datafusion_common::DataFusionError::Execution(
175                    format!("{} expects exactly one row, received {}", #display_name, args.number_rows)
176                ));
177            }
178        }
179    });
180    // Get the handler name in function state by the argument ident
181    // TODO(discord9): consider simple depend injection if more handlers are needed
182    let (handler, snafu_type) = match handler_type.to_string().as_str() {
183        "ProcedureServiceHandlerRef" => (
184            Ident::new("procedure_service_handler", handler_type.span()),
185            Ident::new("MissingProcedureServiceHandlerSnafu", handler_type.span()),
186        ),
187
188        "TableMutationHandlerRef" => (
189            Ident::new("table_mutation_handler", handler_type.span()),
190            Ident::new("MissingTableMutationHandlerSnafu", handler_type.span()),
191        ),
192
193        "FlowServiceHandlerRef" => (
194            Ident::new("flow_service_handler", handler_type.span()),
195            Ident::new("MissingFlowServiceHandlerSnafu", handler_type.span()),
196        ),
197        handler => ok!(error!(
198            handler_type.span(),
199            format!("Unknown handler type: {handler}")
200        )),
201    };
202
203    quote! {
204        #(#attrs)*
205        #vis struct #name {
206            signature: datafusion_expr::Signature,
207            func_ctx: #user_path::function::FunctionContext,
208        }
209
210        impl #name {
211            /// Creates a new instance of the function with function context.
212            fn create(signature: datafusion_expr::Signature, func_ctx: #user_path::function::FunctionContext) -> Self {
213                Self {
214                    signature,
215                    func_ctx,
216                }
217            }
218
219            /// Returns the [`ScalarFunctionFactory`] of the function.
220            pub fn factory() -> impl Into< #user_path::function_factory::ScalarFunctionFactory>  {
221                Self {
222                    signature: #sig_fn().into(),
223                    func_ctx: #user_path::function::FunctionContext::default(),
224                }
225            }
226        }
227
228        impl std::fmt::Display for #name {
229            fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
230                write!(f, #uppcase_display_name)
231            }
232        }
233
234        impl std::fmt::Debug for #name {
235            fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
236                write!(f, "{}({})", #uppcase_display_name, self.func_ctx)
237            }
238        }
239
240        // Implement DataFusion's ScalarUDFImpl trait
241        impl datafusion::logical_expr::ScalarUDFImpl for #name {
242            fn name(&self) -> &str {
243                #display_name
244            }
245
246            fn signature(&self) -> &datafusion_expr::Signature {
247                &self.signature
248            }
249
250            fn return_type(&self, _arg_types: &[datafusion::arrow::datatypes::DataType]) -> datafusion_common::Result<datafusion::arrow::datatypes::DataType> {
251                use datatypes::data_type::DataType;
252                Ok(store_api::storage::ConcreteDataType::#ret().as_arrow_type())
253            }
254
255            fn invoke_with_args(
256                &self,
257                _args: datafusion::logical_expr::ScalarFunctionArgs,
258            ) -> datafusion_common::Result<datafusion_expr::ColumnarValue> {
259                Err(datafusion_common::DataFusionError::NotImplemented(
260                    format!("{} can only be called from async contexts", #display_name)
261                ))
262            }
263        }
264
265        /// Implement From trait for ScalarFunctionFactory
266        impl From<#name> for  #user_path::function_factory::ScalarFunctionFactory {
267            fn from(func: #name) -> Self {
268                 use std::sync::Arc;
269                 use datafusion_expr::ScalarUDFImpl;
270                 use datafusion_expr::async_udf::AsyncScalarUDF;
271
272                let name = func.name().to_string();
273
274                let func = Arc::new(move |ctx: #user_path::function::FunctionContext| {
275                    // create the UDF dynamically with function context
276                    let udf_impl = #name::create(func.signature.clone(), ctx);
277                    let async_udf = AsyncScalarUDF::new(Arc::new(udf_impl));
278                    async_udf.into_scalar_udf()
279                });
280                Self {
281                    name,
282                    factory: func,
283                }
284            }
285        }
286
287        // Implement DataFusion's AsyncScalarUDFImpl trait
288        #[async_trait::async_trait]
289        impl datafusion_expr::async_udf::AsyncScalarUDFImpl for #name {
290            async fn invoke_async_with_args(
291                &self,
292                args: datafusion::logical_expr::ScalarFunctionArgs,
293            ) -> datafusion_common::Result<datafusion_expr::ColumnarValue> {
294                use common_error::ext::ErrorExt;
295
296                let columns = args.args
297                    .iter()
298                    .map(|arg| {
299                        common_query::prelude::ColumnarValue::try_from(arg)
300                            .and_then(|cv| match cv {
301                                common_query::prelude::ColumnarValue::Vector(v) => Ok(v),
302                                common_query::prelude::ColumnarValue::Scalar(s) => {
303                                    datatypes::vectors::Helper::try_from_scalar_value(s, args.number_rows, None)
304                                        .context(common_query::error::FromScalarValueSnafu)
305                                }
306                            })
307                    })
308                    .collect::<common_query::error::Result<Vec<_>>>()
309                    .map_err(|e| datafusion_common::DataFusionError::Execution(format!("Column conversion error: {}", e.output_msg())))?;
310
311                // Safety check: Ensure under the `greptime` catalog for security
312                #user_path::ensure_greptime!(self.func_ctx);
313
314                let columns_num = columns.len();
315                let rows_num = if columns.is_empty() {
316                    1
317                } else {
318                    columns[0].len()
319                };
320                #validate_rows
321
322                use snafu::{OptionExt, ResultExt};
323                use datatypes::data_type::DataType;
324
325                let query_ctx = &self.func_ctx.query_ctx;
326                let handler = self.func_ctx
327                    .state
328                    .#handler
329                    .as_ref()
330                    .context(#snafu_type)
331                    .map_err(|e| datafusion_common::DataFusionError::Execution(e.output_msg()))?;
332
333                let mut builder = store_api::storage::ConcreteDataType::#ret()
334                    .create_mutable_vector(rows_num);
335
336                if columns_num == 0 {
337                    let result = #fn_name(handler, query_ctx, &[]).await
338                        .map_err(|e| datafusion_common::DataFusionError::Execution(e.output_msg()))?;
339
340                    builder.push_value_ref(&result.as_value_ref());
341                } else {
342                    for i in 0..rows_num {
343                        let args: Vec<_> = columns.iter()
344                            .map(|vector| vector.get_ref(i))
345                            .collect();
346
347                        let result = #fn_name(handler, query_ctx, &args).await
348                            .map_err(|e| datafusion_common::DataFusionError::Execution(e.output_msg()))?;
349
350                        builder.push_value_ref(&result.as_value_ref());
351                    }
352                }
353
354                let result_vector = builder.to_vector();
355
356                // Convert result back to DataFusion ColumnarValue
357                Ok(datafusion_expr::ColumnarValue::Array(result_vector.to_arrow_array()))
358            }
359        }
360
361        impl PartialEq for #name {
362            fn eq(&self, other: &Self) -> bool {
363                self.signature == other.signature
364            }
365        }
366
367        impl Eq for #name {}
368
369        impl std::hash::Hash for #name {
370            fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
371                self.signature.hash(state)
372            }
373        }
374    }
375    .into()
376}
common_macro/admin_fn.rs

common_macro/
admin_fn.rs
