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, SubsecRound};
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                        // MySQL temporal types support at most microsecond precision.
279                        datetime.trunc_subsecs(6).format("%Y-%m-%d %H:%M:%S%.f")
280                    )
281                    .map_err(|_| {
282                        InternalSnafu {
283                            err_msg: "timestamp formatting failed",
284                        }
285                        .build()
286                    })?;
287                    staged_timestamp.datetime = Some(datetime);
288                }
289            }
290
291            for (j, column) in record_batch.columns().iter().enumerate() {
292                if column.is_null(i) {
293                    row_writer.write_col(None::<u8>)?;
294                    continue;
295                }
296
297                match column.data_type() {
298                    DataType::Null => {
299                        row_writer.write_col(None::<u8>)?;
300                    }
301                    DataType::Boolean => {
302                        let array = column.as_boolean();
303                        row_writer.write_col(array.value(i) as i8)?;
304                    }
305                    DataType::UInt8 => {
306                        let array = column.as_primitive::<UInt8Type>();
307                        row_writer.write_col(array.value(i))?;
308                    }
309                    DataType::UInt16 => {
310                        let array = column.as_primitive::<UInt16Type>();
311                        row_writer.write_col(array.value(i))?;
312                    }
313                    DataType::UInt32 => {
314                        let array = column.as_primitive::<UInt32Type>();
315                        row_writer.write_col(array.value(i))?;
316                    }
317                    DataType::UInt64 => {
318                        let array = column.as_primitive::<UInt64Type>();
319                        row_writer.write_col(array.value(i))?;
320                    }
321                    DataType::Int8 => {
322                        let array = column.as_primitive::<Int8Type>();
323                        row_writer.write_col(array.value(i))?;
324                    }
325                    DataType::Int16 => {
326                        let array = column.as_primitive::<Int16Type>();
327                        row_writer.write_col(array.value(i))?;
328                    }
329                    DataType::Int32 => {
330                        let array = column.as_primitive::<Int32Type>();
331                        row_writer.write_col(array.value(i))?;
332                    }
333                    DataType::Int64 => {
334                        let array = column.as_primitive::<Int64Type>();
335                        row_writer.write_col(array.value(i))?;
336                    }
337                    DataType::Float32 => {
338                        let array = column.as_primitive::<Float32Type>();
339                        row_writer.write_col(array.value(i))?;
340                    }
341                    DataType::Float64 => {
342                        let array = column.as_primitive::<Float64Type>();
343                        row_writer.write_col(array.value(i))?;
344                    }
345                    DataType::Utf8 | DataType::Utf8View | DataType::LargeUtf8 => {
346                        let v = datatypes::arrow_array::string_array_value(column, i);
347                        row_writer.write_col(v)?;
348                    }
349                    DataType::Binary | DataType::BinaryView | DataType::LargeBinary => {
350                        let v = datatypes::arrow_array::binary_array_value(column, i);
351                        if let ConcreteDataType::Json(_) = &schema.column_schemas()[j].data_type {
352                            let s = jsonb_to_string(v).context(ConvertSqlValueSnafu)?;
353                            row_writer.write_col(s)?;
354                        } else {
355                            row_writer.write_col(v)?;
356                        }
357                    }
358                    DataType::Date32 => {
359                        let array = column.as_primitive::<Date32Type>();
360                        let v = Date::new(array.value(i));
361                        row_writer.write_col(v.to_chrono_date())?;
362                    }
363                    DataType::Timestamp(_, _) => {
364                        let slot = timestamp_slots
365                            .get(j)
366                            .context(InternalSnafu {
367                                err_msg: "timestamp column index is invalid",
368                            })?
369                            .as_ref()
370                            .context(InternalSnafu {
371                                err_msg: "timestamp column has no staging slot",
372                            })?;
373                        let (_, staged_timestamp) =
374                            staged_timestamps.get(*slot).context(InternalSnafu {
375                                err_msg: "timestamp staging slot is missing",
376                            })?;
377                        let datetime = staged_timestamp.datetime.context(InternalSnafu {
378                            err_msg: "timestamp staging value is missing",
379                        })?;
380                        row_writer.write_col(PrecisionTimestamp {
381                            formatted: staged_timestamp.formatted.as_str(),
382                            datetime,
383                        })?;
384                    }
385                    DataType::Interval(interval_unit) => match interval_unit {
386                        IntervalUnit::YearMonth => {
387                            let array = column.as_primitive::<IntervalYearMonthType>();
388                            let v: IntervalYearMonth = array.value(i).into();
389                            row_writer.write_col(v.to_iso8601_string())?;
390                        }
391                        IntervalUnit::DayTime => {
392                            let array = column.as_primitive::<IntervalDayTimeType>();
393                            let v: IntervalDayTime = array.value(i).into();
394                            row_writer.write_col(v.to_iso8601_string())?;
395                        }
396                        IntervalUnit::MonthDayNano => {
397                            let array = column.as_primitive::<IntervalMonthDayNanoType>();
398                            let v: IntervalMonthDayNano = array.value(i).into();
399                            row_writer.write_col(v.to_iso8601_string())?;
400                        }
401                    },
402                    DataType::Duration(_) => {
403                        let v: Duration =
404                            datatypes::arrow_array::duration_array_value(column, i).into();
405                        row_writer.write_col(v)?;
406                    }
407                    DataType::List(_) | DataType::Struct(_) => {
408                        let v = ScalarValue::try_from_array(column, i).context(DataFusionSnafu)?;
409                        row_writer.write_col(v.to_string())?;
410                    }
411                    DataType::Time32(_) | DataType::Time64(_) => {
412                        let time = datatypes::arrow_array::time_array_value(column, i);
413                        let v = time.to_timezone_aware_string(Some(&query_context.timezone()));
414                        row_writer.write_col(v)?;
415                    }
416                    DataType::Decimal128(precision, scale) => {
417                        let array = column.as_primitive::<Decimal128Type>();
418                        let v = Decimal128::new(array.value(i), *precision, *scale);
419                        row_writer.write_col(v.to_string())?;
420                    }
421                    _ => {
422                        return NotSupportedSnafu {
423                            feat: format!("convert {} to MySQL value", column.data_type()),
424                        }
425                        .fail();
426                    }
427                }
428            }
429            row_writer.end_row().await?;
430        }
431        Ok(())
432    }
433
434    async fn write_query_error<'a, W: AsyncWrite + Unpin>(
435        error: impl ErrorExt,
436        w: QueryResultWriter<'a, W>,
437        query_context: QueryContextRef,
438    ) -> io::Result<()> {
439        METRIC_ERROR_COUNTER
440            .with_label_values(&[METRIC_ERROR_COUNTER_LABEL_MYSQL])
441            .inc();
442
443        let (kind, err) = handle_err(error, query_context);
444        debug!("Write query error, kind: {:?}, err: {}", kind, err);
445        w.error(kind, err.as_bytes()).await?;
446        Ok(())
447    }
448}
449
450pub fn create_mysql_column(data_type: &ConcreteDataType, column_name: &str) -> Result<Column> {
451    let data_type = &map_dictionary_to_values_data_type(data_type);
452    let column_type = match data_type {
453        ConcreteDataType::Null(_) => Ok(ColumnType::MYSQL_TYPE_NULL),
454        ConcreteDataType::Boolean(_) | ConcreteDataType::Int8(_) | ConcreteDataType::UInt8(_) => {
455            Ok(ColumnType::MYSQL_TYPE_TINY)
456        }
457        ConcreteDataType::Int16(_) | ConcreteDataType::UInt16(_) => {
458            Ok(ColumnType::MYSQL_TYPE_SHORT)
459        }
460        ConcreteDataType::Int32(_) | ConcreteDataType::UInt32(_) => Ok(ColumnType::MYSQL_TYPE_LONG),
461        ConcreteDataType::Int64(_) | ConcreteDataType::UInt64(_) => {
462            Ok(ColumnType::MYSQL_TYPE_LONGLONG)
463        }
464        ConcreteDataType::Float32(_) => Ok(ColumnType::MYSQL_TYPE_FLOAT),
465        ConcreteDataType::Float64(_) => Ok(ColumnType::MYSQL_TYPE_DOUBLE),
466        ConcreteDataType::Binary(_) | ConcreteDataType::String(_) => {
467            Ok(ColumnType::MYSQL_TYPE_VARCHAR)
468        }
469        ConcreteDataType::Timestamp(_) => Ok(ColumnType::MYSQL_TYPE_TIMESTAMP),
470        ConcreteDataType::Time(_) => Ok(ColumnType::MYSQL_TYPE_TIME),
471        ConcreteDataType::Date(_) => Ok(ColumnType::MYSQL_TYPE_DATE),
472        ConcreteDataType::Interval(_) => Ok(ColumnType::MYSQL_TYPE_VARCHAR),
473        ConcreteDataType::Duration(_) => Ok(ColumnType::MYSQL_TYPE_TIME),
474        ConcreteDataType::Decimal128(_) => Ok(ColumnType::MYSQL_TYPE_DECIMAL),
475        ConcreteDataType::Json(_) => Ok(ColumnType::MYSQL_TYPE_JSON),
476        ConcreteDataType::Vector(_) => Ok(ColumnType::MYSQL_TYPE_BLOB),
477        ConcreteDataType::List(_) => Ok(ColumnType::MYSQL_TYPE_VARCHAR),
478        ConcreteDataType::Struct(_) => Ok(ColumnType::MYSQL_TYPE_VARCHAR),
479        _ => error::UnsupportedDataTypeSnafu {
480            data_type,
481            reason: "not implemented",
482        }
483        .fail(),
484    };
485    let mut colflags = ColumnFlags::empty();
486    match data_type {
487        ConcreteDataType::UInt16(_)
488        | ConcreteDataType::UInt8(_)
489        | ConcreteDataType::UInt32(_)
490        | ConcreteDataType::UInt64(_) => colflags |= ColumnFlags::UNSIGNED_FLAG,
491        _ => {}
492    };
493    column_type.map(|column_type| Column {
494        column: column_name.to_string(),
495        coltype: column_type,
496        // TODO(LFC): Currently "table" and "colflags" are not relevant in MySQL server
497        //   implementation, will revisit them again in the future.
498        table: String::default(),
499        collen: 0, // 0 means "use default".
500        colflags,
501    })
502}
503
504/// Creates MySQL columns definition from our column schema.
505pub fn create_mysql_column_def(schema: &SchemaRef) -> Result<Vec<Column>> {
506    schema
507        .column_schemas()
508        .iter()
509        .map(|column_schema| create_mysql_column(&column_schema.data_type, &column_schema.name))
510        .collect()
511}
512
513fn mysql_error_kind(status_code: &StatusCode) -> ErrorKind {
514    match status_code {
515        StatusCode::Success => ErrorKind::ER_YES,
516        StatusCode::PartialFailure | StatusCode::Unknown | StatusCode::External => {
517            ErrorKind::ER_UNKNOWN_ERROR
518        }
519        StatusCode::Unsupported => ErrorKind::ER_NOT_SUPPORTED_YET,
520        StatusCode::Cancelled | StatusCode::DeadlineExceeded => ErrorKind::ER_QUERY_INTERRUPTED,
521        StatusCode::RuntimeResourcesExhausted => ErrorKind::ER_OUT_OF_RESOURCES,
522        StatusCode::InvalidSyntax => ErrorKind::ER_SYNTAX_ERROR,
523        StatusCode::RegionAlreadyExists | StatusCode::TableAlreadyExists => {
524            ErrorKind::ER_TABLE_EXISTS_ERROR
525        }
526        StatusCode::RegionNotFound | StatusCode::TableNotFound => ErrorKind::ER_NO_SUCH_TABLE,
527        StatusCode::RegionReadonly => ErrorKind::ER_READ_ONLY_MODE,
528        StatusCode::DatabaseNotFound => ErrorKind::ER_WRONG_DB_NAME,
529        StatusCode::UserNotFound => ErrorKind::ER_NO_SUCH_USER,
530        StatusCode::UnsupportedPasswordType => ErrorKind::ER_PASSWORD_FORMAT,
531        StatusCode::PermissionDenied | StatusCode::AccessDenied => {
532            ErrorKind::ER_ACCESS_DENIED_ERROR
533        }
534        StatusCode::UserPasswordMismatch => ErrorKind::ER_DBACCESS_DENIED_ERROR,
535        StatusCode::InvalidAuthHeader | StatusCode::AuthHeaderNotFound => {
536            ErrorKind::ER_NOT_SUPPORTED_AUTH_MODE
537        }
538        StatusCode::Unexpected
539        | StatusCode::Internal
540        | StatusCode::IllegalState
541        | StatusCode::PlanQuery
542        | StatusCode::EngineExecuteQuery
543        | StatusCode::RegionNotReady
544        | StatusCode::RegionBusy
545        | StatusCode::TableUnavailable
546        | StatusCode::StorageUnavailable
547        | StatusCode::RequestOutdated => ErrorKind::ER_INTERNAL_ERROR,
548        StatusCode::InvalidArguments => ErrorKind::ER_WRONG_ARGUMENTS,
549        StatusCode::TableColumnNotFound => ErrorKind::ER_BAD_FIELD_ERROR,
550        StatusCode::TableColumnExists => ErrorKind::ER_DUP_FIELDNAME,
551        StatusCode::DatabaseAlreadyExists => ErrorKind::ER_DB_CREATE_EXISTS,
552        StatusCode::RateLimited => ErrorKind::ER_TOO_MANY_CONCURRENT_TRXS,
553        StatusCode::FlowAlreadyExists => ErrorKind::ER_TABLE_EXISTS_ERROR,
554        StatusCode::FlowNotFound => ErrorKind::ER_NO_SUCH_TABLE,
555        StatusCode::TriggerAlreadyExists => ErrorKind::ER_TABLE_EXISTS_ERROR,
556        StatusCode::TriggerNotFound => ErrorKind::ER_NO_SUCH_TABLE,
557        StatusCode::Suspended => ErrorKind::ER_SERVER_SHUTDOWN,
558    }
559}
servers/mysql/writer.rs

servers/mysql/
writer.rs
