fix(query): add udf name(in meta) into UdfServerCall

TCeason · TCeason · commit ef0cb55c0d7a · 2023-12-22T21:06:51.000+08:00
chore: fix review suggestion in https://github.com/datafuselabs/databend/pull/14098\#discussion_r1433048221
diff --git a/.github/actions/test_stateless_standalone_linux/action.yml b/.github/actions/test_stateless_standalone_linux/action.yml
@@ -8,6 +8,14 @@ inputs:
 runs:
   using: "composite"
   steps:
+    - uses: actions/checkout@v3
+    - name: Start UDF Server
+      shell: bash
+      run: |
+        pip install databend-udf
+        python3 tests/udf/udf_server.py &
+        sleep 2
+
     - uses: ./.github/actions/setup_bendsql
 
     - name: Download artifact
diff --git a/src/query/sql/src/planner/binder/ddl/table.rs b/src/query/sql/src/planner/binder/ddl/table.rs
@@ -992,23 +992,24 @@ impl Binder {
         } = stmt;
 
         let tenant = self.ctx.get_tenant();
-        let (catalog, database, table) =
+        let (catalog, db_name, table) =
             self.normalize_object_identifier_triple(catalog, database, table);
 
         let (new_catalog, new_database, new_table) =
             self.normalize_object_identifier_triple(new_catalog, new_database, new_table);
 
-        if new_catalog != catalog || new_database != database {
+        if new_catalog != catalog || new_database != db_name {
             return Err(ErrorCode::BadArguments(
                 "Rename table not allow modify catalog or database",
-            ));
+            )
+            .set_span(database.as_ref().and_then(|ident| ident.span)));
         }
 
         Ok(Plan::RenameTable(Box::new(RenameTablePlan {
             tenant,
             if_exists: *if_exists,
             catalog,
-            database,
+            database: db_name,
             table,
             new_database,
             new_table,
diff --git a/src/query/sql/src/planner/binder/sort.rs b/src/query/sql/src/planner/binder/sort.rs
@@ -392,6 +392,7 @@ impl Binder {
                         .collect::<Result<Vec<_>>>()?;
                     Ok(UDFServerCall {
                         span: udf.span,
+                        name: udf.name.clone(),
                         func_name: udf.func_name.clone(),
                         display_name: udf.display_name.clone(),
                         server_addr: udf.server_addr.clone(),
diff --git a/src/query/sql/src/planner/optimizer/decorrelate/flatten_scalar.rs b/src/query/sql/src/planner/optimizer/decorrelate/flatten_scalar.rs
@@ -96,6 +96,7 @@ impl SubqueryRewriter {
                     .collect::<Result<Vec<_>>>()?;
                 Ok(ScalarExpr::UDFServerCall(UDFServerCall {
                     span: udf.span,
+                    name: udf.name.clone(),
                     func_name: udf.func_name.clone(),
                     display_name: udf.display_name.clone(),
                     server_addr: udf.server_addr.clone(),
diff --git a/src/query/sql/src/planner/optimizer/decorrelate/subquery_rewriter.rs b/src/query/sql/src/planner/optimizer/decorrelate/subquery_rewriter.rs
@@ -353,6 +353,7 @@ impl SubqueryRewriter {
 
                 let expr: ScalarExpr = UDFServerCall {
                     span: udf.span,
+                    name: udf.name.clone(),
                     func_name: udf.func_name.clone(),
                     display_name: udf.display_name.clone(),
                     server_addr: udf.server_addr.clone(),
diff --git a/src/query/sql/src/planner/optimizer/rule/rewrite/agg_index/query_rewrite.rs b/src/query/sql/src/planner/optimizer/rule/rewrite/agg_index/query_rewrite.rs
@@ -1012,6 +1012,7 @@ fn rewrite_query_item(
             Some(
                 UDFServerCall {
                     span: udf.span,
+                    name: udf.name.clone(),
                     func_name: udf.func_name.clone(),
                     display_name: udf.display_name.clone(),
                     server_addr: udf.server_addr.clone(),
diff --git a/src/query/sql/src/planner/optimizer/rule/rewrite/rule_push_down_filter_eval_scalar.rs b/src/query/sql/src/planner/optimizer/rule/rewrite/rule_push_down_filter_eval_scalar.rs
@@ -229,6 +229,7 @@ impl RulePushDownFilterEvalScalar {
 
                 Ok(ScalarExpr::UDFServerCall(UDFServerCall {
                     span: udf.span,
+                    name: udf.name.clone(),
                     func_name: udf.func_name.clone(),
                     display_name: udf.display_name.clone(),
                     server_addr: udf.server_addr.clone(),
diff --git a/src/query/sql/src/planner/optimizer/rule/rewrite/rule_push_down_filter_scan.rs b/src/query/sql/src/planner/optimizer/rule/rewrite/rule_push_down_filter_scan.rs
@@ -260,6 +260,7 @@ impl RulePushDownFilterScan {
 
                 Ok(ScalarExpr::UDFServerCall(UDFServerCall {
                     span: udf.span,
+                    name: udf.name.clone(),
                     func_name: udf.func_name.clone(),
                     display_name: udf.display_name.clone(),
                     server_addr: udf.server_addr.clone(),
diff --git a/src/query/sql/src/planner/optimizer/s_expr.rs b/src/query/sql/src/planner/optimizer/s_expr.rs
@@ -504,7 +504,7 @@ pub fn get_udf_names(scalar: &ScalarExpr) -> Result<HashSet<&String>> {
                 self.visit(expr)?;
             }
 
-            self.udfs.insert(&udf.func_name);
+            self.udfs.insert(&udf.name);
             Ok(())
         }
 
diff --git a/src/query/sql/src/planner/plans/scalar_expr.rs b/src/query/sql/src/planner/plans/scalar_expr.rs
@@ -589,6 +589,9 @@ fn hash_column_set<H: Hasher>(columns: &ColumnSet, state: &mut H) {
 pub struct UDFServerCall {
     #[educe(Hash(ignore), PartialEq(ignore), Eq(ignore))]
     pub span: Span,
+    // name in meta
+    pub name: String,
+    // name in handler
     pub func_name: String,
     pub display_name: String,
     pub server_addr: String,
diff --git a/src/query/sql/src/planner/semantic/type_check.rs b/src/query/sql/src/planner/semantic/type_check.rs
@@ -2796,7 +2796,8 @@ impl<'a> TypeChecker<'a> {
                     .await?,
             )),
             UDFDefinition::UDFServer(udf_def) => Ok(Some(
-                self.resolve_udf_server(span, arguments, udf_def).await?,
+                self.resolve_udf_server(span, name, arguments, udf_def)
+                    .await?,
             )),
         }
     }
@@ -2806,6 +2807,7 @@ impl<'a> TypeChecker<'a> {
     async fn resolve_udf_server(
         &mut self,
         span: Span,
+        name: String,
         arguments: &[Expr],
         udf_definition: UDFServer,
     ) -> Result<Box<(ScalarExpr, DataType)>> {
@@ -2852,6 +2854,7 @@ impl<'a> TypeChecker<'a> {
         Ok(Box::new((
             UDFServerCall {
                 span,
+                name,
                 func_name: udf_definition.handler,
                 display_name,
                 server_addr: udf_definition.address,
diff --git a/tests/suites/0_stateless/18_rbac/18_0001_udf_priv.result b/tests/suites/0_stateless/18_rbac/18_0001_udf_priv.result
@@ -54,3 +54,5 @@ Error: APIError: ResponseError with 1063: Permission denied, privilege [Usage] i
 200
 100
 200
+4
+Error: APIError: ResponseError with 1063: Permission denied, privilege [Usage] is required on UDF b for user 'test-user'@'%' with roles [public]
diff --git a/tests/suites/0_stateless/18_rbac/18_0001_udf_priv.sh b/tests/suites/0_stateless/18_rbac/18_0001_udf_priv.sh
@@ -123,9 +123,22 @@ echo "select case when i > 100 then 200 else f2(f1(100)) end as c1 from t;" | $T
 echo "select case when i > 100 then f2(f1(200)) else 100 end as c1 from t;" | $TEST_USER_CONNECT
 echo "delete from t;" | $TEST_USER_CONNECT
 
+#udf server test
+echo "drop function if exists a;" | $BENDSQL_CLIENT_CONNECT
+echo "drop function if exists b;" | $BENDSQL_CLIENT_CONNECT
+echo "CREATE FUNCTION a (TINYINT, SMALLINT, INT, BIGINT) RETURNS BIGINT LANGUAGE python HANDLER = 'add_signed' ADDRESS = 'http://0.0.0.0:8815';" | $BENDSQL_CLIENT_CONNECT
+echo "CREATE FUNCTION b (TINYINT, SMALLINT, INT, BIGINT) RETURNS BIGINT LANGUAGE python HANDLER = 'add_signed' ADDRESS = 'http://0.0.0.0:8815';" | $BENDSQL_CLIENT_CONNECT
+
+echo "grant usage on udf a to 'test-user'" | $BENDSQL_CLIENT_CONNECT
+echo "select a(1,1,1,1)" | $TEST_USER_CONNECT
+echo "select b(1,1,1,1)" | $TEST_USER_CONNECT
+
+
 echo "drop user if exists 'test-user'" | $BENDSQL_CLIENT_CONNECT
 echo "DROP FUNCTION IF EXISTS f1;" |  $BENDSQL_CLIENT_CONNECT
 echo "DROP FUNCTION IF EXISTS f2;" |  $BENDSQL_CLIENT_CONNECT
+echo "drop function if exists a;" | $BENDSQL_CLIENT_CONNECT
+echo "drop function if exists b;" | $BENDSQL_CLIENT_CONNECT
 echo "drop table if exists default.t;" | $BENDSQL_CLIENT_CONNECT
 echo "drop table if exists default.t2;" | $BENDSQL_CLIENT_CONNECT
 echo "unset enable_experimental_rbac_check" | $BENDSQL_CLIENT_CONNECT

Original file line number	Diff line number	Diff line change
`@@ -504,7 +504,7 @@ pub fn get_udf_names(scalar: &ScalarExpr) -> Result<HashSet<&String>> {`
`504`	`504`	`self.visit(expr)?;`
`505`	`505`	`}`
`506`	`506`
`507`		`- self.udfs.insert(&udf.func_name);`
	`507`	`+ self.udfs.insert(&udf.name);`
`508`	`508`	`Ok(())`
`509`	`509`	`}`
`510`	`510`
Original file line number	Diff line number	Diff line change
`@@ -54,3 +54,5 @@ Error: APIError: ResponseError with 1063: Permission denied, privilege [Usage] i`
`54`	`54`	`200`
`55`	`55`	`100`
`56`	`56`	`200`
	`57`	`+4`
	`58`	`+Error: APIError: ResponseError with 1063: Permission denied, privilege [Usage] is required on UDF b for user 'test-user'@'%' with roles [public]`