Merge pull request #9175 from BohuTANG/dev-unescape-string

feat: change the unescape to enquota crate and support \\x01 -> \x01

Author: mergify[bot]

Cargo.lock

@@ -21,6 +21,7 @@ bincode = { version = "2.0.0-rc.1", features = ["serde", "std"] }
 bytes = "1.2.1"
 chrono = { workspace = true }
 chrono-tz = { workspace = true }
+enquote = "1.1.0"
 lexical-core = "0.8.5"
 micromarshal = "0.2.1"
 ordered-float = "3.1.0"

src/common/io/Cargo.toml

@@ -14,6 +14,7 @@
 
 use std::cmp;
 
+use common_exception::ErrorCode;
 use common_exception::Result;
 
 pub fn convert_byte_size(num: f64) -> String {
@@ -80,63 +81,11 @@ pub fn deserialize_from_slice<T: serde::de::DeserializeOwned>(slice: &mut &[u8])
     Ok(value)
 }
 
-pub fn is_control_ascii(c: u8) -> bool {
-    c <= 31
-}
-
-pub fn parse_escape_string(bs: &[u8]) -> String {
-    let bs = parse_escape_bytes(bs);
-
-    let mut cs = Vec::with_capacity(bs.len());
-    for b in bs {
-        cs.push(b as char);
-    }
-    cs.iter().collect()
-}
-
-pub fn parse_escape_bytes(bs: &[u8]) -> Vec<u8> {
-    let mut vs = Vec::with_capacity(bs.len());
-    let mut i = 0;
-    while i < bs.len() {
-        if bs[i] == b'\\' {
-            if i + 1 < bs.len() {
-                let c = parse_escape_byte(bs[i + 1]);
-                if c != b'\\'
-                    && c != b'\''
-                    && c != b'"'
-                    && c != b'`'
-                    && c != b'/'
-                    && !is_control_ascii(c)
-                {
-                    vs.push(b'\\');
-                }
-
-                vs.push(c);
-                i += 2;
-            } else {
-                // end with \
-                vs.push(b'\\');
-                break;
-            }
-        } else {
-            vs.push(bs[i]);
-            i += 1;
-        }
-    }
-
-    vs
-}
-
-// https://doc.rust-lang.org/reference/tokens.html
-pub fn parse_escape_byte(b: u8) -> u8 {
-    match b {
-        b'e' => b'\x1B',
-        b'n' => b'\n',
-        b'r' => b'\r',
-        b't' => b'\t',
-        b'0' => b'\0',
-        _ => b,
-    }
+/// Returns string after processing escapes.
+/// This used for settings string unescape, like unescape format_field_delimiter from `\\x01` to `\x01`.
+pub fn unescape_string(escape_str: &str) -> Result<String> {
+    enquote::unescape(escape_str, None)
+        .map_err(|e| ErrorCode::Internal(format!("unescape:{} error:{:?}", escape_str, e)))
 }
 
 /// Mask a string by "******", but keep `unmask_len` of suffix.

src/common/io/src/utils.rs

@@ -33,15 +33,22 @@ fn convert_test() {
 }
 
 #[test]
-fn parse_escape() {
+fn test_unescape_string() {
     let cases = vec![
         vec!["a", "a"],
         vec!["abc", "abc"],
+        vec!["\\x01", "\x01"],
+        vec!["\x01", "\x01"],
         vec!["\t\nabc", "\t\nabc"],
+        vec!["\"\t\nabc\"", "\"\t\nabc\""],
+        vec!["\"\\t\nabc\"", "\"\t\nabc\""],
+        vec!["'\\t\nabc'", "'\t\nabc'"],
         vec!["\\t\\nabc", "\t\nabc"],
+        vec!["\\\\", r"\"],
+        vec!["\\\\", "\\"],
     ];
 
     for c in cases {
-        assert_eq!(parse_escape_bytes(c[0].as_bytes()), c[1].as_bytes());
+        assert_eq!(unescape_string(c[0]).unwrap(), c[1]);
     }
 }

src/common/io/tests/it/utils.rs

@@ -22,7 +22,7 @@ use common_base::base::ProgressValues;
 use common_base::base::TrySpawn;
 use common_exception::ErrorCode;
 use common_exception::Result;
-use common_io::prelude::parse_escape_string;
+use common_io::prelude::unescape_string;
 use common_pipeline_sources::processors::sources::input_formats::InputContext;
 use common_pipeline_sources::processors::sources::input_formats::StreamingReadBatch;
 use futures::StreamExt;
@@ -103,9 +103,11 @@ pub async fn streaming_load(
     for (key, value) in req.headers().iter() {
         if settings.has_setting(key.as_str()) {
             let value = value.to_str().map_err(InternalServerError)?;
-            let value = parse_escape_string(remove_quote(value.as_bytes()));
+            let unquote =
+                std::str::from_utf8(remove_quote(value.as_bytes())).map_err(InternalServerError)?;
+            let value = unescape_string(unquote).map_err(InternalServerError)?;
             settings
-                .set_settings(key.to_string(), value, false)
+                .set_settings(key.to_string(), value.to_string(), false)
                 .map_err(InternalServerError)?
         }
     }

src/query/service/src/servers/http/v1/load.rs

@@ -32,7 +32,7 @@ use common_catalog::table_context::TableContext;
 use common_config::GlobalConfig;
 use common_exception::ErrorCode;
 use common_exception::Result;
-use common_io::prelude::parse_escape_string;
+use common_io::prelude::unescape_string;
 use common_meta_types::FileFormatOptions;
 use common_meta_types::StageFileFormatType;
 use common_meta_types::UserStageInfo;
@@ -611,54 +611,44 @@ pub fn parse_copy_file_format_options(
         .parse::<u64>()?;
 
     // Field delimiter.
-    let field_delimiter = parse_escape_string(
+    let field_delimiter = unescape_string(
         file_format_options
             .get("field_delimiter")
-            .unwrap_or(&"".to_string())
-            .as_bytes(),
-    );
+            .unwrap_or(&"".to_string()),
+    )?;
 
     // Record delimiter.
-    let record_delimiter = parse_escape_string(
+    let record_delimiter = unescape_string(
         file_format_options
             .get("record_delimiter")
-            .unwrap_or(&"".to_string())
-            .as_bytes(),
-    );
+            .unwrap_or(&"".to_string()),
+    )?;
 
     // NaN display.
-    let nan_display = parse_escape_string(
+    let nan_display = unescape_string(
         file_format_options
             .get("nan_display")
-            .unwrap_or(&"".to_string())
-            .as_bytes(),
-    );
+            .unwrap_or(&"".to_string()),
+    )?;
 
     // Escape
-    let escape = parse_escape_string(
-        file_format_options
-            .get("escape")
-            .unwrap_or(&"".to_string())
-            .as_bytes(),
-    );
+    let escape = unescape_string(file_format_options.get("escape").unwrap_or(&"".to_string()))?;
 
     // Compression delimiter.
-    let compression = parse_escape_string(
+    let compression = unescape_string(
         file_format_options
             .get("compression")
-            .unwrap_or(&"none".to_string())
-            .as_bytes(),
-    )
+            .unwrap_or(&"none".to_string()),
+    )?
     .parse()
     .map_err(ErrorCode::UnknownCompressionType)?;
 
     // Row tag in xml.
-    let row_tag = parse_escape_string(
+    let row_tag = unescape_string(
         file_format_options
             .get("row_tag")
-            .unwrap_or(&"".to_string())
-            .as_bytes(),
-    );
+            .unwrap_or(&"".to_string()),
+    )?;
 
     Ok(FileFormatOptions {
         format: file_format,

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

@@ -3,6 +3,10 @@
 CURDIR=$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)
 . "$CURDIR"/../../../../shell_env.sh
 
+
+# Should be <root>/tests/data/
+DATADIR=$(realpath $CURDIR/../../../data/)
+
 echo "drop table if exists test_csv" | $MYSQL_CLIENT_CONNECT
 echo "drop table if exists test_csv2" | $MYSQL_CLIENT_CONNECT
 
@@ -41,16 +45,16 @@ curl -sH "insert_sql:insert into test_csv format CSV" -F "upload=@/tmp/escape_no
 echo "select * from test_csv" | $MYSQL_CLIENT_CONNECT
 echo "truncate table test_csv" | $MYSQL_CLIENT_CONNECT
 
-curl -sH "insert_sql:insert into test_csv format CSV" -H "format_escape:'\\'" -F "upload=@/tmp/escape_slash.csv" -u root: -XPUT "http://localhost:${QUERY_HTTP_HANDLER_PORT}/v1/streaming_load" | grep -c "SUCCESS"
+curl -sH "insert_sql:insert into test_csv format CSV" -H "format_escape:'\\\'" -F "upload=@/tmp/escape_slash.csv" -u root: -XPUT "http://localhost:${QUERY_HTTP_HANDLER_PORT}/v1/streaming_load" | grep -c "SUCCESS"
 echo "select * from test_csv" | $MYSQL_CLIENT_CONNECT
 echo "truncate table test_csv" | $MYSQL_CLIENT_CONNECT
 
 aws --endpoint-url http://127.0.0.1:9900/ s3 cp /tmp/escape_slash2.csv s3://testbucket/admin/data/csv/escape_slash2.csv > /dev/null 2>&1
-echo "copy into test_csv from 's3://testbucket/admin/data/csv/escape_slash2.csv' connection=(aws_key_id='minioadmin' aws_secret_key='minioadmin' endpoint_url='http://127.0.0.1:9900/') FILE_FORMAT = (type = 'CSV' escape='\\\')" | $MYSQL_CLIENT_CONNECT
+echo "copy into test_csv from 'fs:///tmp/escape_slash2.csv' FILE_FORMAT = (type = 'CSV' escape='\\\\\\\\')" | $MYSQL_CLIENT_CONNECT
 echo "select * from test_csv" | $MYSQL_CLIENT_CONNECT
 
 aws --endpoint-url http://127.0.0.1:9900/ s3 cp /tmp/escape_slash3.csv s3://testbucket/admin/data/csv/escape_slash3.csv > /dev/null 2>&1
-echo "copy into test_csv2 from 's3://testbucket/admin/data/csv/escape_slash3.csv' connection=(aws_key_id='minioadmin' aws_secret_key='minioadmin' endpoint_url='http://127.0.0.1:9900/') FILE_FORMAT = (type = 'CSV' escape='\\\' skip_header=1)" | $MYSQL_CLIENT_CONNECT
+echo "copy into test_csv2 from 'fs:///tmp/escape_slash3.csv' FILE_FORMAT = (type = 'CSV' escape='\\\\\\\\' skip_header=1)" | $MYSQL_CLIENT_CONNECT
 echo "select * from test_csv2" | $MYSQL_CLIENT_CONNECT
 
 echo "drop table if exists test_csv" | $MYSQL_CLIENT_CONNECT

tests/suites/1_stateful/05_formats/05_02_csv/05_02_01_csv_escape.sh

@@ -0,0 +1,2 @@
+1
+4

tests/suites/1_stateful/05_formats/05_02_csv/05_02_05_csv_with_0x01_delimiter.result

@@ -0,0 +1,17 @@
+#!/usr/bin/env bash
+
+CURDIR=$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)
+. "$CURDIR"/../../../../shell_env.sh
+
+ echo "drop table if exists test_x01_csv" | $MYSQL_CLIENT_CONNECT
+
+ echo "CREATE TABLE test_x01_csv
+ (
+     a VARCHAR,
+     b Int,
+     c VARCHAR
+ );" | $MYSQL_CLIENT_CONNECT
+
+ curl -sH "insert_sql:insert into test_x01_csv format CSV" -H "format_field_delimiter:\x01" -H "format_skip_header:1" -H "format_record_delimiter:\r\n" -F "upload=@${CURDIR}/x01_field_delimiter.csv" \
+ -u root: -XPUT "http://localhost:${QUERY_HTTP_HANDLER_PORT}/v1/streaming_load" | grep -c "SUCCESS"
+ echo "select count() from test_x01_csv" | $MYSQL_CLIENT_CONNECT

tests/suites/1_stateful/05_formats/05_02_csv/05_02_05_csv_with_0x01_delimiter.sh

@@ -0,0 +1,5 @@
+Name
Age
Occupation
+Donald John Trump
80
president
+马斯克
43
特斯拉
+测试
39
TT
+测试
39
test