tests_fuzz/translator/
csv.rs

1// Copyright 2023 Greptime Team
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use crate::error::Error;
16use crate::ir::insert_expr::{InsertIntoExpr, RowValue};
17use crate::translator::DslTranslator;
18
19/// One CSV record converted from an insert row.
20#[derive(Debug, Clone, PartialEq, Eq)]
21pub struct CsvRecord {
22    /// Cell values in column order.
23    pub values: Vec<String>,
24}
25
26/// CSV records converted from an insert expression.
27#[derive(Debug, Clone, PartialEq, Eq)]
28pub struct CsvRecords {
29    /// Target table name from insert expression.
30    pub table_name: String,
31    /// Header values from insert columns.
32    pub headers: Vec<String>,
33    /// Converted row records.
34    pub records: Vec<CsvRecord>,
35}
36
37/// Translates `InsertIntoExpr` into CSV-writer-ready records.
38pub struct InsertExprToCsvRecordsTranslator;
39
40impl DslTranslator<InsertIntoExpr, CsvRecords> for InsertExprToCsvRecordsTranslator {
41    type Error = Error;
42
43    fn translate(&self, input: &InsertIntoExpr) -> Result<CsvRecords, Self::Error> {
44        let headers = input
45            .columns
46            .iter()
47            .map(|column| column.name.to_string())
48            .collect::<Vec<_>>();
49        let records = input
50            .values_list
51            .iter()
52            .map(|row| CsvRecord {
53                values: row.iter().map(Self::format_row_value).collect(),
54            })
55            .collect::<Vec<_>>();
56
57        Ok(CsvRecords {
58            table_name: input.table_name.to_string(),
59            headers,
60            records,
61        })
62    }
63}
64
65impl InsertExprToCsvRecordsTranslator {
66    fn format_row_value(value: &RowValue) -> String {
67        match value {
68            RowValue::Value(datatypes::value::Value::Null) => String::new(),
69            RowValue::Value(v) => v.to_string(),
70            RowValue::Default => "DEFAULT".to_string(),
71        }
72    }
73}
74
75#[cfg(test)]
76mod tests {
77    use datatypes::data_type::ConcreteDataType;
78
79    use super::InsertExprToCsvRecordsTranslator;
80    use crate::ir::create_expr::ColumnOption;
81    use crate::ir::insert_expr::{InsertIntoExpr, RowValue};
82    use crate::ir::{Column, Ident};
83    use crate::translator::DslTranslator;
84
85    #[test]
86    fn test_translate_insert_expr_to_csv_records() {
87        let input = InsertIntoExpr {
88            table_name: Ident::new("metric_a"),
89            omit_column_list: false,
90            columns: vec![
91                Column {
92                    name: "host".into(),
93                    column_type: ConcreteDataType::string_datatype(),
94                    options: vec![ColumnOption::PrimaryKey],
95                },
96                Column {
97                    name: "value".into(),
98                    column_type: ConcreteDataType::float64_datatype(),
99                    options: vec![],
100                },
101            ],
102            values_list: vec![
103                vec![
104                    RowValue::Value(datatypes::value::Value::String("web-1".into())),
105                    RowValue::Value(datatypes::value::Value::Int32(15)),
106                ],
107                vec![
108                    RowValue::Value(datatypes::value::Value::Null),
109                    RowValue::Default,
110                ],
111            ],
112        };
113
114        let output = InsertExprToCsvRecordsTranslator.translate(&input).unwrap();
115        assert_eq!(output.table_name, "metric_a");
116        assert_eq!(output.headers, vec!["host", "value"]);
117        assert_eq!(output.records.len(), 2);
118        assert_eq!(output.records[0].values, vec!["web-1", "15"]);
119        assert_eq!(output.records[1].values, vec!["", "DEFAULT"]);
120    }
121}