chore(query): improve udf server headers (#17349)

* chore(query): improve udf server headers

* chore(query): improve udf server headers

* chore(query): improve udf server headers

Author: sundy-li

src/common/base/src/headers.rs

@@ -17,6 +17,7 @@ pub const HEADER_QUERY_ID: &str = "X-DATABEND-QUERY-ID";
 pub const HEADER_USER: &str = "X-DATABEND-USER";
 
 pub const HEADER_FUNCTION: &str = "X-DATABEND-FUNCTION";
+pub const HEADER_FUNCTION_HANDLER: &str = "X-DATABEND-FUNCTION-HANDLER";
 
 pub const HEADER_DEDUPLICATE_LABEL: &str = "X-DATABEND-DEDUPLICATE-LABEL";
 pub const HEADER_NODE_ID: &str = "X-DATABEND-NODE-ID";

src/query/expression/src/utils/udf_client.rs

@@ -23,6 +23,7 @@ use arrow_flight::flight_service_client::FlightServiceClient;
 use arrow_flight::FlightDescriptor;
 use arrow_select::concat::concat_batches;
 use databend_common_base::headers::HEADER_FUNCTION;
+use databend_common_base::headers::HEADER_FUNCTION_HANDLER;
 use databend_common_base::headers::HEADER_QUERY_ID;
 use databend_common_base::headers::HEADER_TENANT;
 use databend_common_base::version::DATABEND_SEMVER;
@@ -120,43 +121,38 @@ impl UDFFlightClient {
         })
     }
 
-    /// Set tenant for the UDF client.
-    pub fn with_tenant(mut self, tenant: &str) -> Result<Self> {
-        let key = HEADER_TENANT.to_lowercase();
-        let key = MetadataKey::from_str(key.as_str()).map_err(|err| {
-            ErrorCode::UDFDataError(format!("Parse tenant header key error: {err}"))
-        })?;
-        let value = MetadataValue::from_str(tenant).map_err(|err| {
-            ErrorCode::UDFDataError(format!("Parse tenant header value error: {err}"))
-        })?;
-        self.headers.insert(key, value);
+    pub fn with_headers<'a, H: IntoIterator<Item = (&'a str, &'a str)>>(
+        mut self,
+        headers: H,
+    ) -> Result<Self> {
+        for (key, value) in headers.into_iter() {
+            let key = MetadataKey::from_str(key)
+                .map_err(|err| ErrorCode::UDFDataError(format!("Parse key {key} error: {err}")))?;
+            let value = MetadataValue::from_str(value).map_err(|err| {
+                ErrorCode::UDFDataError(format!("Parse value {value} error: {err}"))
+            })?;
+            self.headers.insert(key, value);
+        }
         Ok(self)
     }
 
+    /// Set tenant for the UDF client.
+    pub fn with_tenant(self, tenant: &str) -> Result<Self> {
+        self.with_headers([(HEADER_TENANT, tenant)])
+    }
+
     /// Set function name for the UDF client.
-    pub fn with_func_name(mut self, func_name: &str) -> Result<Self> {
-        let key = HEADER_FUNCTION.to_lowercase();
-        let key = MetadataKey::from_str(key.as_str()).map_err(|err| {
-            ErrorCode::UDFDataError(format!("Parse function name header key error: {err}"))
-        })?;
-        let value = MetadataValue::from_str(func_name).map_err(|err| {
-            ErrorCode::UDFDataError(format!("Parse function name header value error: {err}"))
-        })?;
-        self.headers.insert(key, value);
-        Ok(self)
+    pub fn with_func_name(self, func_name: &str) -> Result<Self> {
+        self.with_headers([(HEADER_FUNCTION, func_name)])
+    }
+
+    pub fn with_handler_name(self, handler_name: &str) -> Result<Self> {
+        self.with_headers([(HEADER_FUNCTION_HANDLER, handler_name)])
     }
 
     /// Set query id for the UDF client.
-    pub fn with_query_id(mut self, query_id: &str) -> Result<Self> {
-        let key = HEADER_QUERY_ID.to_lowercase();
-        let key = MetadataKey::from_str(key.as_str()).map_err(|err| {
-            ErrorCode::UDFDataError(format!("Parse query id header key error: {err}"))
-        })?;
-        let value = MetadataValue::from_str(query_id).map_err(|err| {
-            ErrorCode::UDFDataError(format!("Parse query id header value error: {err}"))
-        })?;
-        self.headers.insert(key, value);
-        Ok(self)
+    pub fn with_query_id(self, query_id: &str) -> Result<Self> {
+        self.with_headers([(HEADER_QUERY_ID, query_id)])
     }
 
     fn make_request<T>(&self, t: T) -> Request<T> {

src/query/service/src/pipelines/processors/transforms/transform_udf_server.rs

@@ -157,7 +157,8 @@ impl TransformUdfServer {
         let mut client = UDFFlightClient::connect(endpoint, connect_timeout, 65536)
             .await?
             .with_tenant(ctx.get_tenant().tenant_name())?
-            .with_func_name(&func.func_name)?
+            .with_func_name(&func.name)?
+            .with_handler_name(&func.func_name)?
             .with_query_id(&ctx.get_id())?;
 
         let connect_duration = instant.elapsed();

src/query/service/src/table_functions/others/udf.rs

@@ -148,6 +148,7 @@ impl Table for UdfEchoTable {
             .await?
             .with_tenant(ctx.get_tenant().tenant_name())?
             .with_func_name("builtin_echo")?
+            .with_handler_name("builtin_echo")?
             .with_query_id(&ctx.get_id())?;
 
         let array = arrow_array::LargeStringArray::from(vec![self.arg.clone()]);

src/query/sql/src/planner/binder/udf.rs

@@ -106,8 +106,13 @@ impl Binder {
 
                 let endpoint =
                     UDFFlightClient::build_endpoint(address, connect_timeout, request_timeout)?;
-                let mut client =
-                    UDFFlightClient::connect(endpoint, connect_timeout, batch_rows).await?;
+
+                let mut client = UDFFlightClient::connect(endpoint, connect_timeout, batch_rows)
+                    .await?
+                    .with_tenant(self.ctx.get_tenant().tenant_name())?
+                    .with_func_name(&name)?
+                    .with_handler_name(handler)?
+                    .with_query_id(&self.ctx.get_id())?;
                 client
                     .check_schema(handler, &arg_datatypes, &return_type)
                     .await?;