diff --git a/src/common/function/Cargo.toml b/src/common/function/Cargo.toml index fdfd9b4e63a2..7a4c968a3e03 100644 --- a/src/common/function/Cargo.toml +++ b/src/common/function/Cargo.toml @@ -8,6 +8,7 @@ license.workspace = true workspace = true [features] +testing = [] default = ["geo"] geo = ["geohash", "h3o", "s2", "wkt", "geo-types", "dep:geo"] diff --git a/src/common/function/src/function.rs b/src/common/function/src/function.rs index 999361dc1905..196521898339 100644 --- a/src/common/function/src/function.rs +++ b/src/common/function/src/function.rs @@ -32,7 +32,7 @@ pub struct FunctionContext { impl FunctionContext { /// Create a mock [`FunctionContext`] for test. - #[cfg(test)] + #[cfg(any(test, feature = "testing"))] pub fn mock() -> Self { Self { query_ctx: QueryContextBuilder::default().build().into(), diff --git a/src/common/function/src/state.rs b/src/common/function/src/state.rs index a2c63f75ac62..66f5463fa2b5 100644 --- a/src/common/function/src/state.rs +++ b/src/common/function/src/state.rs @@ -28,7 +28,7 @@ pub struct FunctionState { impl FunctionState { /// Create a mock [`FunctionState`] for test. - #[cfg(test)] + #[cfg(any(test, feature = "testing"))] pub fn mock() -> Self { use std::sync::Arc; diff --git a/src/common/macro/src/admin_fn.rs b/src/common/macro/src/admin_fn.rs index b95374540cd3..902bede81496 100644 --- a/src/common/macro/src/admin_fn.rs +++ b/src/common/macro/src/admin_fn.rs @@ -16,7 +16,7 @@ use proc_macro::TokenStream; use quote::quote; use syn::spanned::Spanned; use syn::{ - parse_macro_input, Attribute, Ident, ItemFn, Signature, Type, TypePath, TypeReference, + parse_macro_input, Attribute, Ident, ItemFn, Path, Signature, Type, TypePath, TypeReference, Visibility, }; @@ -44,6 +44,7 @@ pub(crate) fn process_admin_fn(args: TokenStream, input: TokenStream) -> TokenSt let mut display_name: Option = None; let mut sig_fn: Option = None; let mut ret: Option = None; + let mut user_path: Option = None; let parser = syn::meta::parser(|meta| { if meta.path.is_ident("name") { @@ -58,6 +59,9 @@ pub(crate) fn process_admin_fn(args: TokenStream, input: TokenStream) -> TokenSt } else if meta.path.is_ident("ret") { ret = Some(meta.value()?.parse()?); Ok(()) + } else if meta.path.is_ident("user_path") { + user_path = Some(meta.value()?.parse()?); + Ok(()) } else { Err(meta.error("unsupported property")) } @@ -66,6 +70,10 @@ pub(crate) fn process_admin_fn(args: TokenStream, input: TokenStream) -> TokenSt // extract arg map parse_macro_input!(args with parser); + if user_path.is_none() { + user_path = Some(syn::parse_str("crate").expect("failed to parse user path")); + } + // decompose the fn block let compute_fn = parse_macro_input!(input as ItemFn); let ItemFn { @@ -104,6 +112,7 @@ pub(crate) fn process_admin_fn(args: TokenStream, input: TokenStream) -> TokenSt ret.expect("ret required"), handler_type, display_name, + user_path.expect("user_path required"), ); result.extend(struct_code); } @@ -148,6 +157,7 @@ fn build_struct( ret: Ident, handler_type: &Ident, display_name_ident: Ident, + user_path: Path, ) -> TokenStream { let display_name = display_name_ident.to_string(); let ret = Ident::new(&format!("{ret}_datatype"), ret.span()); @@ -188,7 +198,7 @@ fn build_struct( #[async_trait::async_trait] - impl crate::function::AsyncFunction for #name { + impl #user_path::function::AsyncFunction for #name { fn name(&self) -> &'static str { #display_name } @@ -201,9 +211,9 @@ fn build_struct( #sig_fn() } - async fn eval(&self, func_ctx: crate::function::FunctionContext, columns: &[datatypes::vectors::VectorRef]) -> common_query::error::Result { + async fn eval(&self, func_ctx: #user_path::function::FunctionContext, columns: &[datatypes::vectors::VectorRef]) -> common_query::error::Result { // Ensure under the `greptime` catalog for security - crate::ensure_greptime!(func_ctx); + #user_path::ensure_greptime!(func_ctx); let columns_num = columns.len(); let rows_num = if columns.is_empty() { diff --git a/src/common/macro/src/lib.rs b/src/common/macro/src/lib.rs index 29371fb24f35..a6adc3534c06 100644 --- a/src/common/macro/src/lib.rs +++ b/src/common/macro/src/lib.rs @@ -91,8 +91,8 @@ pub fn range_fn(args: TokenStream, input: TokenStream) -> TokenStream { /// - `ret`: The return type of the generated SQL function, it will be transformed into `ConcreteDataType::{ret}_datatype()` result. /// - `display_name`: The display name of the generated SQL function. /// - `sig_fn`: the function to returns `Signature` of generated `Function`. -/// -/// Note that this macro should only be used in `common-function` crate for now +/// - `user_path`: Optional path to the trait and context (e.g., `crate`); +/// defaults to `crate` if not provided. #[proc_macro_attribute] pub fn admin_fn(args: TokenStream, input: TokenStream) -> TokenStream { process_admin_fn(args, input)