Skip to main content

servers/mysql/
helper.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::ops::ControlFlow;
16use std::time::Duration;
17
18use chrono::{NaiveDate, NaiveDateTime};
19use common_query::prelude::ScalarValue;
20use common_sql::convert::sql_value_to_value;
21use common_time::timestamp::TimeUnit;
22use common_time::{Date, Timestamp, Timezone};
23use datatypes::prelude::{ConcreteDataType, DataType};
24use datatypes::schema::ColumnSchema;
25use datatypes::types::TimestampType;
26use datatypes::value::{self, Value};
27#[cfg(test)]
28use itertools::Itertools;
29use opensrv_mysql::{ParamValue, ValueInner, to_naive_datetime};
30use snafu::ResultExt;
31use sql::ast::{Expr, Value as ValueExpr, ValueWithSpan, VisitMut, visit_expressions_mut};
32use sql::statements::statement::Statement;
33
34use crate::error::{self, Result};
35
36/// Location of a prepared-statement placeholder in the original SQL text.
37#[derive(Debug, Clone, Copy, PartialEq, Eq)]
38pub(crate) struct PlaceholderSpan {
39    /// 1-based placeholder index.
40    pub(crate) index: usize,
41    pub(crate) start_line: u64,
42    pub(crate) start_column: u64,
43    pub(crate) end_line: u64,
44    pub(crate) end_column: u64,
45}
46
47/// Returns the placeholder string "$i".
48pub fn format_placeholder(i: usize) -> String {
49    format!("${}", i)
50}
51
52/// Replace all the "?" placeholder into "$i" in SQL,
53/// returns the new SQL and the last placeholder index.
54#[cfg(test)]
55pub fn replace_placeholders(query: &str) -> (String, usize) {
56    let query_parts = query.split('?').collect::<Vec<_>>();
57    let parts_len = query_parts.len();
58    let mut index = 0;
59    let query = query_parts
60        .into_iter()
61        .enumerate()
62        .map(|(i, part)| {
63            if i == parts_len - 1 {
64                return part.to_string();
65            }
66
67            index += 1;
68            format!("{part}{}", format_placeholder(index))
69        })
70        .join("");
71
72    (query, index + 1)
73}
74
75/// Transform all the "?" placeholders into "$i" and return the number of
76/// transformed placeholders.
77pub fn transform_placeholders_with_count(mut stmt: Statement) -> (Statement, usize) {
78    let count = visit_placeholders(&mut stmt);
79    (stmt, count)
80}
81
82/// Collect spans of "$i" placeholders in a statement.
83pub(crate) fn placeholder_spans(mut stmt: Statement) -> Vec<PlaceholderSpan> {
84    let mut spans = Vec::new();
85    collect_placeholder_spans(&mut stmt, &mut spans);
86    spans
87}
88
89fn collect_placeholder_spans<V>(v: &mut V, spans: &mut Vec<PlaceholderSpan>)
90where
91    V: VisitMut,
92{
93    let _ = visit_expressions_mut(v, |expr| {
94        if let Expr::Value(ValueWithSpan {
95            value: ValueExpr::Placeholder(s),
96            span,
97        }) = expr
98            && let Some(index) = placeholder_index(s)
99        {
100            spans.push(PlaceholderSpan {
101                index,
102                start_line: span.start.line,
103                start_column: span.start.column,
104                end_line: span.end.line,
105                end_column: span.end.column,
106            });
107        }
108        ControlFlow::<()>::Continue(())
109    });
110}
111
112fn placeholder_index(s: &str) -> Option<usize> {
113    s.strip_prefix('$')?
114        .parse::<usize>()
115        .ok()
116        .filter(|i| *i > 0)
117}
118
119fn visit_placeholders<V>(v: &mut V) -> usize
120where
121    V: VisitMut,
122{
123    let mut index = 1;
124    let _ = visit_expressions_mut(v, |expr| {
125        if let Expr::Value(ValueWithSpan {
126            value: ValueExpr::Placeholder(s),
127            ..
128        }) = expr
129            && s == "?"
130        {
131            *s = format_placeholder(index);
132            index += 1;
133        }
134        ControlFlow::<()>::Continue(())
135    });
136    index - 1
137}
138
139/// Convert [`ParamValue`] into [`Value`] according to param type.
140/// It will try it's best to do type conversions if possible
141pub fn convert_value(
142    param: &ParamValue,
143    t: &ConcreteDataType,
144    timezone: &Timezone,
145) -> Result<ScalarValue> {
146    if let ConcreteDataType::Dictionary(dictionary) = t {
147        return Ok(ScalarValue::Dictionary(
148            Box::new(dictionary.key_type().as_arrow_type()),
149            Box::new(convert_value(param, dictionary.value_type(), timezone)?),
150        ));
151    }
152
153    match param.value.into_inner() {
154        ValueInner::Int(i) => match t {
155            ConcreteDataType::Int8(_) => Ok(ScalarValue::Int8(Some(i as i8))),
156            ConcreteDataType::Int16(_) => Ok(ScalarValue::Int16(Some(i as i16))),
157            ConcreteDataType::Int32(_) => Ok(ScalarValue::Int32(Some(i as i32))),
158            ConcreteDataType::Int64(_) => Ok(ScalarValue::Int64(Some(i))),
159            ConcreteDataType::UInt8(_) => Ok(ScalarValue::UInt8(Some(i as u8))),
160            ConcreteDataType::UInt16(_) => Ok(ScalarValue::UInt16(Some(i as u16))),
161            ConcreteDataType::UInt32(_) => Ok(ScalarValue::UInt32(Some(i as u32))),
162            ConcreteDataType::UInt64(_) => Ok(ScalarValue::UInt64(Some(i as u64))),
163            ConcreteDataType::Float32(_) => Ok(ScalarValue::Float32(Some(i as f32))),
164            ConcreteDataType::Float64(_) => Ok(ScalarValue::Float64(Some(i as f64))),
165            ConcreteDataType::Boolean(_) => Ok(ScalarValue::Boolean(Some(i != 0))),
166            ConcreteDataType::Timestamp(ts_type) => Value::Timestamp(ts_type.create_timestamp(i))
167                .try_to_scalar_value(t)
168                .context(error::ConvertScalarValueSnafu),
169
170            _ => error::PreparedStmtTypeMismatchSnafu {
171                expected: t,
172                actual: param.coltype,
173            }
174            .fail(),
175        },
176        ValueInner::UInt(u) => match t {
177            ConcreteDataType::Int8(_) => Ok(ScalarValue::Int8(Some(u as i8))),
178            ConcreteDataType::Int16(_) => Ok(ScalarValue::Int16(Some(u as i16))),
179            ConcreteDataType::Int32(_) => Ok(ScalarValue::Int32(Some(u as i32))),
180            ConcreteDataType::Int64(_) => Ok(ScalarValue::Int64(Some(u as i64))),
181            ConcreteDataType::UInt8(_) => Ok(ScalarValue::UInt8(Some(u as u8))),
182            ConcreteDataType::UInt16(_) => Ok(ScalarValue::UInt16(Some(u as u16))),
183            ConcreteDataType::UInt32(_) => Ok(ScalarValue::UInt32(Some(u as u32))),
184            ConcreteDataType::UInt64(_) => Ok(ScalarValue::UInt64(Some(u))),
185            ConcreteDataType::Float32(_) => Ok(ScalarValue::Float32(Some(u as f32))),
186            ConcreteDataType::Float64(_) => Ok(ScalarValue::Float64(Some(u as f64))),
187            ConcreteDataType::Boolean(_) => Ok(ScalarValue::Boolean(Some(u != 0))),
188            ConcreteDataType::Timestamp(ts_type) => {
189                Value::Timestamp(ts_type.create_timestamp(u as i64))
190                    .try_to_scalar_value(t)
191                    .context(error::ConvertScalarValueSnafu)
192            }
193
194            _ => error::PreparedStmtTypeMismatchSnafu {
195                expected: t,
196                actual: param.coltype,
197            }
198            .fail(),
199        },
200        ValueInner::Double(f) => match t {
201            ConcreteDataType::Int8(_) => Ok(ScalarValue::Int8(Some(f as i8))),
202            ConcreteDataType::Int16(_) => Ok(ScalarValue::Int16(Some(f as i16))),
203            ConcreteDataType::Int32(_) => Ok(ScalarValue::Int32(Some(f as i32))),
204            ConcreteDataType::Int64(_) => Ok(ScalarValue::Int64(Some(f as i64))),
205            ConcreteDataType::UInt8(_) => Ok(ScalarValue::UInt8(Some(f as u8))),
206            ConcreteDataType::UInt16(_) => Ok(ScalarValue::UInt16(Some(f as u16))),
207            ConcreteDataType::UInt32(_) => Ok(ScalarValue::UInt32(Some(f as u32))),
208            ConcreteDataType::UInt64(_) => Ok(ScalarValue::UInt64(Some(f as u64))),
209            ConcreteDataType::Float32(_) => Ok(ScalarValue::Float32(Some(f as f32))),
210            ConcreteDataType::Float64(_) => Ok(ScalarValue::Float64(Some(f))),
211
212            _ => error::PreparedStmtTypeMismatchSnafu {
213                expected: t,
214                actual: param.coltype,
215            }
216            .fail(),
217        },
218        ValueInner::NULL => value::to_null_scalar_value(t).context(error::ConvertScalarValueSnafu),
219        ValueInner::Bytes(b) => match t {
220            ConcreteDataType::String(t) => {
221                let s = String::from_utf8_lossy(b).to_string();
222                if t.is_large() {
223                    Ok(ScalarValue::LargeUtf8(Some(s)))
224                } else {
225                    Ok(ScalarValue::Utf8(Some(s)))
226                }
227            }
228            ConcreteDataType::Binary(_) => Ok(ScalarValue::Binary(Some(b.to_vec()))),
229            ConcreteDataType::Timestamp(ts_type) => {
230                convert_bytes_to_timestamp(b, ts_type, timezone)
231            }
232            ConcreteDataType::Date(_) => convert_bytes_to_date(b),
233            _ => error::PreparedStmtTypeMismatchSnafu {
234                expected: t,
235                actual: param.coltype,
236            }
237            .fail(),
238        },
239        ValueInner::Date(_) => {
240            let date: common_time::Date = NaiveDate::from(param.value).into();
241            Ok(ScalarValue::Date32(Some(date.val())))
242        }
243        ValueInner::Datetime(_) => match t {
244            ConcreteDataType::Timestamp(_) => {
245                let datetime = to_naive_datetime(param.value).map_err(|e| {
246                    error::MysqlValueConversionSnafu {
247                        err_msg: e.to_string(),
248                    }
249                    .build()
250                })?;
251                let timestamp_millis = datetime_to_timestamp_millis(datetime, timezone)?;
252                Ok(ScalarValue::TimestampMillisecond(
253                    Some(timestamp_millis),
254                    None,
255                ))
256            }
257            _ => error::PreparedStmtTypeMismatchSnafu {
258                expected: t,
259                actual: param.coltype,
260            }
261            .fail(),
262        },
263        ValueInner::Time(_) => Ok(ScalarValue::Time64Nanosecond(Some(
264            Duration::from(param.value).as_millis() as i64,
265        ))),
266    }
267}
268
269/// Convert an MySQL expression to a scalar value.
270/// It automatically handles the conversion of strings to numeric values.
271pub fn convert_expr_to_scalar_value(
272    param: &Expr,
273    t: &ConcreteDataType,
274    timezone: &Timezone,
275) -> Result<ScalarValue> {
276    if let ConcreteDataType::Dictionary(dictionary) = t {
277        return Ok(ScalarValue::Dictionary(
278            Box::new(dictionary.key_type().as_arrow_type()),
279            Box::new(convert_expr_to_scalar_value(
280                param,
281                dictionary.value_type(),
282                timezone,
283            )?),
284        ));
285    }
286
287    let column_schema = ColumnSchema::new("", t.clone(), true);
288    match param {
289        Expr::Value(v) => {
290            let v = sql_value_to_value(&column_schema, &v.value, Some(timezone), None, true);
291            match v {
292                Ok(v) => v
293                    .try_to_scalar_value(t)
294                    .context(error::ConvertScalarValueSnafu),
295                Err(e) => error::InvalidParameterSnafu {
296                    reason: e.to_string(),
297                }
298                .fail(),
299            }
300        }
301        Expr::UnaryOp { op, expr } if let Expr::Value(v) = &**expr => {
302            let v = sql_value_to_value(&column_schema, &v.value, Some(timezone), Some(*op), true);
303            match v {
304                Ok(v) => v
305                    .try_to_scalar_value(t)
306                    .context(error::ConvertScalarValueSnafu),
307                Err(e) => error::InvalidParameterSnafu {
308                    reason: e.to_string(),
309                }
310                .fail(),
311            }
312        }
313        _ => error::InvalidParameterSnafu {
314            reason: format!("cannot convert {:?} to scalar value of type {}", param, t),
315        }
316        .fail(),
317    }
318}
319
320/// Interprets a timezone-less datetime in the given timezone and returns the
321/// corresponding epoch timestamp in milliseconds.
322fn datetime_to_timestamp_millis(datetime: NaiveDateTime, timezone: &Timezone) -> Result<i64> {
323    let ts = Timestamp::from_naive_datetime(datetime, timezone)
324        .map_err(|e| {
325            error::MysqlValueConversionSnafu {
326                err_msg: e.to_string(),
327            }
328            .build()
329        })?
330        .convert_to(TimeUnit::Millisecond)
331        .ok_or_else(|| {
332            error::MysqlValueConversionSnafu {
333                err_msg: "Overflow when converting datetime to milliseconds".to_string(),
334            }
335            .build()
336        })?;
337    Ok(ts.value())
338}
339
340fn convert_bytes_to_timestamp(
341    bytes: &[u8],
342    ts_type: &TimestampType,
343    timezone: &Timezone,
344) -> Result<ScalarValue> {
345    let ts = Timestamp::from_str(&String::from_utf8_lossy(bytes), Some(timezone))
346        .map_err(|e| {
347            error::MysqlValueConversionSnafu {
348                err_msg: e.to_string(),
349            }
350            .build()
351        })?
352        .convert_to(ts_type.unit())
353        .ok_or_else(|| {
354            error::MysqlValueConversionSnafu {
355                err_msg: "Overflow when converting timestamp to target unit".to_string(),
356            }
357            .build()
358        })?;
359    match ts_type {
360        TimestampType::Nanosecond(_) => {
361            Ok(ScalarValue::TimestampNanosecond(Some(ts.value()), None))
362        }
363        TimestampType::Microsecond(_) => {
364            Ok(ScalarValue::TimestampMicrosecond(Some(ts.value()), None))
365        }
366        TimestampType::Millisecond(_) => {
367            Ok(ScalarValue::TimestampMillisecond(Some(ts.value()), None))
368        }
369        TimestampType::Second(_) => Ok(ScalarValue::TimestampSecond(Some(ts.value()), None)),
370    }
371}
372
373fn convert_bytes_to_date(bytes: &[u8]) -> Result<ScalarValue> {
374    let date = Date::from_str_utc(&String::from_utf8_lossy(bytes)).map_err(|e| {
375        error::MysqlValueConversionSnafu {
376            err_msg: e.to_string(),
377        }
378        .build()
379    })?;
380
381    Ok(ScalarValue::Date32(Some(date.val())))
382}
383
384#[cfg(test)]
385mod tests {
386    use datatypes::types::{
387        TimestampMicrosecondType, TimestampMillisecondType, TimestampNanosecondType,
388        TimestampSecondType,
389    };
390    use sql::dialect::MySqlDialect;
391    use sql::parser::{ParseOptions, ParserContext};
392
393    use super::*;
394
395    #[test]
396    fn test_format_placeholder() {
397        assert_eq!("$1", format_placeholder(1));
398        assert_eq!("$3", format_placeholder(3));
399    }
400
401    #[test]
402    fn test_replace_placeholders() {
403        let create = "create table demo(host string, ts timestamp time index)";
404        let (sql, index) = replace_placeholders(create);
405        assert_eq!(create, sql);
406        assert_eq!(1, index);
407
408        let insert = "insert into demo values(?,?,?)";
409        let (sql, index) = replace_placeholders(insert);
410        assert_eq!("insert into demo values($1,$2,$3)", sql);
411        assert_eq!(4, index);
412
413        let query = "select from demo where host=? and idc in (select idc from idcs where name=?) and cpu>?";
414        let (sql, index) = replace_placeholders(query);
415        assert_eq!(
416            "select from demo where host=$1 and idc in (select idc from idcs where name=$2) and cpu>$3",
417            sql
418        );
419        assert_eq!(4, index);
420    }
421
422    fn parse_sql(sql: &str) -> Statement {
423        let mut stmts =
424            ParserContext::create_with_dialect(sql, &MySqlDialect {}, ParseOptions::default())
425                .unwrap();
426        stmts.remove(0)
427    }
428
429    #[test]
430    fn test_transform_placeholders() {
431        let insert = parse_sql("insert into demo values(?,?,?)");
432        let (stmt, count) = transform_placeholders_with_count(insert);
433        let Statement::Insert(insert) = stmt else {
434            unreachable!()
435        };
436        assert_eq!(
437            "INSERT INTO demo VALUES ($1, $2, $3)",
438            insert.inner.to_string()
439        );
440        assert_eq!(3, count);
441
442        let delete = parse_sql("delete from demo where host=? and idc=?");
443        let (stmt, count) = transform_placeholders_with_count(delete);
444        let Statement::Delete(delete) = stmt else {
445            unreachable!()
446        };
447        assert_eq!(
448            "DELETE FROM demo WHERE host = $1 AND idc = $2",
449            delete.inner.to_string()
450        );
451        assert_eq!(2, count);
452
453        let select = parse_sql(
454            "select * from demo where host=? and idc in (select idc from idcs where name=?) and cpu>?",
455        );
456        let (stmt, count) = transform_placeholders_with_count(select);
457        let Statement::Query(select) = stmt else {
458            unreachable!()
459        };
460        assert_eq!(
461            "SELECT * FROM demo WHERE host = $1 AND idc IN (SELECT idc FROM idcs WHERE name = $2) AND cpu > $3",
462            select.inner.to_string()
463        );
464        assert_eq!(3, count);
465
466        let select = parse_sql("select '?', ?");
467        let (stmt, count) = transform_placeholders_with_count(select);
468        let Statement::Query(select) = stmt else {
469            unreachable!()
470        };
471        assert_eq!("SELECT '?', $1", select.inner.to_string());
472        assert_eq!(1, count);
473
474        let set = parse_sql("set time_zone = ?");
475        let (stmt, count) = transform_placeholders_with_count(set);
476        assert_eq!("SET time_zone = $1", stmt.to_string());
477        assert_eq!(1, count);
478    }
479
480    #[test]
481    fn test_convert_expr_to_scalar_value() {
482        let utc = Timezone::from_tz_string("UTC").unwrap();
483        let expr = Expr::Value(ValueExpr::Number("123".to_string(), false).into());
484        let t = ConcreteDataType::int32_datatype();
485        let v = convert_expr_to_scalar_value(&expr, &t, &utc).unwrap();
486        assert_eq!(ScalarValue::Int32(Some(123)), v);
487
488        let expr = Expr::Value(ValueExpr::Number("123.456789".to_string(), false).into());
489        let t = ConcreteDataType::float64_datatype();
490        let v = convert_expr_to_scalar_value(&expr, &t, &utc).unwrap();
491        assert_eq!(ScalarValue::Float64(Some(123.456789)), v);
492
493        let expr = Expr::Value(ValueExpr::SingleQuotedString("2001-01-02".to_string()).into());
494        let t = ConcreteDataType::date_datatype();
495        let v = convert_expr_to_scalar_value(&expr, &t, &utc).unwrap();
496        let scalar_v = ScalarValue::Utf8(Some("2001-01-02".to_string()))
497            .cast_to(&arrow_schema::DataType::Date32)
498            .unwrap();
499        assert_eq!(scalar_v, v);
500
501        let expr =
502            Expr::Value(ValueExpr::SingleQuotedString("2001-01-02 03:04:05".to_string()).into());
503        let t = ConcreteDataType::timestamp_microsecond_datatype();
504        let v = convert_expr_to_scalar_value(&expr, &t, &utc).unwrap();
505        let scalar_v = ScalarValue::Utf8(Some("2001-01-02 03:04:05".to_string()))
506            .cast_to(&arrow_schema::DataType::Timestamp(
507                arrow_schema::TimeUnit::Microsecond,
508                None,
509            ))
510            .unwrap();
511        assert_eq!(scalar_v, v);
512
513        let expr = Expr::Value(ValueExpr::SingleQuotedString("hello".to_string()).into());
514        let t = ConcreteDataType::string_datatype();
515        let v = convert_expr_to_scalar_value(&expr, &t, &utc).unwrap();
516        assert_eq!(ScalarValue::Utf8(Some("hello".to_string())), v);
517
518        let t = ConcreteDataType::dictionary_datatype(
519            ConcreteDataType::uint32_datatype(),
520            ConcreteDataType::string_datatype(),
521        );
522        let v = convert_expr_to_scalar_value(&expr, &t, &utc).unwrap();
523        assert_eq!(
524            ScalarValue::Dictionary(
525                Box::new(arrow_schema::DataType::UInt32),
526                Box::new(ScalarValue::Utf8(Some("hello".to_string()))),
527            ),
528            v
529        );
530
531        let expr = Expr::Value(ValueExpr::Null.into());
532        let t = ConcreteDataType::time_microsecond_datatype();
533        let v = convert_expr_to_scalar_value(&expr, &t, &utc).unwrap();
534        assert_eq!(ScalarValue::Time64Microsecond(None), v);
535    }
536
537    #[test]
538    fn test_convert_bytes_to_timestamp() {
539        let test_cases = vec![
540            // input unix timestamp in seconds -> nanosecond.
541            (
542                "2024-12-26 12:00:00",
543                TimestampType::Nanosecond(TimestampNanosecondType),
544                ScalarValue::TimestampNanosecond(Some(1735214400000000000), None),
545            ),
546            // input unix timestamp in seconds -> microsecond.
547            (
548                "2024-12-26 12:00:00",
549                TimestampType::Microsecond(TimestampMicrosecondType),
550                ScalarValue::TimestampMicrosecond(Some(1735214400000000), None),
551            ),
552            // input unix timestamp in seconds -> millisecond.
553            (
554                "2024-12-26 12:00:00",
555                TimestampType::Millisecond(TimestampMillisecondType),
556                ScalarValue::TimestampMillisecond(Some(1735214400000), None),
557            ),
558            // input unix timestamp in seconds -> second.
559            (
560                "2024-12-26 12:00:00",
561                TimestampType::Second(TimestampSecondType),
562                ScalarValue::TimestampSecond(Some(1735214400), None),
563            ),
564            // input unix timestamp in milliseconds -> nanosecond.
565            (
566                "2024-12-26 12:00:00.123",
567                TimestampType::Nanosecond(TimestampNanosecondType),
568                ScalarValue::TimestampNanosecond(Some(1735214400123000000), None),
569            ),
570            // input unix timestamp in milliseconds -> microsecond.
571            (
572                "2024-12-26 12:00:00.123",
573                TimestampType::Microsecond(TimestampMicrosecondType),
574                ScalarValue::TimestampMicrosecond(Some(1735214400123000), None),
575            ),
576            // input unix timestamp in milliseconds -> millisecond.
577            (
578                "2024-12-26 12:00:00.123",
579                TimestampType::Millisecond(TimestampMillisecondType),
580                ScalarValue::TimestampMillisecond(Some(1735214400123), None),
581            ),
582            // input unix timestamp in milliseconds -> second.
583            (
584                "2024-12-26 12:00:00.123",
585                TimestampType::Second(TimestampSecondType),
586                ScalarValue::TimestampSecond(Some(1735214400), None),
587            ),
588            // input unix timestamp in microseconds -> nanosecond.
589            (
590                "2024-12-26 12:00:00.123456",
591                TimestampType::Nanosecond(TimestampNanosecondType),
592                ScalarValue::TimestampNanosecond(Some(1735214400123456000), None),
593            ),
594            // input unix timestamp in microseconds -> microsecond.
595            (
596                "2024-12-26 12:00:00.123456",
597                TimestampType::Microsecond(TimestampMicrosecondType),
598                ScalarValue::TimestampMicrosecond(Some(1735214400123456), None),
599            ),
600            // input unix timestamp in microseconds -> millisecond.
601            (
602                "2024-12-26 12:00:00.123456",
603                TimestampType::Millisecond(TimestampMillisecondType),
604                ScalarValue::TimestampMillisecond(Some(1735214400123), None),
605            ),
606            // input unix timestamp in milliseconds -> second.
607            (
608                "2024-12-26 12:00:00.123456",
609                TimestampType::Second(TimestampSecondType),
610                ScalarValue::TimestampSecond(Some(1735214400), None),
611            ),
612        ];
613
614        let utc = Timezone::from_tz_string("UTC").unwrap();
615        for (input, ts_type, expected) in test_cases {
616            let result = convert_bytes_to_timestamp(input.as_bytes(), &ts_type, &utc).unwrap();
617            assert_eq!(result, expected);
618        }
619    }
620
621    fn utc_millis(year: i32, month: u32, day: u32, hour: u32, minute: u32, second: u32) -> i64 {
622        NaiveDate::from_ymd_opt(year, month, day)
623            .unwrap()
624            .and_hms_opt(hour, minute, second)
625            .unwrap()
626            .and_utc()
627            .timestamp_millis()
628    }
629
630    #[test]
631    fn test_datetime_to_timestamp_millis() {
632        let datetime = NaiveDate::from_ymd_opt(2026, 8, 13)
633            .unwrap()
634            .and_hms_opt(8, 0, 0)
635            .unwrap();
636
637        let utc = Timezone::from_tz_string("UTC").unwrap();
638        assert_eq!(
639            datetime_to_timestamp_millis(datetime, &utc).unwrap(),
640            utc_millis(2026, 8, 13, 8, 0, 0)
641        );
642
643        let shanghai = Timezone::from_tz_string("Asia/Shanghai").unwrap();
644        assert_eq!(
645            datetime_to_timestamp_millis(datetime, &shanghai).unwrap(),
646            utc_millis(2026, 8, 13, 0, 0, 0)
647        );
648
649        let minus_7 = Timezone::from_tz_string("-07:00").unwrap();
650        assert_eq!(
651            datetime_to_timestamp_millis(datetime, &minus_7).unwrap(),
652            utc_millis(2026, 8, 13, 15, 0, 0)
653        );
654
655        // 2026-03-08 02:30 does not exist in America/New_York (DST gap).
656        let new_york = Timezone::from_tz_string("America/New_York").unwrap();
657        let gap = NaiveDate::from_ymd_opt(2026, 3, 8)
658            .unwrap()
659            .and_hms_opt(2, 30, 0)
660            .unwrap();
661        assert!(datetime_to_timestamp_millis(gap, &new_york).is_err());
662
663        // 2026-11-01 01:30 is ambiguous in America/New_York; picks the first (EDT, UTC-4).
664        let ambiguous = NaiveDate::from_ymd_opt(2026, 11, 1)
665            .unwrap()
666            .and_hms_opt(1, 30, 0)
667            .unwrap();
668        assert_eq!(
669            datetime_to_timestamp_millis(ambiguous, &new_york).unwrap(),
670            utc_millis(2026, 11, 1, 5, 30, 0)
671        );
672    }
673
674    #[test]
675    fn test_convert_bytes_to_timestamp_with_timezone() {
676        let shanghai = Timezone::from_tz_string("Asia/Shanghai").unwrap();
677        let result = convert_bytes_to_timestamp(
678            "2026-08-13 08:00:00".as_bytes(),
679            &TimestampType::Millisecond(TimestampMillisecondType),
680            &shanghai,
681        )
682        .unwrap();
683        assert_eq!(
684            result,
685            ScalarValue::TimestampMillisecond(Some(utc_millis(2026, 8, 13, 0, 0, 0)), None)
686        );
687    }
688
689    #[test]
690    fn test_convert_bytes_to_date() {
691        let test_cases = vec![
692            // Standard date format: YYYY-MM-DD
693            ("1970-01-01", ScalarValue::Date32(Some(0))),
694            ("1969-12-31", ScalarValue::Date32(Some(-1))),
695            ("2024-02-29", ScalarValue::Date32(Some(19782))),
696            ("2024-01-01", ScalarValue::Date32(Some(19723))),
697            ("2024-12-31", ScalarValue::Date32(Some(20088))),
698            ("2001-01-02", ScalarValue::Date32(Some(11324))),
699            ("2050-06-14", ScalarValue::Date32(Some(29384))),
700            ("2020-03-15", ScalarValue::Date32(Some(18336))),
701        ];
702
703        for (input, expected) in test_cases {
704            let result = convert_bytes_to_date(input.as_bytes()).unwrap();
705            assert_eq!(result, expected, "Failed for input: {}", input);
706        }
707    }
708}