Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions src/common/function/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,7 @@ license.workspace = true
workspace = true

[features]
testing = []
default = ["geo"]
geo = ["geohash", "h3o", "s2", "wkt", "geo-types", "dep:geo"]

Expand Down
2 changes: 1 addition & 1 deletion src/common/function/src/function.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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(),
Expand Down
2 changes: 1 addition & 1 deletion src/common/function/src/state.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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;

Expand Down
18 changes: 14 additions & 4 deletions src/common/macro/src/admin_fn.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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,
};

Expand Down Expand Up @@ -44,6 +44,7 @@ pub(crate) fn process_admin_fn(args: TokenStream, input: TokenStream) -> TokenSt
let mut display_name: Option<Ident> = None;
let mut sig_fn: Option<Ident> = None;
let mut ret: Option<Ident> = None;
let mut user_path: Option<Path> = None;

let parser = syn::meta::parser(|meta| {
if meta.path.is_ident("name") {
Expand All @@ -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"))
}
Expand All @@ -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 {
Expand Down Expand Up @@ -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);
}
Expand Down Expand Up @@ -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());
Expand Down Expand Up @@ -188,7 +198,7 @@ fn build_struct(


#[async_trait::async_trait]
impl crate::function::AsyncFunction for #name {
impl #user_path::function::AsyncFunction for #name {
Comment thread
WenyXu marked this conversation as resolved.
fn name(&self) -> &'static str {
#display_name
}
Expand All @@ -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<datatypes::vectors::VectorRef> {
async fn eval(&self, func_ctx: #user_path::function::FunctionContext, columns: &[datatypes::vectors::VectorRef]) -> common_query::error::Result<datatypes::vectors::VectorRef> {
// 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() {
Expand Down
4 changes: 2 additions & 2 deletions src/common/macro/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down
Loading