servers/mysql/
writer.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::fmt::Write as FmtWrite;
16use std::io;
17use std::time::Duration;
18
19use arrow::array::{Array, AsArray};
20use arrow::datatypes::{
21    Date32Type, Decimal128Type, Float32Type, Float64Type, Int8Type, Int16Type, Int32Type,
22    Int64Type, IntervalDayTimeType, IntervalMonthDayNanoType, IntervalYearMonthType, UInt8Type,
23    UInt16Type, UInt32Type, UInt64Type,
24};
25use arrow_schema::{DataType, IntervalUnit};
26use chrono::{Datelike, NaiveDateTime};
27use common_decimal::Decimal128;
28use common_error::ext::ErrorExt;
29use common_error::status_code::StatusCode;
30use common_query::{Output, OutputData};
31use common_recordbatch::{
32    RecordBatch, SendableRecordBatchStream, map_dictionary_to_values_data_type,
33};
34use common_telemetry::{debug, error};
35use common_time::{Date, IntervalDayTime, IntervalMonthDayNano, IntervalYearMonth};
36use datafusion_common::ScalarValue;
37use datatypes::prelude::ConcreteDataType;
38use datatypes::schema::SchemaRef;
39use datatypes::types::jsonb_to_string;
40use futures::StreamExt;
41use opensrv_mysql::{
42    Column, ColumnFlags, ColumnType, ErrorKind, OkResponse, QueryResultWriter, RowWriter,
43    ToMysqlValue,
44};
45use session::SessionRef;
46use session::context::QueryContextRef;
47use snafu::prelude::*;
48use tokio::io::AsyncWrite;
49
50use crate::error::{
51    self, ConvertSqlValueSnafu, DataFusionSnafu, InternalSnafu, NotSupportedSnafu, Result,
52    TimestampOverflowSnafu,
53};
54use crate::metrics::*;
55
56const MYSQL_DATETIME_MIN_YEAR: i32 = 1000;
57const MYSQL_DATETIME_MAX_YEAR: i32 = 9999;
58
59/// Try to write multiple output to the writer if possible.
60pub async fn write_output<W: AsyncWrite + Send + Sync + Unpin>(
61    mut writer: QueryResultWriter<'_, W>,
62    query_context: QueryContextRef,
63    session: SessionRef,
64    outputs: Vec<Result<Output>>,
65) -> Result<()> {
66    if let Some(warning) = query_context.warning() {
67        session.add_warning(warning);
68    }
69
70    enum Response {
71        ResultSet {
72            columns: Vec<Column>,
73            stream: SendableRecordBatchStream,
74        },
75        AffectedRows(usize),
76    }
77
78    let mut responses = Vec::with_capacity(outputs.len());
79    for output in outputs {
80        match output {
81            Ok(x) => {
82                let output = match x.data {
83                    OutputData::Stream(stream) => either::Left(stream),
84                    OutputData::RecordBatches(record_batches) => {
85                        either::Left(record_batches.as_stream())
86                    }
87                    OutputData::AffectedRows(rows) => either::Right(rows),
88                };
89                responses.push(match output {
90                    either::Left(stream) => {
91                        let schema = stream.schema();
92                        let columns = match create_mysql_column_def(&schema) {
93                            Ok(columns) => columns,
94                            Err(e) => {
95                                MysqlResultWriter::write_query_error(
96                                    e,
97                                    writer,
98                                    query_context.clone(),
99                                )
100                                .await?;
101                                return Ok(());
102                            }
103                        };
104                        Response::ResultSet { columns, stream }
105                    }
106                    either::Right(rows) => Response::AffectedRows(rows),
107                });
108            }
109            Err(e) => {
110                MysqlResultWriter::write_query_error(e, writer, query_context.clone()).await?;
111                return Ok(());
112            }
113        }
114    }
115
116    for response in &mut responses {
117        writer = match response {
118            Response::ResultSet { columns, stream } => {
119                let mut row_writer = writer.start(columns).await?;
120                while let Some(record_batch) = stream.next().await {
121                    match record_batch {
122                        Ok(record_batch) => {
123                            if let Err(e) = MysqlResultWriter::write_recordbatch(
124                                &mut row_writer,
125                                record_batch,
126                                query_context.clone(),
127                            )
128                            .await
129                            {
130                                let (kind, err) = handle_err(e, query_context);
131                                row_writer.finish_error(kind, &err.as_bytes()).await?;
132                                return Ok(());
133                            }
134                        }
135                        Err(e) => {
136                            let (kind, err) = handle_err(e, query_context);
137                            row_writer.finish_error(kind, &err.as_bytes()).await?;
138                            return Ok(());
139                        }
140                    }
141                }
142                row_writer.finish_one().await?
143            }
144            Response::AffectedRows(rows) => {
145                MysqlResultWriter::write_affected_rows(writer, *rows, &session).await?
146            }
147        }
148    }
149
150    writer.no_more_results().await?;
151    Ok(())
152}
153
154/// Handle GreptimeDB error, convert it to MySQL error
155pub fn handle_err(e: impl ErrorExt, query_ctx: QueryContextRef) -> (ErrorKind, String) {
156    let status_code = e.status_code();
157    let kind = mysql_error_kind(&status_code);
158
159    if status_code.should_log_error() {
160        let root_error = e.root_cause().unwrap_or(&e);
161        error!(e; "Failed to handle mysql query, code: {}, error: {}, db: {}", status_code, root_error.to_string(), query_ctx.get_db_string());
162    } else {
163        debug!(
164            "Failed to handle mysql query, code: {}, db: {}, error: {:?}",
165            status_code,
166            query_ctx.get_db_string(),
167            e
168        );
169    };
170    let msg = e.output_msg();
171    // Inline the status code to output message for MySQL
172    let err_msg = format!("({status_code}): {msg}");
173
174    (kind, err_msg)
175}
176
177struct MysqlResultWriter;
178
179struct PrecisionTimestamp<'a> {
180    formatted: &'a str,
181    datetime: chrono::NaiveDateTime,
182}
183
184struct StagedTimestamp {
185    datetime: Option<NaiveDateTime>,
186    formatted: String,
187}
188
189impl StagedTimestamp {
190    fn new() -> Self {
191        Self {
192            datetime: None,
193            formatted: String::with_capacity(32),
194        }
195    }
196}
197
198impl<'a> ToMysqlValue for PrecisionTimestamp<'a> {
199    fn to_mysql_text<W: std::io::Write>(&self, w: &mut W) -> io::Result<()> {
200        self.formatted.to_mysql_text(w)
201    }
202
203    fn to_mysql_bin<W: std::io::Write>(&self, w: &mut W, c: &Column) -> io::Result<()> {
204        self.datetime.to_mysql_bin(w, c)
205    }
206}
207
208impl MysqlResultWriter {
209    async fn write_affected_rows<'a, W: AsyncWrite + Unpin>(
210        w: QueryResultWriter<'a, W>,
211        rows: usize,
212        session: &SessionRef,
213    ) -> io::Result<QueryResultWriter<'a, W>> {
214        let warnings = session.warnings_count() as u16;
215
216        let next_writer = w
217            .complete_one(OkResponse {
218                affected_rows: rows as u64,
219                warnings,
220                ..Default::default()
221            })
222            .await?;
223        Ok(next_writer)
224    }
225
226    async fn write_recordbatch<W: AsyncWrite + Unpin>(
227        row_writer: &mut RowWriter<'_, '_, W>,
228        record_batch: RecordBatch,
229        query_context: QueryContextRef,
230    ) -> Result<()> {
231        let schema = record_batch.schema.clone();
232        let record_batch = record_batch.into_df_record_batch();
233        let mut timestamp_slots = vec![None; record_batch.num_columns()];
234        let mut staged_timestamps = record_batch
235            .columns()
236            .iter()
237            .enumerate()
238            .filter(|(_, column)| matches!(column.data_type(), DataType::Timestamp(_, _)))
239            .enumerate()
240            .map(|(slot, (column_index, column))| {
241                timestamp_slots[column_index] = Some(slot);
242                (column, StagedTimestamp::new())
243            })
244            .collect::<Vec<_>>();
245        for i in 0..record_batch.num_rows() {
246            for (column, staged_timestamp) in &mut staged_timestamps {
247                let column = *column;
248                staged_timestamp.datetime = None;
249                staged_timestamp.formatted.clear();
250                if !column.is_null(i) {
251                    let timestamp = datatypes::arrow_array::timestamp_array_value(column, i);
252                    let datetime = timestamp
253                        .to_chrono_datetime_with_timezone(Some(&query_context.timezone()))
254                        .with_context(|| TimestampOverflowSnafu {
255                            error: format!(
256                                "timestamp {} overflow with unit {}",
257                                timestamp.value(),
258                                timestamp.unit()
259                            ),
260                        })?;
261                    let year = datetime.year();
262                    if !(MYSQL_DATETIME_MIN_YEAR..=MYSQL_DATETIME_MAX_YEAR).contains(&year) {
263                        return TimestampOverflowSnafu {
264                            error: format!(
265                                "timestamp {} with unit {} has local year {}, outside MySQL DATETIME range {}..={}",
266                                timestamp.value(),
267                                timestamp.unit(),
268                                year,
269                                MYSQL_DATETIME_MIN_YEAR,
270                                MYSQL_DATETIME_MAX_YEAR,
271                            ),
272                        }
273                        .fail();
274                    }
275                    write!(
276                        &mut staged_timestamp.formatted,
277                        "{}",
278                        datetime.format("%Y-%m-%d %H:%M:%S%.f")
279                    )
280                    .map_err(|_| {
281                        InternalSnafu {
282                            err_msg: "timestamp formatting failed",
283                        }
284                        .build()
285                    })?;
286                    staged_timestamp.datetime = Some(datetime);
287                }
288            }
289
290            for (j, column) in record_batch.columns().iter().enumerate() {
291                if column.is_null(i) {
292                    row_writer.write_col(None::<u8>)?;
293                    continue;
294                }
295
296                match column.data_type() {
297                    DataType::Null => {
298                        row_writer.write_col(None::<u8>)?;
299                    }
300                    DataType::Boolean => {
301                        let array = column.as_boolean();
302                        row_writer.write_col(array.value(i) as i8)?;
303                    }
304                    DataType::UInt8 => {
305                        let array = column.as_primitive::<UInt8Type>();
306                        row_writer.write_col(array.value(i))?;
307                    }
308                    DataType::UInt16 => {
309                        let array = column.as_primitive::<UInt16Type>();
310                        row_writer.write_col(array.value(i))?;
311                    }
312                    DataType::UInt32 => {
313                        let array = column.as_primitive::<UInt32Type>();
314                        row_writer.write_col(array.value(i))?;
315                    }
316                    DataType::UInt64 => {
317                        let array = column.as_primitive::<UInt64Type>();
318                        row_writer.write_col(array.value(i))?;
319                    }
320                    DataType::Int8 => {
321                        let array = column.as_primitive::<Int8Type>();
322                        row_writer.write_col(array.value(i))?;
323                    }
324                    DataType::Int16 => {
325                        let array = column.as_primitive::<Int16Type>();
326                        row_writer.write_col(array.value(i))?;
327                    }
328                    DataType::Int32 => {
329                        let array = column.as_primitive::<Int32Type>();
330                        row_writer.write_col(array.value(i))?;
331                    }
332                    DataType::Int64 => {
333                        let array = column.as_primitive::<Int64Type>();
334                        row_writer.write_col(array.value(i))?;
335                    }
336                    DataType::Float32 => {
337                        let array = column.as_primitive::<Float32Type>();
338                        row_writer.write_col(array.value(i))?;
339                    }
340                    DataType::Float64 => {
341                        let array = column.as_primitive::<Float64Type>();
342                        row_writer.write_col(array.value(i))?;
343                    }
344                    DataType::Utf8 | DataType::Utf8View | DataType::LargeUtf8 => {
345                        let v = datatypes::arrow_array::string_array_value(column, i);
346                        row_writer.write_col(v)?;
347                    }
348                    DataType::Binary | DataType::BinaryView | DataType::LargeBinary => {
349                        let v = datatypes::arrow_array::binary_array_value(column, i);
350                        if let ConcreteDataType::Json(_) = &schema.column_schemas()[j].data_type {
351                            let s = jsonb_to_string(v).context(ConvertSqlValueSnafu)?;
352                            row_writer.write_col(s)?;
353                        } else {
354                            row_writer.write_col(v)?;
355                        }
356                    }
357                    DataType::Date32 => {
358                        let array = column.as_primitive::<Date32Type>();
359                        let v = Date::new(array.value(i));
360                        row_writer.write_col(v.to_chrono_date())?;
361                    }
362                    DataType::Timestamp(_, _) => {
363                        let slot = timestamp_slots
364                            .get(j)
365                            .context(InternalSnafu {
366                                err_msg: "timestamp column index is invalid",
367                            })?
368                            .as_ref()
369                            .context(InternalSnafu {
370                                err_msg: "timestamp column has no staging slot",
371                            })?;
372                        let (_, staged_timestamp) =
373                            staged_timestamps.get(*slot).context(InternalSnafu {
374                                err_msg: "timestamp staging slot is missing",
375                            })?;
376                        let datetime = staged_timestamp.datetime.context(InternalSnafu {
377                            err_msg: "timestamp staging value is missing",
378                        })?;
379                        row_writer.write_col(PrecisionTimestamp {
380                            formatted: staged_timestamp.formatted.as_str(),
381                            datetime,
382                        })?;
383                    }
384                    DataType::Interval(interval_unit) => match interval_unit {
385                        IntervalUnit::YearMonth => {
386                            let array = column.as_primitive::<IntervalYearMonthType>();
387                            let v: IntervalYearMonth = array.value(i).into();
388                            row_writer.write_col(v.to_iso8601_string())?;
389                        }
390                        IntervalUnit::DayTime => {
391                            let array = column.as_primitive::<IntervalDayTimeType>();
392                            let v: IntervalDayTime = array.value(i).into();
393                            row_writer.write_col(v.to_iso8601_string())?;
394                        }
395                        IntervalUnit::MonthDayNano => {
396                            let array = column.as_primitive::<IntervalMonthDayNanoType>();
397                            let v: IntervalMonthDayNano = array.value(i).into();
398                            row_writer.write_col(v.to_iso8601_string())?;
399                        }
400                    },
401                    DataType::Duration(_) => {
402                        let v: Duration =
403                            datatypes::arrow_array::duration_array_value(column, i).into();
404                        row_writer.write_col(v)?;
405                    }
406                    DataType::List(_) | DataType::Struct(_) => {
407                        let v = ScalarValue::try_from_array(column, i).context(DataFusionSnafu)?;
408                        row_writer.write_col(v.to_string())?;
409                    }
410                    DataType::Time32(_) | DataType::Time64(_) => {
411                        let time = datatypes::arrow_array::time_array_value(column, i);
412                        let v = time.to_timezone_aware_string(Some(&query_context.timezone()));
413                        row_writer.write_col(v)?;
414                    }
415                    DataType::Decimal128(precision, scale) => {
416                        let array = column.as_primitive::<Decimal128Type>();
417                        let v = Decimal128::new(array.value(i), *precision, *scale);
418                        row_writer.write_col(v.to_string())?;
419                    }
420                    _ => {
421                        return NotSupportedSnafu {
422                            feat: format!("convert {} to MySQL value", column.data_type()),
423                        }
424                        .fail();
425                    }
426                }
427            }
428            row_writer.end_row().await?;
429        }
430        Ok(())
431    }
432
433    async fn write_query_error<'a, W: AsyncWrite + Unpin>(
434        error: impl ErrorExt,
435        w: QueryResultWriter<'a, W>,
436        query_context: QueryContextRef,
437    ) -> io::Result<()> {
438        METRIC_ERROR_COUNTER
439            .with_label_values(&[METRIC_ERROR_COUNTER_LABEL_MYSQL])
440            .inc();
441
442        let (kind, err) = handle_err(error, query_context);
443        debug!("Write query error, kind: {:?}, err: {}", kind, err);
444        w.error(kind, err.as_bytes()).await?;
445        Ok(())
446    }
447}
448
449pub fn create_mysql_column(data_type: &ConcreteDataType, column_name: &str) -> Result<Column> {
450    let data_type = &map_dictionary_to_values_data_type(data_type);
451    let column_type = match data_type {
452        ConcreteDataType::Null(_) => Ok(ColumnType::MYSQL_TYPE_NULL),
453        ConcreteDataType::Boolean(_) | ConcreteDataType::Int8(_) | ConcreteDataType::UInt8(_) => {
454            Ok(ColumnType::MYSQL_TYPE_TINY)
455        }
456        ConcreteDataType::Int16(_) | ConcreteDataType::UInt16(_) => {
457            Ok(ColumnType::MYSQL_TYPE_SHORT)
458        }
459        ConcreteDataType::Int32(_) | ConcreteDataType::UInt32(_) => Ok(ColumnType::MYSQL_TYPE_LONG),
460        ConcreteDataType::Int64(_) | ConcreteDataType::UInt64(_) => {
461            Ok(ColumnType::MYSQL_TYPE_LONGLONG)
462        }
463        ConcreteDataType::Float32(_) => Ok(ColumnType::MYSQL_TYPE_FLOAT),
464        ConcreteDataType::Float64(_) => Ok(ColumnType::MYSQL_TYPE_DOUBLE),
465        ConcreteDataType::Binary(_) | ConcreteDataType::String(_) => {
466            Ok(ColumnType::MYSQL_TYPE_VARCHAR)
467        }
468        ConcreteDataType::Timestamp(_) => Ok(ColumnType::MYSQL_TYPE_TIMESTAMP),
469        ConcreteDataType::Time(_) => Ok(ColumnType::MYSQL_TYPE_TIME),
470        ConcreteDataType::Date(_) => Ok(ColumnType::MYSQL_TYPE_DATE),
471        ConcreteDataType::Interval(_) => Ok(ColumnType::MYSQL_TYPE_VARCHAR),
472        ConcreteDataType::Duration(_) => Ok(ColumnType::MYSQL_TYPE_TIME),
473        ConcreteDataType::Decimal128(_) => Ok(ColumnType::MYSQL_TYPE_DECIMAL),
474        ConcreteDataType::Json(_) => Ok(ColumnType::MYSQL_TYPE_JSON),
475        ConcreteDataType::Vector(_) => Ok(ColumnType::MYSQL_TYPE_BLOB),
476        ConcreteDataType::List(_) => Ok(ColumnType::MYSQL_TYPE_VARCHAR),
477        ConcreteDataType::Struct(_) => Ok(ColumnType::MYSQL_TYPE_VARCHAR),
478        _ => error::UnsupportedDataTypeSnafu {
479            data_type,
480            reason: "not implemented",
481        }
482        .fail(),
483    };
484    let mut colflags = ColumnFlags::empty();
485    match data_type {
486        ConcreteDataType::UInt16(_)
487        | ConcreteDataType::UInt8(_)
488        | ConcreteDataType::UInt32(_)
489        | ConcreteDataType::UInt64(_) => colflags |= ColumnFlags::UNSIGNED_FLAG,
490        _ => {}
491    };
492    column_type.map(|column_type| Column {
493        column: column_name.to_string(),
494        coltype: column_type,
495        // TODO(LFC): Currently "table" and "colflags" are not relevant in MySQL server
496        //   implementation, will revisit them again in the future.
497        table: String::default(),
498        collen: 0, // 0 means "use default".
499        colflags,
500    })
501}
502
503/// Creates MySQL columns definition from our column schema.
504pub fn create_mysql_column_def(schema: &SchemaRef) -> Result<Vec<Column>> {
505    schema
506        .column_schemas()
507        .iter()
508        .map(|column_schema| create_mysql_column(&column_schema.data_type, &column_schema.name))
509        .collect()
510}
511
512fn mysql_error_kind(status_code: &StatusCode) -> ErrorKind {
513    match status_code {
514        StatusCode::Success => ErrorKind::ER_YES,
515        StatusCode::Unknown | StatusCode::External => ErrorKind::ER_UNKNOWN_ERROR,
516        StatusCode::Unsupported => ErrorKind::ER_NOT_SUPPORTED_YET,
517        StatusCode::Cancelled | StatusCode::DeadlineExceeded => ErrorKind::ER_QUERY_INTERRUPTED,
518        StatusCode::RuntimeResourcesExhausted => ErrorKind::ER_OUT_OF_RESOURCES,
519        StatusCode::InvalidSyntax => ErrorKind::ER_SYNTAX_ERROR,
520        StatusCode::RegionAlreadyExists | StatusCode::TableAlreadyExists => {
521            ErrorKind::ER_TABLE_EXISTS_ERROR
522        }
523        StatusCode::RegionNotFound | StatusCode::TableNotFound => ErrorKind::ER_NO_SUCH_TABLE,
524        StatusCode::RegionReadonly => ErrorKind::ER_READ_ONLY_MODE,
525        StatusCode::DatabaseNotFound => ErrorKind::ER_WRONG_DB_NAME,
526        StatusCode::UserNotFound => ErrorKind::ER_NO_SUCH_USER,
527        StatusCode::UnsupportedPasswordType => ErrorKind::ER_PASSWORD_FORMAT,
528        StatusCode::PermissionDenied | StatusCode::AccessDenied => {
529            ErrorKind::ER_ACCESS_DENIED_ERROR
530        }
531        StatusCode::UserPasswordMismatch => ErrorKind::ER_DBACCESS_DENIED_ERROR,
532        StatusCode::InvalidAuthHeader | StatusCode::AuthHeaderNotFound => {
533            ErrorKind::ER_NOT_SUPPORTED_AUTH_MODE
534        }
535        StatusCode::Unexpected
536        | StatusCode::Internal
537        | StatusCode::IllegalState
538        | StatusCode::PlanQuery
539        | StatusCode::EngineExecuteQuery
540        | StatusCode::RegionNotReady
541        | StatusCode::RegionBusy
542        | StatusCode::TableUnavailable
543        | StatusCode::StorageUnavailable
544        | StatusCode::RequestOutdated => ErrorKind::ER_INTERNAL_ERROR,
545        StatusCode::InvalidArguments => ErrorKind::ER_WRONG_ARGUMENTS,
546        StatusCode::TableColumnNotFound => ErrorKind::ER_BAD_FIELD_ERROR,
547        StatusCode::TableColumnExists => ErrorKind::ER_DUP_FIELDNAME,
548        StatusCode::DatabaseAlreadyExists => ErrorKind::ER_DB_CREATE_EXISTS,
549        StatusCode::RateLimited => ErrorKind::ER_TOO_MANY_CONCURRENT_TRXS,
550        StatusCode::FlowAlreadyExists => ErrorKind::ER_TABLE_EXISTS_ERROR,
551        StatusCode::FlowNotFound => ErrorKind::ER_NO_SUCH_TABLE,
552        StatusCode::TriggerAlreadyExists => ErrorKind::ER_TABLE_EXISTS_ERROR,
553        StatusCode::TriggerNotFound => ErrorKind::ER_NO_SUCH_TABLE,
554        StatusCode::Suspended => ErrorKind::ER_SERVER_SHUTDOWN,
555    }
556}
servers/mysql/writer.rs

servers/mysql/
writer.rs
