chore: add udf compat

databendlabs · Jan 7, 2025 · 4ed2fff · 4ed2fff
1 parent 99cb6e1
commit 4ed2fff
Show file tree

Hide file tree

Showing 5 changed files with 63 additions and 4 deletions.
diff --git a/.github/actions/test_compat_fuse/action.yml b/.github/actions/test_compat_fuse/action.yml
@@ -17,6 +17,13 @@ runs:
         path: ./bins/current
         artifacts: sqllogictests,meta,query
 
+    - name: Start UDF Server
+      shell: bash
+      run: |
+        pip install databend-udf>=0.2.6
+        python3 tests/udf/udf_server.py &
+        sleep 3
+
     - name: Test compatibility
       shell: bash
       # test-*.sh <old-query-ver> <meta-ver> <test-suite>
@@ -31,6 +38,8 @@ runs:
           bash ./tests/compat_fuse/test_compat_fuse.sh         1.2.318 1.2.527 rbac
           bash ./tests/compat_fuse/test_compat_fuse_forward.sh 1.2.307 1.2.527 rbac
           bash ./tests/compat_fuse/test_compat_fuse_forward.sh 1.2.318 1.2.527 rbac
+          bash ./tests/compat_fuse/test_compat_fuse.sh         1.2.527 1.2.680 udf
+          bash ./tests/compat_fuse/test_compat_fuse_forward.sh 1.2.527 1.2.680 udf
     - name: Upload failure
       if: failure()
       uses: ./.github/actions/artifact_failure

diff --git a/src/meta/api/src/kv_pb_api/mod.rs b/src/meta/api/src/kv_pb_api/mod.rs
@@ -41,6 +41,7 @@ use futures::stream::BoxStream;
 use futures::stream::StreamExt;
 use futures::TryStreamExt;
 use itertools::Itertools;
+use log::error;
 
 pub(crate) use self::codec::decode_non_empty_item;
 pub(crate) use self::codec::decode_seqv;
@@ -446,8 +447,31 @@ pub trait KVPbApi: KVApi {
         K::ValueType: FromToProto,
         Self::Error: From<PbApiReadError<Self::Error>>,
     {
-        self.list_pb(prefix)
-            .map_ok(|strm| strm.map_ok(|x| x.seqv.data).boxed())
+        let prefix_str = format!("{:?}", prefix);
+        let fut = self.list_pb(prefix)
+            .map_ok(move |strm| {
+                strm.map_ok(move |x| {
+                    let seqv_data_result =  K::ValueType::from_proto(&x.seqv.data);
+                    match seqv_data_result {
+                        Ok(data) => Ok(data),
+                        Err(e) =>{
+                            error!("Failed to convert seqv.data {:?} with seq: {:?}, error: {:?}", x.seqv.data, x.seqv.seq, e);
+                            Err(e.into())
+                        }
+                    }
+                }).boxed()
+            });
+
+        let fut = async move {
+            match fut.await {
+                Ok(stream) => Ok(stream),
+                Err(e) => {
+                    error!("Failed to list values with prefix: {}, error: {:?}",prefix_str, e);
+                    Err(e)
+                }
+            }
+        };
+        fut
     }
 
     /// List protobuf encoded values by prefix and returns a stream.

diff --git a/src/query/management/src/user/user_mgr.rs b/src/query/management/src/user/user_mgr.rs
@@ -144,8 +144,15 @@ impl UserApi for UserMgr {
         let values = self.get_raw_users().await?;
         let mut r = vec![];
         for (_key, val) in values {
-            let u = deserialize_struct(&val.data, ErrorCode::IllegalUserInfoFormat, || "")?;
-            r.push(SeqV::new(val.seq, u));
+            match deserialize_struct(&val.data, ErrorCode::IllegalUserInfoFormat, || "") {
+                Ok(u) => r.push(SeqV::new(val.seq, u)),
+                Err(e) => {
+                    Err(ErrorCode::UnknownUser(format!(
+                        "user {} with invalid info format: {}", val.data
+                        e
+                    )))
+                }
+            }
         }
 
         Ok(r)

diff --git a/tests/compat_fuse/compat-logictest/udf/fuse_compat_read b/tests/compat_fuse/compat-logictest/udf/fuse_compat_read
@@ -0,0 +1,9 @@
+query
+SELECT definition FROM SYSTEM.USER_FUNCTIONS ORDER BY name;
+----
+(Int8 NULL, Int16 NULL, Int32 NULL, Int64 NULL) RETURNS Int64 NULL LANGUAGE python HANDLER = add_signed ADDRESS = http://0.0.0.0:8815
+
+statement ok
+call add_signed(1,1,1,1);
+----
+4
diff --git a/tests/compat_fuse/compat-logictest/udf/fuse_compat_write b/tests/compat_fuse/compat-logictest/udf/fuse_compat_write
@@ -0,0 +1,10 @@
+statement ok
+DROP FUNCTION IF EXISTS add_signed;
+
+statement ok
+CREATE FUNCTION add_signed (TINYINT, SMALLINT, INT, BIGINT) RETURNS BIGINT LANGUAGE python HANDLER = 'add_signed' ADDRESS = 'http://0.0.0.0:8815'
+
+query
+SELECT definition FROM SYSTEM.USER_FUNCTIONS ORDER BY name;
+----
+(Int8 NULL, Int16 NULL, Int32 NULL, Int64 NULL) RETURNS Int64 NULL LANGUAGE python HANDLER = add_signed ADDRESS = http://0.0.0.0:8815