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 as_any(&self) -> &dyn std::any::Any {
243                self
244            }
245
246            fn name(&self) -> &str {
247                #display_name
248            }
249
250            fn signature(&self) -> &datafusion_expr::Signature {
251                &self.signature
252            }
253
254            fn return_type(&self, _arg_types: &[datafusion::arrow::datatypes::DataType]) -> datafusion_common::Result<datafusion::arrow::datatypes::DataType> {
255                use datatypes::data_type::DataType;
256                Ok(store_api::storage::ConcreteDataType::#ret().as_arrow_type())
257            }
258
259            fn invoke_with_args(
260                &self,
261                _args: datafusion::logical_expr::ScalarFunctionArgs,
262            ) -> datafusion_common::Result<datafusion_expr::ColumnarValue> {
263                Err(datafusion_common::DataFusionError::NotImplemented(
264                    format!("{} can only be called from async contexts", #display_name)
265                ))
266            }
267        }
268
269        /// Implement From trait for ScalarFunctionFactory
270        impl From<#name> for  #user_path::function_factory::ScalarFunctionFactory {
271            fn from(func: #name) -> Self {
272                 use std::sync::Arc;
273                 use datafusion_expr::ScalarUDFImpl;
274                 use datafusion_expr::async_udf::AsyncScalarUDF;
275
276                let name = func.name().to_string();
277
278                let func = Arc::new(move |ctx: #user_path::function::FunctionContext| {
279                    // create the UDF dynamically with function context
280                    let udf_impl = #name::create(func.signature.clone(), ctx);
281                    let async_udf = AsyncScalarUDF::new(Arc::new(udf_impl));
282                    async_udf.into_scalar_udf()
283                });
284                Self {
285                    name,
286                    factory: func,
287                }
288            }
289        }
290
291        // Implement DataFusion's AsyncScalarUDFImpl trait
292        #[async_trait::async_trait]
293        impl datafusion_expr::async_udf::AsyncScalarUDFImpl for #name {
294            async fn invoke_async_with_args(
295                &self,
296                args: datafusion::logical_expr::ScalarFunctionArgs,
297            ) -> datafusion_common::Result<datafusion_expr::ColumnarValue> {
298                use common_error::ext::ErrorExt;
299
300                let columns = args.args
301                    .iter()
302                    .map(|arg| {
303                        common_query::prelude::ColumnarValue::try_from(arg)
304                            .and_then(|cv| match cv {
305                                common_query::prelude::ColumnarValue::Vector(v) => Ok(v),
306                                common_query::prelude::ColumnarValue::Scalar(s) => {
307                                    datatypes::vectors::Helper::try_from_scalar_value(s, args.number_rows)
308                                        .context(common_query::error::FromScalarValueSnafu)
309                                }
310                            })
311                    })
312                    .collect::<common_query::error::Result<Vec<_>>>()
313                    .map_err(|e| datafusion_common::DataFusionError::Execution(format!("Column conversion error: {}", e.output_msg())))?;
314
315                // Safety check: Ensure under the `greptime` catalog for security
316                #user_path::ensure_greptime!(self.func_ctx);
317
318                let columns_num = columns.len();
319                let rows_num = if columns.is_empty() {
320                    1
321                } else {
322                    columns[0].len()
323                };
324                #validate_rows
325
326                use snafu::{OptionExt, ResultExt};
327                use datatypes::data_type::DataType;
328
329                let query_ctx = &self.func_ctx.query_ctx;
330                let handler = self.func_ctx
331                    .state
332                    .#handler
333                    .as_ref()
334                    .context(#snafu_type)
335                    .map_err(|e| datafusion_common::DataFusionError::Execution(e.output_msg()))?;
336
337                let mut builder = store_api::storage::ConcreteDataType::#ret()
338                    .create_mutable_vector(rows_num);
339
340                if columns_num == 0 {
341                    let result = #fn_name(handler, query_ctx, &[]).await
342                        .map_err(|e| datafusion_common::DataFusionError::Execution(e.output_msg()))?;
343
344                    builder.push_value_ref(&result.as_value_ref());
345                } else {
346                    for i in 0..rows_num {
347                        let args: Vec<_> = columns.iter()
348                            .map(|vector| vector.get_ref(i))
349                            .collect();
350
351                        let result = #fn_name(handler, query_ctx, &args).await
352                            .map_err(|e| datafusion_common::DataFusionError::Execution(e.output_msg()))?;
353
354                        builder.push_value_ref(&result.as_value_ref());
355                    }
356                }
357
358                let result_vector = builder.to_vector();
359
360                // Convert result back to DataFusion ColumnarValue
361                Ok(datafusion_expr::ColumnarValue::Array(result_vector.to_arrow_array()))
362            }
363        }
364
365        impl PartialEq for #name {
366            fn eq(&self, other: &Self) -> bool {
367                self.signature == other.signature
368            }
369        }
370
371        impl Eq for #name {}
372
373        impl std::hash::Hash for #name {
374            fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
375                self.signature.hash(state)
376            }
377        }
378    }
379    .into()
380}
common_macro/admin_fn.rs

common_macro/
admin_fn.rs
