Merge pull request #9188 from BohuTANG/dev-record-delimiter

feat: custom csv record delimiter

Author: BohuTANG

src/query/formats/src/format_option_checker.rs

@@ -230,17 +230,25 @@ pub fn check_field_delimiter(option: &mut String, default: &str) -> Result<()> {
     Ok(())
 }
 
+/// `\r\n` or u8
 pub fn check_record_delimiter(option: &mut String) -> Result<()> {
-    if option.is_empty() {
-        *option = "\n".to_string()
-    } else {
-        let o = option.as_str();
-        if o != "\n" && o != "\r\n" {
+    match option.len() {
+        0 => *option = "\n".to_string(),
+        1 => {}
+        2 => {
+            if option != "\r\n" {
+                return Err(ErrorCode::InvalidArgument(
+                    "record_delimiter with two chars can only be '\\r\\n'",
+                ));
+            };
+        }
+        _ => {
             return Err(ErrorCode::InvalidArgument(
-                "record_delimiter can only be '\\n' or '\\r\\n'",
+                "record_delimiter can not more than two chars, please use one char or '\\r\\n'",
             ));
-        };
+        }
     }
+
     Ok(())
 }
 

src/query/formats/src/lib.rs

@@ -28,5 +28,6 @@ pub use field_decoder::*;
 pub use file_format_type::parse_timezone;
 pub use file_format_type::FileFormatOptionsExt;
 pub use file_format_type::FileFormatTypeExt;
+pub use format_option_checker::check_record_delimiter;
 
 use crate::common_settings::CommonSettings;

src/query/formats/tests/it/format_option_checker.rs

@@ -0,0 +1,37 @@
+// 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_formats::check_record_delimiter;
+
+/// This test code is written by OpenAI's GPT-3.
+#[test]
+fn test_check_record_delimiter() {
+    let mut option = "".to_string();
+    assert!(check_record_delimiter(&mut option).is_ok());
+    assert_eq!(option, "\n");
+
+    let mut option = "|".to_string();
+    assert!(check_record_delimiter(&mut option).is_ok());
+    assert_eq!(option, "|");
+
+    let mut option = "\r\n".to_string();
+    assert!(check_record_delimiter(&mut option).is_ok());
+    assert_eq!(option, "\r\n");
+
+    let mut option = "foo".to_string();
+    assert!(check_record_delimiter(&mut option).is_err());
+
+    let mut option = "|\r".to_string();
+    assert!(check_record_delimiter(&mut option).is_err());
+}

src/query/formats/tests/it/main.rs

@@ -21,6 +21,7 @@ use common_meta_types::StageFileFormatType;
 use common_settings::Settings;
 
 mod field_encoder;
+mod format_option_checker;
 mod output_format_json_each_row;
 mod output_format_tcsv;
 mod output_format_utils;

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

@@ -11,6 +11,6 @@ CURDIR=$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)
      b Int NULL
  );" | $MYSQL_CLIENT_CONNECT
 
- curl -sH "insert_sql:insert into test_csv format CSV" -F "upload=@${CURDIR}/no_newline.csv" \
+ curl -sH "insert_sql:insert into test_csv format CSV" -F "upload=@${CURDIR}/testdata/no_newline.csv" \
  -u root: -XPUT "http://localhost:${QUERY_HTTP_HANDLER_PORT}/v1/streaming_load" | grep -c "SUCCESS"
  echo "select * from test_csv" | $MYSQL_CLIENT_CONNECT

tests/suites/1_stateful/05_formats/05_02_csv/05_02_05_csv_with_0x01_delimiter.result renamed to tests/suites/1_stateful/05_formats/05_02_csv/05_02_05_csv_0x01_field_0x02_record.result

@@ -12,6 +12,6 @@ CURDIR=$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)
      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" \
+ curl -sH "insert_sql:insert into test_x01_csv format CSV" -H "format_field_delimiter:\x01" -H "format_record_delimiter:\x02" -H "format_skip_header:0" -F "upload=@${CURDIR}/testdata/x01_field_x02_record.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 renamed to tests/suites/1_stateful/05_formats/05_02_csv/05_02_05_csv_0x01_field_0x02_record.sh

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