mirror of
https://github.com/GreptimeTeam/greptimedb.git
synced 2026-09-21 04:35:35 +00:00
fix(mysql): fail closed on unrepresentable timestamps (#8580)
* fix(mysql): reject unrepresentable timestamps Signed-off-by: discord9 <discord9@163.com> * fix(mysql): validate timestamp protocol boundaries Signed-off-by: discord9 <discord9@163.com> --------- Signed-off-by: discord9 <discord9@163.com>
This commit is contained in:
discord9
@@ -20,7 +20,7 @@ use std::time::Duration;
|
||||
|
||||
use arrow::datatypes::TimeUnit as ArrowTimeUnit;
|
||||
use chrono::{
|
||||
DateTime, Days, LocalResult, Months, NaiveDate, NaiveDateTime, NaiveTime, TimeDelta,
|
||||
DateTime, Days, LocalResult, Months, NaiveDate, NaiveDateTime, NaiveTime, Offset, TimeDelta,
|
||||
TimeZone as ChronoTimeZone, Utc,
|
||||
};
|
||||
use serde::{Deserialize, Serialize};
|
||||
@@ -399,12 +399,15 @@ impl Timestamp {
|
||||
}
|
||||
|
||||
pub fn to_chrono_datetime_with_timezone(&self, tz: Option<&Timezone>) -> Option<NaiveDateTime> {
|
||||
let datetime = self.to_chrono_datetime();
|
||||
datetime.map(|v| match tz {
|
||||
Some(Timezone::Offset(offset)) => offset.from_utc_datetime(&v).naive_local(),
|
||||
Some(Timezone::Named(tz)) => tz.from_utc_datetime(&v).naive_local(),
|
||||
None => Utc.from_utc_datetime(&v).naive_local(),
|
||||
})
|
||||
let utc = self.to_chrono_datetime()?;
|
||||
match tz {
|
||||
None => Some(utc),
|
||||
Some(Timezone::Offset(offset)) => utc.checked_add_offset(*offset),
|
||||
Some(Timezone::Named(tz)) => {
|
||||
let offset = tz.offset_from_utc_datetime(&utc).fix();
|
||||
utc.checked_add_offset(offset)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Convert timestamp to chrono date.
|
||||
@@ -1321,6 +1324,37 @@ mod tests {
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_to_chrono_datetime_with_timezone_bounds() {
|
||||
let positive_offset = Timezone::from_tz_string("+08:00").unwrap();
|
||||
assert_eq!(
|
||||
None,
|
||||
Timestamp::MAX_SECOND.to_chrono_datetime_with_timezone(Some(&positive_offset))
|
||||
);
|
||||
|
||||
let negative_offset = Timezone::from_tz_string("-08:00").unwrap();
|
||||
assert_eq!(
|
||||
None,
|
||||
Timestamp::MIN_SECOND.to_chrono_datetime_with_timezone(Some(&negative_offset))
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_to_chrono_datetime_with_named_timezone_summer_offset() {
|
||||
let timestamp = Timestamp::from_str_utc("2024-07-01 12:00:00Z").unwrap();
|
||||
let berlin = Timezone::from_tz_string("Europe/Berlin").unwrap();
|
||||
|
||||
assert_eq!(
|
||||
Some(
|
||||
NaiveDate::from_ymd_opt(2024, 7, 1)
|
||||
.unwrap()
|
||||
.and_hms_opt(14, 0, 0)
|
||||
.unwrap()
|
||||
),
|
||||
timestamp.to_chrono_datetime_with_timezone(Some(&berlin))
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_from_arrow_time_unit() {
|
||||
assert_eq!(TimeUnit::Second, TimeUnit::from(ArrowTimeUnit::Second));
|
||||
|
||||
@@ -23,6 +23,7 @@ use arrow::datatypes::{
|
||||
UInt16Type, UInt32Type, UInt64Type,
|
||||
};
|
||||
use arrow_schema::{DataType, IntervalUnit};
|
||||
use chrono::{Datelike, NaiveDateTime};
|
||||
use common_decimal::Decimal128;
|
||||
use common_error::ext::ErrorExt;
|
||||
use common_error::status_code::StatusCode;
|
||||
@@ -46,9 +47,15 @@ use session::context::QueryContextRef;
|
||||
use snafu::prelude::*;
|
||||
use tokio::io::AsyncWrite;
|
||||
|
||||
use crate::error::{self, ConvertSqlValueSnafu, DataFusionSnafu, NotSupportedSnafu, Result};
|
||||
use crate::error::{
|
||||
self, ConvertSqlValueSnafu, DataFusionSnafu, InternalSnafu, NotSupportedSnafu, Result,
|
||||
TimestampOverflowSnafu,
|
||||
};
|
||||
use crate::metrics::*;
|
||||
|
||||
const MYSQL_DATETIME_MIN_YEAR: i32 = 1000;
|
||||
const MYSQL_DATETIME_MAX_YEAR: i32 = 9999;
|
||||
|
||||
/// Try to write multiple output to the writer if possible.
|
||||
pub async fn write_output<W: AsyncWrite + Send + Sync + Unpin>(
|
||||
mut writer: QueryResultWriter<'_, W>,
|
||||
@@ -174,6 +181,20 @@ struct PrecisionTimestamp<'a> {
|
||||
datetime: chrono::NaiveDateTime,
|
||||
}
|
||||
|
||||
struct StagedTimestamp {
|
||||
datetime: Option<NaiveDateTime>,
|
||||
formatted: String,
|
||||
}
|
||||
|
||||
impl StagedTimestamp {
|
||||
fn new() -> Self {
|
||||
Self {
|
||||
datetime: None,
|
||||
formatted: String::with_capacity(32),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl<'a> ToMysqlValue for PrecisionTimestamp<'a> {
|
||||
fn to_mysql_text<W: std::io::Write>(&self, w: &mut W) -> io::Result<()> {
|
||||
self.formatted.to_mysql_text(w)
|
||||
@@ -209,10 +230,63 @@ impl MysqlResultWriter {
|
||||
) -> Result<()> {
|
||||
let schema = record_batch.schema.clone();
|
||||
let record_batch = record_batch.into_df_record_batch();
|
||||
// Reusable buffer for timestamp formatting — avoids one heap allocation per
|
||||
// timestamp value per row. Cleared and refilled in the Timestamp arm below.
|
||||
let mut format_buf = String::new();
|
||||
let mut timestamp_slots = vec![None; record_batch.num_columns()];
|
||||
let mut staged_timestamps = record_batch
|
||||
.columns()
|
||||
.iter()
|
||||
.enumerate()
|
||||
.filter(|(_, column)| matches!(column.data_type(), DataType::Timestamp(_, _)))
|
||||
.enumerate()
|
||||
.map(|(slot, (column_index, column))| {
|
||||
timestamp_slots[column_index] = Some(slot);
|
||||
(column, StagedTimestamp::new())
|
||||
})
|
||||
.collect::<Vec<_>>();
|
||||
for i in 0..record_batch.num_rows() {
|
||||
for (column, staged_timestamp) in &mut staged_timestamps {
|
||||
let column = *column;
|
||||
staged_timestamp.datetime = None;
|
||||
staged_timestamp.formatted.clear();
|
||||
if !column.is_null(i) {
|
||||
let timestamp = datatypes::arrow_array::timestamp_array_value(column, i);
|
||||
let datetime = timestamp
|
||||
.to_chrono_datetime_with_timezone(Some(&query_context.timezone()))
|
||||
.with_context(|| TimestampOverflowSnafu {
|
||||
error: format!(
|
||||
"timestamp {} overflow with unit {}",
|
||||
timestamp.value(),
|
||||
timestamp.unit()
|
||||
),
|
||||
})?;
|
||||
let year = datetime.year();
|
||||
if !(MYSQL_DATETIME_MIN_YEAR..=MYSQL_DATETIME_MAX_YEAR).contains(&year) {
|
||||
return TimestampOverflowSnafu {
|
||||
error: format!(
|
||||
"timestamp {} with unit {} has local year {}, outside MySQL DATETIME range {}..={}",
|
||||
timestamp.value(),
|
||||
timestamp.unit(),
|
||||
year,
|
||||
MYSQL_DATETIME_MIN_YEAR,
|
||||
MYSQL_DATETIME_MAX_YEAR,
|
||||
),
|
||||
}
|
||||
.fail();
|
||||
}
|
||||
write!(
|
||||
&mut staged_timestamp.formatted,
|
||||
"{}",
|
||||
datetime.format("%Y-%m-%d %H:%M:%S%.f")
|
||||
)
|
||||
.map_err(|_| {
|
||||
InternalSnafu {
|
||||
err_msg: "timestamp formatting failed",
|
||||
}
|
||||
.build()
|
||||
})?;
|
||||
staged_timestamp.datetime = Some(datetime);
|
||||
}
|
||||
}
|
||||
|
||||
for (j, column) in record_batch.columns().iter().enumerate() {
|
||||
if column.is_null(i) {
|
||||
row_writer.write_col(None::<u8>)?;
|
||||
@@ -286,28 +360,26 @@ impl MysqlResultWriter {
|
||||
row_writer.write_col(v.to_chrono_date())?;
|
||||
}
|
||||
DataType::Timestamp(_, _) => {
|
||||
let v = datatypes::arrow_array::timestamp_array_value(column, i);
|
||||
// Reuse `format_buf` to avoid a per-row heap allocation.
|
||||
// This mirrors what `Timestamp::as_formatted_string` does
|
||||
// internally, but writes directly into our pre-allocated buffer.
|
||||
format_buf.clear();
|
||||
if let Some(datetime) =
|
||||
v.to_chrono_datetime_with_timezone(Some(&query_context.timezone()))
|
||||
{
|
||||
let _ = write!(
|
||||
&mut format_buf,
|
||||
"{}",
|
||||
datetime.format("%Y-%m-%d %H:%M:%S%.f")
|
||||
);
|
||||
row_writer.write_col(PrecisionTimestamp {
|
||||
formatted: &format_buf,
|
||||
datetime,
|
||||
let slot = timestamp_slots
|
||||
.get(j)
|
||||
.context(InternalSnafu {
|
||||
err_msg: "timestamp column index is invalid",
|
||||
})?
|
||||
.as_ref()
|
||||
.context(InternalSnafu {
|
||||
err_msg: "timestamp column has no staging slot",
|
||||
})?;
|
||||
} else {
|
||||
let _ =
|
||||
write!(&mut format_buf, "[Timestamp{}: {}]", v.unit(), v.value());
|
||||
row_writer.write_col(&*format_buf)?;
|
||||
}
|
||||
let (_, staged_timestamp) =
|
||||
staged_timestamps.get(*slot).context(InternalSnafu {
|
||||
err_msg: "timestamp staging slot is missing",
|
||||
})?;
|
||||
let datetime = staged_timestamp.datetime.context(InternalSnafu {
|
||||
err_msg: "timestamp staging value is missing",
|
||||
})?;
|
||||
row_writer.write_col(PrecisionTimestamp {
|
||||
formatted: staged_timestamp.formatted.as_str(),
|
||||
datetime,
|
||||
})?;
|
||||
}
|
||||
DataType::Interval(interval_unit) => match interval_unit {
|
||||
IntervalUnit::YearMonth => {
|
||||
|
||||
@@ -16,21 +16,32 @@ use std::net::SocketAddr;
|
||||
use std::sync::Arc;
|
||||
use std::time::Duration;
|
||||
|
||||
use async_trait::async_trait;
|
||||
use auth::tests::{DatabaseAuthInfo, MockUserProvider};
|
||||
use chrono::{Datelike, NaiveDate};
|
||||
use common_catalog::consts::DEFAULT_SCHEMA_NAME;
|
||||
use common_query::Output;
|
||||
use common_recordbatch::RecordBatch;
|
||||
use common_runtime::Builder as RuntimeBuilder;
|
||||
use common_runtime::runtime::BuilderBuild;
|
||||
use common_time::{Timestamp, Timezone};
|
||||
use datafusion_expr::LogicalPlan;
|
||||
use datatypes::prelude::VectorRef;
|
||||
use datatypes::schema::{ColumnSchema, Schema};
|
||||
use datatypes::value::Value;
|
||||
use datatypes::vectors::{Int32Vector, TimestampMicrosecondVector, TimestampSecondVector};
|
||||
use mysql_async::prelude::*;
|
||||
use mysql_async::{Conn, Row, SslOpts};
|
||||
use query::parser::PromQuery;
|
||||
use query::query_engine::DescribeResult;
|
||||
use servers::error::Result;
|
||||
use servers::install_default_crypto_provider;
|
||||
use servers::mysql::server::{MysqlServer, MysqlSpawnConfig, MysqlSpawnRef};
|
||||
use servers::query_handler::sql::{ServerSqlQueryHandlerRef, SqlQueryHandler};
|
||||
use servers::server::Server;
|
||||
use servers::tls::{ReloadableTlsServerConfig, TlsOption};
|
||||
use session::context::QueryContextRef;
|
||||
use sql::statements::statement::Statement;
|
||||
use table::TableRef;
|
||||
use table::test_util::MemTable;
|
||||
|
||||
@@ -45,8 +56,15 @@ struct MysqlOpts<'a> {
|
||||
}
|
||||
|
||||
fn create_mysql_server(table: TableRef, opts: MysqlOpts<'_>) -> Result<Box<dyn Server>> {
|
||||
let _ = install_default_crypto_provider();
|
||||
let query_handler = create_testing_sql_query_handler(table);
|
||||
create_mysql_server_with_query_handler(query_handler, opts)
|
||||
}
|
||||
|
||||
fn create_mysql_server_with_query_handler(
|
||||
query_handler: ServerSqlQueryHandlerRef,
|
||||
opts: MysqlOpts<'_>,
|
||||
) -> Result<Box<dyn Server>> {
|
||||
let _ = install_default_crypto_provider();
|
||||
let io_runtime = RuntimeBuilder::default()
|
||||
.worker_threads(4)
|
||||
.thread_name("mysql-io-handlers")
|
||||
@@ -77,6 +95,59 @@ fn create_mysql_server(table: TableRef, opts: MysqlOpts<'_>) -> Result<Box<dyn S
|
||||
))
|
||||
}
|
||||
|
||||
struct SessionTimezoneQueryHandler {
|
||||
inner: ServerSqlQueryHandlerRef,
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl SqlQueryHandler for SessionTimezoneQueryHandler {
|
||||
async fn do_query(&self, query: &str, query_ctx: QueryContextRef) -> Vec<Result<Output>> {
|
||||
if query == "SET time_zone = '+08:00'" {
|
||||
query_ctx.set_timezone(Timezone::hours_mins_opt(8, 0).unwrap());
|
||||
return vec![Ok(Output::new_with_affected_rows(0))];
|
||||
}
|
||||
|
||||
self.inner.do_query(query, query_ctx).await
|
||||
}
|
||||
|
||||
async fn do_analyze_stream_query(
|
||||
&self,
|
||||
query: &str,
|
||||
query_ctx: QueryContextRef,
|
||||
) -> Result<Output> {
|
||||
self.inner.do_analyze_stream_query(query, query_ctx).await
|
||||
}
|
||||
|
||||
async fn do_exec_plan(
|
||||
&self,
|
||||
plan: LogicalPlan,
|
||||
stmt: Option<Statement>,
|
||||
query_ctx: QueryContextRef,
|
||||
) -> Result<Output> {
|
||||
self.inner.do_exec_plan(plan, stmt, query_ctx).await
|
||||
}
|
||||
|
||||
async fn do_promql_query(
|
||||
&self,
|
||||
query: &PromQuery,
|
||||
query_ctx: QueryContextRef,
|
||||
) -> Vec<Result<Output>> {
|
||||
self.inner.do_promql_query(query, query_ctx).await
|
||||
}
|
||||
|
||||
async fn do_describe(
|
||||
&self,
|
||||
stmt: Statement,
|
||||
query_ctx: QueryContextRef,
|
||||
) -> Result<Option<DescribeResult>> {
|
||||
self.inner.do_describe(stmt, query_ctx).await
|
||||
}
|
||||
|
||||
async fn is_valid_schema(&self, catalog: &str, schema: &str) -> Result<bool> {
|
||||
self.inner.is_valid_schema(catalog, schema).await
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_start_mysql_server() -> Result<()> {
|
||||
let table = MemTable::default_numbers_table();
|
||||
@@ -422,6 +493,386 @@ async fn do_test_query_all_datatypes(server_tls: TlsOption, client_tls: bool) ->
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_mysql_text_protocol_out_of_range_timestamp_fails_closed() -> Result<()> {
|
||||
let timestamp = i64::MAX;
|
||||
assert!(
|
||||
Timestamp::new_second(timestamp)
|
||||
.to_chrono_datetime_with_timezone(None)
|
||||
.is_none(),
|
||||
"the target value must be outside Chrono's timestamp range"
|
||||
);
|
||||
|
||||
let (table, schema) = timestamp_table("out_of_range_timestamp", Some(timestamp));
|
||||
assert!(
|
||||
schema.column_schemas()[1].is_nullable(),
|
||||
"the timestamp field must be nullable while Arrow marks this value valid"
|
||||
);
|
||||
|
||||
let (result, health_check) =
|
||||
query_timestamp_with_mysql_text_protocol(table, "out_of_range_timestamp").await;
|
||||
assert_timestamp_overflow(result);
|
||||
assert_eq!(Some(1), health_check.unwrap());
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_mysql_text_protocol_max_timestamp_with_session_timezone_fails_closed() -> Result<()> {
|
||||
let timestamp = Timestamp::MAX_SECOND.value();
|
||||
let (table, _) = timestamp_table("max_timestamp_with_session_timezone", Some(timestamp));
|
||||
let query_handler = Arc::new(SessionTimezoneQueryHandler {
|
||||
inner: create_testing_sql_query_handler(table),
|
||||
});
|
||||
let mut mysql_server =
|
||||
create_mysql_server_with_query_handler(query_handler, Default::default())?;
|
||||
let listening = "127.0.0.1:0".parse::<SocketAddr>().unwrap();
|
||||
mysql_server.start(listening).await.unwrap();
|
||||
|
||||
let server_addr = mysql_server.bind_addr().unwrap();
|
||||
let mut connection = create_connection_default_db_name(server_addr.port(), false)
|
||||
.await
|
||||
.unwrap();
|
||||
connection
|
||||
.query_drop("SET time_zone = '+08:00'")
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let result = match connection
|
||||
.query_iter("SELECT id, ts FROM max_timestamp_with_session_timezone")
|
||||
.await
|
||||
{
|
||||
Ok(mut result) => result.collect::<MysqlTextRow>().await,
|
||||
Err(error) => Err(error),
|
||||
};
|
||||
assert_timestamp_overflow_with_value(result, timestamp);
|
||||
assert_eq!(Some(1), connection.query_first("SELECT 1").await.unwrap());
|
||||
|
||||
mysql_server.shutdown().await.unwrap();
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_mysql_binary_protocol_out_of_range_timestamp_fails_closed() -> Result<()> {
|
||||
let timestamp = i64::MAX;
|
||||
assert!(
|
||||
Timestamp::new_second(timestamp)
|
||||
.to_chrono_datetime_with_timezone(None)
|
||||
.is_none(),
|
||||
"the target value must be outside Chrono's timestamp range"
|
||||
);
|
||||
|
||||
let (table, schema) = timestamp_table("out_of_range_timestamp", Some(timestamp));
|
||||
assert!(
|
||||
schema.column_schemas()[1].is_nullable(),
|
||||
"the timestamp field must be nullable while Arrow marks this value valid"
|
||||
);
|
||||
|
||||
let (result, health_check) =
|
||||
query_timestamp_with_mysql_binary_protocol(table, "out_of_range_timestamp").await;
|
||||
assert_timestamp_overflow(result);
|
||||
assert_eq!(Some(1), health_check.unwrap());
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_mysql_binary_protocol_chrono_representable_year_overflow_fails_closed() -> Result<()>
|
||||
{
|
||||
let timestamp = Timestamp::new_second(
|
||||
NaiveDate::from_ymd_opt(100_000, 1, 1)
|
||||
.unwrap()
|
||||
.and_hms_opt(0, 0, 0)
|
||||
.unwrap()
|
||||
.and_utc()
|
||||
.timestamp(),
|
||||
);
|
||||
assert_eq!(
|
||||
100_000,
|
||||
timestamp
|
||||
.to_chrono_datetime_with_timezone(None)
|
||||
.unwrap()
|
||||
.year(),
|
||||
"the target value must be Chrono-representable"
|
||||
);
|
||||
|
||||
let (table, schema) = timestamp_table(
|
||||
"chrono_representable_year_overflow",
|
||||
Some(timestamp.value()),
|
||||
);
|
||||
assert!(
|
||||
schema.column_schemas()[1].is_nullable(),
|
||||
"the timestamp field must be nullable while Arrow marks this value valid"
|
||||
);
|
||||
|
||||
let (result, health_check) =
|
||||
query_timestamp_with_mysql_binary_protocol(table, "chrono_representable_year_overflow")
|
||||
.await;
|
||||
assert_timestamp_overflow_with_value(result, timestamp.value());
|
||||
assert_eq!(Some(1), health_check.unwrap());
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_mysql_text_protocol_timestamp_controls() -> Result<()> {
|
||||
let (table, _) = timestamp_table("in_range_timestamp", Some(0));
|
||||
let (result, health_check) =
|
||||
query_timestamp_with_mysql_text_protocol(table, "in_range_timestamp").await;
|
||||
let rows = result.unwrap();
|
||||
assert_eq!(
|
||||
vec![
|
||||
Value::from(b"7".to_vec()),
|
||||
Value::from(b"1970-01-01 00:00:00".to_vec())
|
||||
],
|
||||
rows[0].values
|
||||
);
|
||||
assert_eq!(Some(1), health_check.unwrap());
|
||||
|
||||
let (table, _) = timestamp_table("null_timestamp", None);
|
||||
let (result, health_check) =
|
||||
query_timestamp_with_mysql_text_protocol(table, "null_timestamp").await;
|
||||
let rows = result.unwrap();
|
||||
assert_eq!(
|
||||
vec![Value::from(b"7".to_vec()), Value::Null],
|
||||
rows[0].values
|
||||
);
|
||||
assert_eq!(Some(1), health_check.unwrap());
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_mysql_binary_protocol_timestamp_controls() -> Result<()> {
|
||||
let (table, _) = timestamp_table("in_range_timestamp", Some(0));
|
||||
let (result, health_check) =
|
||||
query_timestamp_with_mysql_binary_protocol(table, "in_range_timestamp").await;
|
||||
let rows = result.unwrap();
|
||||
assert_eq!(1, rows.len());
|
||||
assert_eq!(Some(&mysql_async::Value::Int(7)), rows[0].as_ref(0));
|
||||
assert_eq!(
|
||||
Some(&mysql_async::Value::Date(1970, 1, 1, 0, 0, 0, 0)),
|
||||
rows[0].as_ref(1)
|
||||
);
|
||||
assert_eq!(Some(1), health_check.unwrap());
|
||||
|
||||
let (table, _) = timestamp_table("null_timestamp", None);
|
||||
let (result, health_check) =
|
||||
query_timestamp_with_mysql_binary_protocol(table, "null_timestamp").await;
|
||||
let rows = result.unwrap();
|
||||
assert_eq!(1, rows.len());
|
||||
assert_eq!(Some(&mysql_async::Value::Int(7)), rows[0].as_ref(0));
|
||||
assert_eq!(Some(&mysql_async::Value::NULL), rows[0].as_ref(1));
|
||||
assert_eq!(Some(1), health_check.unwrap());
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_mysql_timestamp_precision_slot_isolation() -> Result<()> {
|
||||
let (result, health_check) = query_mysql_text_protocol(
|
||||
precision_timestamp_table("precision_timestamps"),
|
||||
"SELECT id, ts_a, ts_b FROM precision_timestamps".to_string(),
|
||||
)
|
||||
.await;
|
||||
let rows = result.unwrap();
|
||||
assert_eq!(
|
||||
vec![
|
||||
Value::from(b"7".to_vec()),
|
||||
Value::from(b"1970-01-01 00:00:00.123456".to_vec()),
|
||||
Value::Null,
|
||||
],
|
||||
rows[0].values
|
||||
);
|
||||
assert_eq!(
|
||||
vec![
|
||||
Value::from(b"8".to_vec()),
|
||||
Value::Null,
|
||||
Value::from(b"1970-01-01 00:00:00.987654".to_vec()),
|
||||
],
|
||||
rows[1].values
|
||||
);
|
||||
assert_eq!(Some(1), health_check.unwrap());
|
||||
|
||||
let (result, health_check) = query_mysql_binary_protocol(
|
||||
precision_timestamp_table("precision_timestamps"),
|
||||
"SELECT id, ts_a, ts_b FROM precision_timestamps".to_string(),
|
||||
)
|
||||
.await;
|
||||
let rows = result.unwrap();
|
||||
assert_eq!(2, rows.len());
|
||||
assert_eq!(Some(&mysql_async::Value::Int(7)), rows[0].as_ref(0));
|
||||
assert_eq!(
|
||||
Some(&mysql_async::Value::Date(1970, 1, 1, 0, 0, 0, 123_456)),
|
||||
rows[0].as_ref(1)
|
||||
);
|
||||
assert_eq!(Some(&mysql_async::Value::NULL), rows[0].as_ref(2));
|
||||
assert_eq!(Some(&mysql_async::Value::Int(8)), rows[1].as_ref(0));
|
||||
assert_eq!(Some(&mysql_async::Value::NULL), rows[1].as_ref(1));
|
||||
assert_eq!(
|
||||
Some(&mysql_async::Value::Date(1970, 1, 1, 0, 0, 0, 987_654)),
|
||||
rows[1].as_ref(2)
|
||||
);
|
||||
assert_eq!(Some(1), health_check.unwrap());
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn timestamp_table(table_name: &str, timestamp: Option<i64>) -> (TableRef, Arc<Schema>) {
|
||||
let schema = Arc::new(Schema::new(vec![
|
||||
ColumnSchema::new(
|
||||
"id",
|
||||
datatypes::prelude::ConcreteDataType::int32_datatype(),
|
||||
false,
|
||||
),
|
||||
ColumnSchema::new(
|
||||
"ts",
|
||||
datatypes::prelude::ConcreteDataType::timestamp_second_datatype(),
|
||||
true,
|
||||
),
|
||||
]));
|
||||
let recordbatch = RecordBatch::new(
|
||||
schema.clone(),
|
||||
vec![
|
||||
Arc::new(Int32Vector::from_values(vec![7])) as VectorRef,
|
||||
Arc::new(TimestampSecondVector::from(vec![timestamp])) as VectorRef,
|
||||
],
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
(MemTable::table(table_name, recordbatch), schema)
|
||||
}
|
||||
|
||||
fn precision_timestamp_table(table_name: &str) -> TableRef {
|
||||
let schema = Arc::new(Schema::new(vec![
|
||||
ColumnSchema::new(
|
||||
"id",
|
||||
datatypes::prelude::ConcreteDataType::int32_datatype(),
|
||||
false,
|
||||
),
|
||||
ColumnSchema::new(
|
||||
"ts_a",
|
||||
datatypes::prelude::ConcreteDataType::timestamp_microsecond_datatype(),
|
||||
true,
|
||||
),
|
||||
ColumnSchema::new(
|
||||
"ts_b",
|
||||
datatypes::prelude::ConcreteDataType::timestamp_microsecond_datatype(),
|
||||
true,
|
||||
),
|
||||
]));
|
||||
let recordbatch = RecordBatch::new(
|
||||
schema,
|
||||
vec![
|
||||
Arc::new(Int32Vector::from_values(vec![7, 8])) as VectorRef,
|
||||
Arc::new(TimestampMicrosecondVector::from(vec![Some(123_456), None])) as VectorRef,
|
||||
Arc::new(TimestampMicrosecondVector::from(vec![None, Some(987_654)])) as VectorRef,
|
||||
],
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
MemTable::table(table_name, recordbatch)
|
||||
}
|
||||
|
||||
async fn query_timestamp_with_mysql_text_protocol(
|
||||
table: TableRef,
|
||||
table_name: &str,
|
||||
) -> (
|
||||
mysql_async::Result<Vec<MysqlTextRow>>,
|
||||
mysql_async::Result<Option<u8>>,
|
||||
) {
|
||||
query_mysql_text_protocol(table, format!("SELECT id, ts FROM {table_name}")).await
|
||||
}
|
||||
|
||||
async fn query_mysql_text_protocol(
|
||||
table: TableRef,
|
||||
query: String,
|
||||
) -> (
|
||||
mysql_async::Result<Vec<MysqlTextRow>>,
|
||||
mysql_async::Result<Option<u8>>,
|
||||
) {
|
||||
let mut mysql_server = create_mysql_server(table, Default::default()).unwrap();
|
||||
let listening = "127.0.0.1:0".parse::<SocketAddr>().unwrap();
|
||||
mysql_server.start(listening).await.unwrap();
|
||||
|
||||
let server_addr = mysql_server.bind_addr().unwrap();
|
||||
let mut connection = create_connection_default_db_name(server_addr.port(), false)
|
||||
.await
|
||||
.unwrap();
|
||||
let result = match connection.query_iter(query).await {
|
||||
Ok(mut result) => result.collect::<MysqlTextRow>().await,
|
||||
Err(error) => Err(error),
|
||||
};
|
||||
let health_check = connection.query_first("SELECT 1").await;
|
||||
|
||||
mysql_server.shutdown().await.unwrap();
|
||||
(result, health_check)
|
||||
}
|
||||
|
||||
async fn query_timestamp_with_mysql_binary_protocol(
|
||||
table: TableRef,
|
||||
table_name: &str,
|
||||
) -> (
|
||||
mysql_async::Result<Vec<Row>>,
|
||||
mysql_async::Result<Option<u8>>,
|
||||
) {
|
||||
query_mysql_binary_protocol(table, format!("SELECT id, ts FROM {table_name}")).await
|
||||
}
|
||||
|
||||
async fn query_mysql_binary_protocol(
|
||||
table: TableRef,
|
||||
query: String,
|
||||
) -> (
|
||||
mysql_async::Result<Vec<Row>>,
|
||||
mysql_async::Result<Option<u8>>,
|
||||
) {
|
||||
let mut mysql_server = create_mysql_server(table, Default::default()).unwrap();
|
||||
let listening = "127.0.0.1:0".parse::<SocketAddr>().unwrap();
|
||||
mysql_server.start(listening).await.unwrap();
|
||||
|
||||
let server_addr = mysql_server.bind_addr().unwrap();
|
||||
let mut connection = create_connection_default_db_name(server_addr.port(), false)
|
||||
.await
|
||||
.unwrap();
|
||||
let result = match connection.prep(query).await {
|
||||
Ok(statement) => connection.exec(statement, ()).await,
|
||||
Err(error) => Err(error),
|
||||
};
|
||||
let health_check = connection.query_first("SELECT 1").await;
|
||||
|
||||
mysql_server.shutdown().await.unwrap();
|
||||
(result, health_check)
|
||||
}
|
||||
|
||||
fn assert_timestamp_overflow<T>(result: mysql_async::Result<T>) {
|
||||
assert_timestamp_overflow_with_value(result, i64::MAX);
|
||||
}
|
||||
|
||||
fn assert_timestamp_overflow_with_value<T>(result: mysql_async::Result<T>, timestamp: i64) {
|
||||
match result {
|
||||
Err(mysql_async::Error::Server(error)) => {
|
||||
assert_eq!(1210, error.code, "expected ER_WRONG_ARGUMENTS");
|
||||
assert!(
|
||||
error.message.contains("Timestamp overflow"),
|
||||
"expected timestamp overflow marker, got: {}",
|
||||
error.message
|
||||
);
|
||||
assert!(
|
||||
error.message.contains(×tamp.to_string()),
|
||||
"expected raw timestamp value, got: {}",
|
||||
error.message
|
||||
);
|
||||
assert!(
|
||||
error.message.contains("Second"),
|
||||
"expected timestamp unit, got: {}",
|
||||
error.message
|
||||
);
|
||||
}
|
||||
Err(error) => panic!("expected a MySQL server timestamp overflow error, got: {error}"),
|
||||
Ok(_) => panic!("out-of-range timestamp must fail closed instead of returning a row"),
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
|
||||
async fn test_query_concurrently() -> Result<()> {
|
||||
common_telemetry::init_default_ut_logging();
|
||||
|
||||
Reference in New Issue
Block a user