Add escape for management key

common/base/src/lib.rs

@@ -23,6 +23,7 @@ mod runtime_tracker;
 mod shutdown_signal;
 mod stop_handle;
 mod stoppable;
+mod string_func;
 mod thread;
 mod uniq_id;
 
@@ -43,6 +44,8 @@ pub use shutdown_signal::SignalStream;
 pub use shutdown_signal::SignalType;
 pub use stop_handle::StopHandle;
 pub use stoppable::Stoppable;
+pub use string_func::escape_for_key;
+pub use string_func::unescape_for_key;
 pub use thread::Thread;
 pub use tokio;
 pub use uniq_id::GlobalSequence;

common/base/src/string_func.rs

@@ -0,0 +1,93 @@
+// Copyright 2022 Datafuse Labs.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+use std::string::FromUtf8Error;
+
+/// Function that escapes special characters in a string.
+///
+/// All characters except digit, alphabet and '_' are treated as special characters.
+/// A special character will be converted into "%num" where num is the hexadecimal form of the character.
+///
+/// # Example
+/// ```
+/// let key = "data_bend!!";
+/// let new_key = escape_for_key(&key);
+/// assert_eq!(Ok("data_bend%21%21".to_string()), new_key);
+/// ```
+pub fn escape_for_key(key: &str) -> Result<String, FromUtf8Error> {
+    let mut new_key = Vec::with_capacity(key.len());
+
+    fn hex(num: u8) -> u8 {
+        match num {
+            0..=9 => b'0' + num,
+            10..=15 => b'a' + (num - 10),
+            unreachable => unreachable!("Unreachable branch num = {}", unreachable),
+        }
+    }
+
+    for char in key.as_bytes() {
+        match char {
+            b'0'..=b'9' => new_key.push(*char),
+            b'_' | b'a'..=b'z' | b'A'..=b'Z' => new_key.push(*char),
+            _other => {
+                new_key.push(b'%');
+                new_key.push(hex(*char / 16));
+                new_key.push(hex(*char % 16));
+            }
+        }
+    }
+
+    String::from_utf8(new_key)
+}
+
+/// The reverse function of escape_for_key.
+///
+/// # Example
+/// ```
+/// let key = "data_bend%21%21";
+/// let original_key = unescape_for_key(&key);
+/// assert_eq!(Ok("data_bend!!".to_string()), original_key);
+/// ```
+pub fn unescape_for_key(key: &str) -> Result<String, FromUtf8Error> {
+    let mut new_key = Vec::with_capacity(key.len());
+
+    fn unhex(num: u8) -> u8 {
+        match num {
+            b'0'..=b'9' => num - b'0',
+            b'a'..=b'f' => num - b'a' + 10,
+            unreachable => unreachable!("Unreachable branch num = {}", unreachable),
+        }
+    }
+
+    let bytes = key.as_bytes();
+
+    let mut index = 0;
+    while index < bytes.len() {
+        match bytes[index] {
+            b'%' => {
+                // The last byte of the string won't be '%'
+                let mut num = unhex(bytes[index + 1]) * 16;
+                num += unhex(bytes[index + 2]);
+                new_key.push(num);
+                index += 3;
+            }
+            other => {
+                new_key.push(other);
+                index += 1;
+            }
+        }
+    }
+
+    String::from_utf8(new_key)
+}

common/base/tests/it/main.rs

@@ -15,3 +15,4 @@
 mod progress;
 mod runtime;
 mod stoppable;
+mod string_func;

common/base/tests/it/string_func.rs

@@ -0,0 +1,28 @@
+// Copyright 2022 Datafuse Labs.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+use common_base::*;
+use common_exception::Result;
+
+#[test]
+fn test_progress() -> Result<()> {
+    let original_key = "databend/test_user123!!";
+    let new_key = escape_for_key(original_key);
+    assert_eq!(Ok("databend%2ftest_user123%21%21".to_string()), new_key);
+    assert_eq!(
+        Ok(original_key.to_string()),
+        unescape_for_key(new_key.unwrap().as_str())
+    );
+    Ok(())
+}

common/management/src/cluster/cluster_mgr.rs

@@ -17,6 +17,8 @@ use std::sync::Arc;
 use std::time::Duration;
 use std::time::UNIX_EPOCH;
 
+use common_base::escape_for_key;
+use common_base::unescape_for_key;
 use common_exception::ErrorCode;
 use common_exception::Result;
 use common_meta_api::KVApi;
@@ -58,70 +60,12 @@ impl ClusterMgr {
             cluster_prefix: format!(
                 "{}/{}/{}/databend_query",
                 CLUSTER_API_KEY_PREFIX,
-                Self::escape_for_key(tenant)?,
-                Self::escape_for_key(cluster_id)?
+                escape_for_key(tenant)?,
+                escape_for_key(cluster_id)?
             ),
         })
     }
 
-    fn escape_for_key(key: &str) -> Result<String> {
-        let mut new_key = Vec::with_capacity(key.len());
-
-        fn hex(num: u8) -> u8 {
-            match num {
-                0..=9 => b'0' + num,
-                10..=15 => b'a' + (num - 10),
-                unreachable => unreachable!("Unreachable branch num = {}", unreachable),
-            }
-        }
-
-        for char in key.as_bytes() {
-            match char {
-                b'0'..=b'9' => new_key.push(*char),
-                b'_' | b'a'..=b'z' | b'A'..=b'Z' => new_key.push(*char),
-                _other => {
-                    new_key.push(b'%');
-                    new_key.push(hex(*char / 16));
-                    new_key.push(hex(*char % 16));
-                }
-            }
-        }
-
-        Ok(String::from_utf8(new_key)?)
-    }
-
-    fn unescape_for_key(key: &str) -> Result<String> {
-        let mut new_key = Vec::with_capacity(key.len());
-
-        fn unhex(num: u8) -> u8 {
-            match num {
-                b'0'..=b'9' => num - b'0',
-                b'a'..=b'f' => num - b'a',
-                unreachable => unreachable!("Unreachable branch num = {}", unreachable),
-            }
-        }
-
-        let bytes = key.as_bytes();
-
-        let mut index = 0;
-        while index < bytes.len() {
-            match bytes[index] {
-                b'%' => {
-                    let mut num = unhex(bytes[index + 1]) * 16;
-                    num += unhex(bytes[index + 2]);
-                    new_key.push(num);
-                    index += 3;
-                }
-                other => {
-                    new_key.push(other);
-                    index += 1;
-                }
-            }
-        }
-
-        Ok(String::from_utf8(new_key)?)
-    }
-
     fn new_lift_time(&self) -> KVMeta {
         let now = std::time::SystemTime::now();
         let expire_at = now
@@ -142,11 +86,7 @@ impl ClusterApi for ClusterMgr {
         let seq = MatchSeq::Exact(0);
         let meta = Some(self.new_lift_time());
         let value = Operation::Update(serde_json::to_vec(&node)?);
-        let node_key = format!(
-            "{}/{}",
-            self.cluster_prefix,
-            Self::escape_for_key(&node.id)?
-        );
+        let node_key = format!("{}/{}", self.cluster_prefix, escape_for_key(&node.id)?);
         let upsert_node = self
             .kv_api
             .upsert_kv(UpsertKVAction::new(&node_key, seq, value, meta));
@@ -169,7 +109,7 @@ impl ClusterApi for ClusterMgr {
         for (node_key, value) in values {
             let mut node_info = serde_json::from_slice::<NodeInfo>(&value.data)?;
 
-            let node_key = Self::unescape_for_key(&node_key)?;
+            let node_key = unescape_for_key(&node_key)?;
             node_info.id = node_key[self.cluster_prefix.len() + 1..].to_string();
             nodes_info.push(node_info);
         }
@@ -178,11 +118,7 @@ impl ClusterApi for ClusterMgr {
     }
 
     async fn drop_node(&self, node_id: String, seq: Option<u64>) -> Result<()> {
-        let node_key = format!(
-            "{}/{}",
-            self.cluster_prefix,
-            Self::escape_for_key(&node_id)?
-        );
+        let node_key = format!("{}/{}", self.cluster_prefix, escape_for_key(&node_id)?);
         let upsert_node = self.kv_api.upsert_kv(UpsertKVAction::new(
             &node_key,
             seq.into(),
@@ -205,11 +141,7 @@ impl ClusterApi for ClusterMgr {
 
     async fn heartbeat(&self, node: &NodeInfo, seq: Option<u64>) -> Result<u64> {
         let meta = Some(self.new_lift_time());
-        let node_key = format!(
-            "{}/{}",
-            self.cluster_prefix,
-            Self::escape_for_key(&node.id)?
-        );
+        let node_key = format!("{}/{}", self.cluster_prefix, escape_for_key(&node.id)?);
         let seq = match seq {
             None => MatchSeq::GE(1),
             Some(exact) => MatchSeq::Exact(exact),

common/management/src/stage/stage_mgr.rs

@@ -14,6 +14,7 @@
 
 use std::sync::Arc;
 
+use common_base::escape_for_key;
 use common_exception::ErrorCode;
 use common_exception::Result;
 use common_meta_api::KVApi;
@@ -45,7 +46,7 @@ impl StageMgr {
 
         Ok(StageMgr {
             kv_api,
-            stage_prefix: format!("{}/{}", USER_STAGE_API_KEY_PREFIX, tenant),
+            stage_prefix: format!("{}/{}", USER_STAGE_API_KEY_PREFIX, escape_for_key(tenant)?),
         })
     }
 }
@@ -55,7 +56,11 @@ impl StageApi for StageMgr {
     async fn add_stage(&self, info: UserStageInfo) -> Result<u64> {
         let seq = MatchSeq::Exact(0);
         let val = Operation::Update(serde_json::to_vec(&info)?);
-        let key = format!("{}/{}", self.stage_prefix, info.stage_name);
+        let key = format!(
+            "{}/{}",
+            self.stage_prefix,
+            escape_for_key(&info.stage_name)?
+        );
         let upsert_info = self
             .kv_api
             .upsert_kv(UpsertKVAction::new(&key, seq, val, None));
@@ -72,7 +77,7 @@ impl StageApi for StageMgr {
     }
 
     async fn get_stage(&self, name: &str, seq: Option<u64>) -> Result<SeqV<UserStageInfo>> {
-        let key = format!("{}/{}", self.stage_prefix, name);
+        let key = format!("{}/{}", self.stage_prefix, escape_for_key(name)?);
         let kv_api = self.kv_api.clone();
         let get_kv = async move { kv_api.get_kv(&key).await };
         let res = get_kv.await?;
@@ -97,7 +102,7 @@ impl StageApi for StageMgr {
     }
 
     async fn drop_stage(&self, name: &str, seq: Option<u64>) -> Result<()> {
-        let key = format!("{}/{}", self.stage_prefix, name);
+        let key = format!("{}/{}", self.stage_prefix, escape_for_key(name)?);
         let kv_api = self.kv_api.clone();
         let upsert_kv = async move {
             kv_api

common/management/src/udf/udf_mgr.rs

@@ -15,6 +15,7 @@
 use std::sync::Arc;
 
 use common_ast::udfs::UDFParser;
+use common_base::escape_for_key;
 use common_exception::ErrorCode;
 use common_exception::Result;
 use common_functions::is_builtin_function;
@@ -47,7 +48,7 @@ impl UdfMgr {
 
         Ok(UdfMgr {
             kv_api,
-            udf_prefix: format!("{}/{}", UDF_API_KEY_PREFIX, tenant),
+            udf_prefix: format!("{}/{}", UDF_API_KEY_PREFIX, escape_for_key(tenant)?),
         })
     }
 }
@@ -69,7 +70,7 @@ impl UdfApi for UdfMgr {
 
         let seq = MatchSeq::Exact(0);
         let val = Operation::Update(serde_json::to_vec(&info)?);
-        let key = format!("{}/{}", self.udf_prefix, info.name);
+        let key = format!("{}/{}", self.udf_prefix, escape_for_key(&info.name)?);
         let upsert_info = self
             .kv_api
             .upsert_kv(UpsertKVAction::new(&key, seq, val, None));
@@ -97,7 +98,7 @@ impl UdfApi for UdfMgr {
         let _ = self.get_udf(info.name.as_str(), seq).await?;
 
         let val = Operation::Update(serde_json::to_vec(&info)?);
-        let key = format!("{}/{}", self.udf_prefix, info.name);
+        let key = format!("{}/{}", self.udf_prefix, escape_for_key(&info.name)?);
         let upsert_info =
             self.kv_api
                 .upsert_kv(UpsertKVAction::new(&key, MatchSeq::from(seq), val, None));
@@ -113,7 +114,7 @@ impl UdfApi for UdfMgr {
     }
 
     async fn get_udf(&self, udf_name: &str, seq: Option<u64>) -> Result<SeqV<UserDefinedFunction>> {
-        let key = format!("{}/{}", self.udf_prefix, udf_name);
+        let key = format!("{}/{}", self.udf_prefix, escape_for_key(udf_name)?);
         let kv_api = self.kv_api.clone();
         let get_kv = async move { kv_api.get_kv(&key).await };
         let res = get_kv.await?;
@@ -138,7 +139,7 @@ impl UdfApi for UdfMgr {
     }
 
     async fn drop_udf(&self, udf_name: &str, seq: Option<u64>) -> Result<()> {
-        let key = format!("{}/{}", self.udf_prefix, udf_name);
+        let key = format!("{}/{}", self.udf_prefix, escape_for_key(udf_name)?);
         let kv_api = self.kv_api.clone();
         let upsert_kv = async move {
             kv_api