feat: add server to admin udfs api (#17141)

* feat: add server to admin udfs api * z
databendlabs · Dec 30, 2024 · 6560bf5 · 6560bf5
1 parent 9e71e4d
commit 6560bf5
Show file tree

Hide file tree

Showing 3 changed files with 48 additions and 20 deletions.
diff --git a/src/meta/app/src/principal/user_defined_function.rs b/src/meta/app/src/principal/user_defined_function.rs
@@ -65,6 +65,35 @@ pub enum UDFDefinition {
     UDAFScript(UDAFScript),
 }
 
+impl UDFDefinition {
+    pub fn category(&self) -> &str {
+        match self {
+            Self::LambdaUDF(_) => "LambdaUDF",
+            Self::UDFServer(_) => "UDFServer",
+            Self::UDFScript(_) => "UDFScript",
+            Self::UDAFScript(_) => "UDAFScript",
+        }
+    }
+
+    pub fn is_aggregate(&self) -> bool {
+        match self {
+            Self::LambdaUDF(_) => false,
+            Self::UDFServer(_) => false,
+            Self::UDFScript(_) => false,
+            Self::UDAFScript(_) => true,
+        }
+    }
+
+    pub fn language(&self) -> &str {
+        match self {
+            Self::LambdaUDF(_) => "SQL",
+            Self::UDFServer(x) => x.language.as_str(),
+            Self::UDFScript(x) => x.language.as_str(),
+            Self::UDAFScript(x) => x.language.as_str(),
+        }
+    }
+}
+
 #[derive(Clone, Debug, Eq, PartialEq)]
 pub struct UserDefinedFunction {
     pub name: String,

diff --git a/src/query/storages/system/src/user_functions_table.rs b/src/query/storages/system/src/user_functions_table.rs
@@ -106,6 +106,8 @@ pub struct UserFunctionArguments {
     arg_types: Vec<String>,
     #[serde(skip_serializing_if = "std::option::Option::is_none")]
     return_type: Option<String>,
+    #[serde(skip_serializing_if = "std::option::Option::is_none")]
+    server: Option<String>,
     #[serde(skip_serializing_if = "std::vec::Vec::is_empty")]
     parameters: Vec<String>,
     #[serde(skip_serializing_if = "std::collections::BTreeMap::is_empty")]
@@ -118,6 +120,7 @@ pub struct UserFunction {
     is_aggregate: bool,
     description: String,
     language: String,
+    category: String,
     definition: String,
     created_on: DateTime<Utc>,
     arguments: UserFunctionArguments,
@@ -162,43 +165,39 @@ impl UserFunctionsTable {
             .into_iter()
             .map(|user_function| UserFunction {
                 name: user_function.name,
-                is_aggregate: match user_function.definition {
-                    UDFDefinition::LambdaUDF(_) => false,
-                    UDFDefinition::UDFServer(_) | UDFDefinition::UDFScript(_) => false,
-                    UDFDefinition::UDAFScript(_) => true,
-                },
+                is_aggregate: user_function.definition.is_aggregate(),
                 description: user_function.description,
-                language: match &user_function.definition {
-                    UDFDefinition::LambdaUDF(_) => String::from("SQL"),
-                    UDFDefinition::UDFServer(x) => x.language.clone(),
-                    UDFDefinition::UDFScript(x) => x.language.to_string(),
-                    UDFDefinition::UDAFScript(x) => x.language.to_string(),
-                },
+                language: user_function.definition.language().to_string(),
+                category: user_function.definition.category().to_string(),
                 definition: user_function.definition.to_string(),
                 created_on: user_function.created_on,
                 arguments: match &user_function.definition {
                     UDFDefinition::LambdaUDF(x) => UserFunctionArguments {
-                        return_type: None,
                         arg_types: vec![],
+                        return_type: None,
+                        server: None,
                         parameters: x.parameters.clone(),
                         states: BTreeMap::new(),
                     },
                     UDFDefinition::UDFServer(x) => UserFunctionArguments {
-                        parameters: vec![],
-                        return_type: Some(x.return_type.to_string()),
                         arg_types: x.arg_types.iter().map(ToString::to_string).collect(),
+                        return_type: Some(x.return_type.to_string()),
+                        server: Some(x.address.to_string()),
+                        parameters: vec![],
                         states: BTreeMap::new(),
                     },
                     UDFDefinition::UDFScript(x) => UserFunctionArguments {
-                        parameters: vec![],
-                        return_type: Some(x.return_type.to_string()),
                         arg_types: x.arg_types.iter().map(ToString::to_string).collect(),
+                        return_type: Some(x.return_type.to_string()),
+                        server: None,
+                        parameters: vec![],
                         states: BTreeMap::new(),
                     },
                     UDFDefinition::UDAFScript(x) => UserFunctionArguments {
-                        parameters: vec![],
-                        return_type: Some(x.return_type.to_string()),
                         arg_types: x.arg_types.iter().map(ToString::to_string).collect(),
+                        return_type: Some(x.return_type.to_string()),
+                        server: None,
+                        parameters: vec![],
                         states: x
                             .state_fields
                             .iter()

diff --git a/tests/suites/0_stateless/20+_others/20_0016_udf_timestamp.result b/tests/suites/0_stateless/20+_others/20_0016_udf_timestamp.result
@@ -1,6 +1,6 @@
 ==TEST SHOW USER FUNCTIONS==
 isnotempty	0		{"parameters":["p"]}	SQL	yyyy-mm-dd HH:MM:SS.ssssss
-ping	0	Built-in UDF	{"arg_types":["String NULL"],"return_type":"String NULL"}	python	yyyy-mm-dd HH:MM:SS.ssssss
+ping	0	Built-in UDF	{"arg_types":["String NULL"],"return_type":"String NULL","server":"http://0.0.0.0:8815"}	python	yyyy-mm-dd HH:MM:SS.ssssss
 ==TEST SELECT * FROM SYSTEM.USER_FUNCTIONS==
 isnotempty	0		{"parameters":["p"]}	SQL	 (p) -> NOT is_null(p)	yyyy-mm-dd HH:MM:SS.ssssss
-ping	0	Built-in UDF	{"arg_types":["String NULL"],"return_type":"String NULL"}	python	 (String NULL) RETURNS String NULL LANGUAGE python HANDLER = ping ADDRESS = http://0.0.0.0:8815	yyyy-mm-dd HH:MM:SS.ssssss
+ping	0	Built-in UDF	{"arg_types":["String NULL"],"return_type":"String NULL","server":"http://0.0.0.0:8815"}	python	 (String NULL) RETURNS String NULL LANGUAGE python HANDLER = ping ADDRESS = http://0.0.0.0:8815	yyyy-mm-dd HH:MM:SS.ssssss