sqlness_runner/
formatter.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 std::borrow::Cow;
16use std::fmt::Display;
17use std::sync::Arc;
18
19use client::{Output, OutputData, RecordBatches};
20use common_error::ext::ErrorExt;
21use datatypes::prelude::ConcreteDataType;
22use datatypes::scalars::ScalarVectorBuilder;
23use datatypes::schema::{ColumnSchema, Schema};
24use datatypes::vectors::{StringVectorBuilder, VectorRef};
25use mysql::Row as MySqlRow;
26use tokio_postgres::SimpleQueryMessage as PgRow;
27
28use crate::client::MysqlSqlResult;
29
30/// A formatter for errors.
31pub struct ErrorFormatter<E: ErrorExt>(E);
32
33impl<E: ErrorExt> From<E> for ErrorFormatter<E> {
34    fn from(error: E) -> Self {
35        ErrorFormatter(error)
36    }
37}
38
39impl<E: ErrorExt> Display for ErrorFormatter<E> {
40    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
41        let status_code = self.0.status_code();
42        let root_cause = self.0.output_msg();
43        let root_cause = normalize_meta_client_error(&root_cause, status_code);
44        write!(
45            f,
46            "Error: {}({status_code}), {root_cause}",
47            status_code as u32
48        )
49    }
50}
51
52fn normalize_meta_client_error(
53    root_cause: &str,
54    status_code: common_error::status_code::StatusCode,
55) -> &str {
56    let details_prefix = format!(", code: {status_code}, tonic code: ");
57    match root_cause.rsplit_once(&details_prefix) {
58        Some((message, tonic_code)) if !tonic_code.is_empty() => message,
59        _ => root_cause,
60    }
61}
62
63/// A formatter for [`Output`].
64pub struct OutputFormatter(Output);
65
66impl From<Output> for OutputFormatter {
67    fn from(output: Output) -> Self {
68        OutputFormatter(output)
69    }
70}
71
72impl Display for OutputFormatter {
73    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
74        match &self.0.data {
75            OutputData::AffectedRows(rows) => {
76                write!(f, "Affected Rows: {rows}")
77            }
78            OutputData::RecordBatches(recordbatches) => {
79                let pretty = recordbatches.pretty_print().map_err(|e| e.to_string());
80                match pretty {
81                    Ok(s) => write!(f, "{s}"),
82                    Err(e) => {
83                        write!(f, "Failed to pretty format {recordbatches:?}, error: {e}")
84                    }
85                }
86            }
87            OutputData::Stream(_) => unreachable!(),
88        }
89    }
90}
91
92pub struct PostgresqlFormatter(Vec<PgRow>);
93
94impl From<Vec<PgRow>> for PostgresqlFormatter {
95    fn from(rows: Vec<PgRow>) -> Self {
96        PostgresqlFormatter(rows)
97    }
98}
99
100impl Display for PostgresqlFormatter {
101    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
102        if self.0.is_empty() {
103            return f.write_fmt(format_args!("(Empty response)"));
104        }
105
106        if let PgRow::CommandComplete(affected_rows) = &self.0[0] {
107            return write!(
108                f,
109                "{}",
110                OutputFormatter(Output::new_with_affected_rows(*affected_rows as usize))
111            );
112        };
113
114        let Some(recordbatches) = build_recordbatches_from_postgres_rows(&self.0) else {
115            return Ok(());
116        };
117        write!(
118            f,
119            "{}",
120            OutputFormatter(Output::new_with_record_batches(recordbatches))
121        )
122    }
123}
124
125fn build_recordbatches_from_postgres_rows(rows: &[PgRow]) -> Option<RecordBatches> {
126    // create schema
127    let schema = match &rows[0] {
128        PgRow::RowDescription(desc) => Arc::new(Schema::new(
129            desc.iter()
130                .map(|column| {
131                    ColumnSchema::new(column.name(), ConcreteDataType::string_datatype(), true)
132                })
133                .collect(),
134        )),
135        _ => unreachable!(),
136    };
137    if schema.num_columns() == 0 {
138        return None;
139    }
140
141    // convert to string vectors
142    let mut columns: Vec<StringVectorBuilder> = (0..schema.num_columns())
143        .map(|_| StringVectorBuilder::with_capacity(schema.num_columns()))
144        .collect();
145    for row in rows.iter().skip(1) {
146        if let PgRow::Row(row) = row {
147            for (i, column) in columns.iter_mut().enumerate().take(schema.num_columns()) {
148                column.push(row.get(i));
149            }
150        }
151    }
152    let columns: Vec<VectorRef> = columns
153        .into_iter()
154        .map(|mut col| Arc::new(col.finish()) as VectorRef)
155        .collect();
156
157    // construct recordbatch
158    let recordbatches = RecordBatches::try_from_columns(schema, columns)
159        .expect("Failed to construct recordbatches from columns. Please check the schema.");
160    Some(recordbatches)
161}
162
163/// A formatter for [`MysqlSqlResult`].
164pub struct MysqlFormatter(MysqlSqlResult);
165
166impl From<MysqlSqlResult> for MysqlFormatter {
167    fn from(result: MysqlSqlResult) -> Self {
168        MysqlFormatter(result)
169    }
170}
171
172impl Display for MysqlFormatter {
173    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
174        match &self.0 {
175            MysqlSqlResult::AffectedRows(rows) => {
176                write!(f, "affected_rows: {rows}")
177            }
178            MysqlSqlResult::Rows(rows) => {
179                if rows.is_empty() {
180                    return f.write_fmt(format_args!("(Empty response)"));
181                }
182
183                let recordbatches = build_recordbatches_from_mysql_rows(rows);
184                write!(
185                    f,
186                    "{}",
187                    OutputFormatter(Output::new_with_record_batches(recordbatches))
188                )
189            }
190        }
191    }
192}
193
194pub fn build_recordbatches_from_mysql_rows(rows: &[MySqlRow]) -> RecordBatches {
195    // create schema
196    let head_column = &rows[0];
197    let head_binding = head_column.columns();
198    let names = head_binding
199        .iter()
200        .map(|column| column.name_str())
201        .collect::<Vec<Cow<str>>>();
202    let schema = Arc::new(Schema::new(
203        names
204            .iter()
205            .map(|name| {
206                ColumnSchema::new(name.to_string(), ConcreteDataType::string_datatype(), true)
207            })
208            .collect(),
209    ));
210
211    // convert to string vectors
212    let mut columns: Vec<StringVectorBuilder> = (0..schema.num_columns())
213        .map(|_| StringVectorBuilder::with_capacity(schema.num_columns()))
214        .collect();
215    for row in rows.iter() {
216        for (i, name) in names.iter().enumerate() {
217            // `get::<String, _>` panics on NULL values, so fetch as Option.
218            let value: Option<String> = row.get::<Option<String>, &str>(name).flatten();
219            columns[i].push(value.as_deref());
220        }
221    }
222    let columns: Vec<VectorRef> = columns
223        .into_iter()
224        .map(|mut col| Arc::new(col.finish()) as VectorRef)
225        .collect();
226
227    // construct recordbatch
228    RecordBatches::try_from_columns(schema, columns)
229        .expect("Failed to construct recordbatches from columns. Please check the schema.")
230}
231
232#[cfg(test)]
233mod tests {
234    use common_error::ext::PlainError;
235    use common_error::status_code::StatusCode;
236
237    use super::*;
238
239    #[test]
240    fn test_normalize_meta_client_error() {
241        let error = PlainError::new(
242            "Invalid options, code: InvalidArguments, tonic code: Client specified an invalid argument"
243                .to_string(),
244            StatusCode::InvalidArguments,
245        );
246
247        assert_eq!(
248            "Error: 1004(InvalidArguments), Invalid options",
249            ErrorFormatter::from(error).to_string()
250        );
251    }
252
253    #[test]
254    fn test_preserve_regular_error() {
255        let error = PlainError::new(
256            "Invalid options without transport details".to_string(),
257            StatusCode::InvalidArguments,
258        );
259
260        assert_eq!(
261            "Error: 1004(InvalidArguments), Invalid options without transport details",
262            ErrorFormatter::from(error).to_string()
263        );
264    }
265
266    #[test]
267    fn test_preserve_error_with_mismatched_code() {
268        let error = PlainError::new(
269            "Invalid options, code: Unsupported, tonic code: Operation is not supported"
270                .to_string(),
271            StatusCode::InvalidArguments,
272        );
273
274        assert_eq!(
275            "Error: 1004(InvalidArguments), Invalid options, code: Unsupported, tonic code: Operation is not supported",
276            ErrorFormatter::from(error).to_string()
277        );
278    }
279}
sqlness_runner/formatter.rs

sqlness_runner/
formatter.rs
