use std::collections::HashMap;
use std::ops::ControlFlow;
use std::time::Duration as StdDuration;
use itertools::Itertools;
use lazy_static::lazy_static;
use regex::Regex;
use sqlparser::ast::{DataType, Expr, Interval, Value};
use crate::statements::transform::TransformRule;
lazy_static! {
static ref INTERVAL_ABBREVIATION_PATTERN: Regex = Regex::new(r"([+-]?\d+|[a-zA-Z]+|\+|-)").unwrap();
static ref IS_VALID_ISO_8601_PREFIX_PATTERN: Regex = Regex::new(r"^[-]?[Pp]").unwrap();
static ref INTERVAL_ABBREVIATION_MAPPING: HashMap<&'static str, &'static str> = HashMap::from([
("y","years"),
("mon","months"),
("w","weeks"),
("d","days"),
("h","hours"),
("m","minutes"),
("s","seconds"),
("millis","milliseconds"),
("ms","milliseconds"),
("us","microseconds"),
("ns","nanoseconds"),
]);
}
pub(crate) struct ExpandIntervalTransformRule;
impl TransformRule for ExpandIntervalTransformRule {
fn visit_expr(&self, expr: &mut Expr) -> ControlFlow<()> {
match expr {
Expr::Interval(interval) => match &*interval.value {
Expr::Value(Value::SingleQuotedString(value))
| Expr::Value(Value::DoubleQuotedString(value)) => {
if let Some(normalized_name) = normalize_interval_name(value) {
*expr = update_existing_interval_with_value(
interval,
single_quoted_string_expr(normalized_name),
);
}
}
Expr::BinaryOp { left, op, right } => match &**left {
Expr::Value(Value::SingleQuotedString(value))
| Expr::Value(Value::DoubleQuotedString(value)) => {
if let Some(normalized_name) = normalize_interval_name(value) {
let new_expr_value = Box::new(Expr::BinaryOp {
left: single_quoted_string_expr(normalized_name),
op: op.clone(),
right: right.clone(),
});
*expr = update_existing_interval_with_value(interval, new_expr_value);
}
}
_ => {}
},
_ => {}
},
Expr::Cast {
expr: cast_exp,
data_type,
..
} => {
if DataType::Interval == *data_type {
match &**cast_exp {
Expr::Value(Value::SingleQuotedString(value))
| Expr::Value(Value::DoubleQuotedString(value)) => {
let interval_name =
normalize_interval_name(value).unwrap_or_else(|| value.to_string());
*expr = create_interval(single_quoted_string_expr(interval_name));
}
_ => {}
}
}
}
_ => {}
}
ControlFlow::<()>::Continue(())
}
}
fn single_quoted_string_expr(string: String) -> Box<Expr> {
Box::new(Expr::Value(Value::SingleQuotedString(string)))
}
fn create_interval(value: Box<Expr>) -> Expr {
Expr::Interval(Interval {
value,
leading_field: None,
leading_precision: None,
last_field: None,
fractional_seconds_precision: None,
})
}
fn update_existing_interval_with_value(interval: &Interval, value: Box<Expr>) -> Expr {
Expr::Interval(Interval {
value,
leading_field: interval.leading_field.clone(),
leading_precision: interval.leading_precision,
last_field: interval.last_field.clone(),
fractional_seconds_precision: interval.fractional_seconds_precision,
})
}
fn normalize_interval_name(interval_str: &str) -> Option<String> {
if interval_str.contains(char::is_whitespace) {
return None;
}
if IS_VALID_ISO_8601_PREFIX_PATTERN.is_match(interval_str) {
return parse_iso8601_interval(interval_str);
}
expand_interval_abbreviation(interval_str)
}
fn parse_iso8601_interval(signed_iso: &str) -> Option<String> {
let (is_negative, unsigned_iso) = if let Some(stripped) = signed_iso.strip_prefix('-') {
(true, stripped)
} else {
(false, signed_iso)
};
match iso8601::duration(&unsigned_iso.to_uppercase()) {
Ok(duration) => {
let millis = StdDuration::from(duration).as_millis();
let sign = if is_negative { "-" } else { "" };
Some(format!("{}{} milliseconds", sign, millis))
}
Err(_) => None,
}
}
fn expand_interval_abbreviation(interval_str: &str) -> Option<String> {
Some(
INTERVAL_ABBREVIATION_PATTERN
.find_iter(interval_str)
.map(|mat| {
let mat_str = mat.as_str();
*INTERVAL_ABBREVIATION_MAPPING
.get(mat_str)
.unwrap_or(&mat_str)
})
.join(" "),
)
}
#[cfg(test)]
mod tests {
use std::ops::ControlFlow;
use sqlparser::ast::{BinaryOperator, DataType, Expr, Interval, Value};
use crate::statements::transform::expand_interval::{
create_interval, normalize_interval_name, single_quoted_string_expr,
ExpandIntervalTransformRule,
};
use crate::statements::transform::TransformRule;
#[test]
fn test_transform_interval_basic_conversions() {
let test_cases = vec![
("1y", "1 years"),
("4mon", "4 months"),
("-3w", "-3 weeks"),
("55h", "55 hours"),
("3d", "3 days"),
("5s", "5 seconds"),
("2m", "2 minutes"),
("100millis", "100 milliseconds"),
("200ms", "200 milliseconds"),
("350us", "350 microseconds"),
("400ns", "400 nanoseconds"),
];
for (input, expected) in test_cases {
let result = normalize_interval_name(input).unwrap();
assert_eq!(result, expected);
}
let test_cases = vec!["1 year 2 months 3 days 4 hours", "-2 months"];
for input in test_cases {
assert_eq!(normalize_interval_name(input), None);
}
}
#[test]
fn test_transform_interval_compound_conversions() {
let test_cases = vec![
("2y4mon6w", "2 years 4 months 6 weeks"),
("5d3h1m", "5 days 3 hours 1 minutes"),
(
"10s312ms789ns",
"10 seconds 312 milliseconds 789 nanoseconds",
),
(
"23millis987us754ns",
"23 milliseconds 987 microseconds 754 nanoseconds",
),
("-1d-5h", "-1 days -5 hours"),
("-2y-4mon-6w", "-2 years -4 months -6 weeks"),
("-5d-3h-1m", "-5 days -3 hours -1 minutes"),
(
"-10s-312ms-789ns",
"-10 seconds -312 milliseconds -789 nanoseconds",
),
(
"-23millis-987us-754ns",
"-23 milliseconds -987 microseconds -754 nanoseconds",
),
];
for (input, expected) in test_cases {
let result = normalize_interval_name(input).unwrap();
assert_eq!(result, expected);
}
}
#[test]
fn test_iso8601_format() {
assert_eq!(
normalize_interval_name("P1Y2M3DT4H5M6S"),
Some("36993906000 milliseconds".to_string())
);
assert_eq!(
normalize_interval_name("p3y3m700dt133h17m36.789s"),
Some("163343856789 milliseconds".to_string())
);
assert_eq!(
normalize_interval_name("-P1Y2M3DT4H5M6S"),
Some("-36993906000 milliseconds".to_string())
);
assert_eq!(normalize_interval_name("P1_INVALID_ISO8601"), None);
}
#[test]
fn test_visit_expr_when_interval_is_single_quoted_string_abbr_expr() {
let interval_transformation_rule = ExpandIntervalTransformRule {};
let mut string_expr = create_interval(single_quoted_string_expr("5y".to_string()));
let control_flow = interval_transformation_rule.visit_expr(&mut string_expr);
assert_eq!(control_flow, ControlFlow::Continue(()));
assert_eq!(
string_expr,
Expr::Interval(Interval {
value: Box::new(Expr::Value(Value::SingleQuotedString(
"5 years".to_string()
))),
leading_field: None,
leading_precision: None,
last_field: None,
fractional_seconds_precision: None,
})
);
}
#[test]
fn test_visit_expr_when_interval_is_single_quoted_string_iso8601_expr() {
let interval_transformation_rule = ExpandIntervalTransformRule {};
let mut string_expr =
create_interval(single_quoted_string_expr("P1Y2M3DT4H5M6S".to_string()));
let control_flow = interval_transformation_rule.visit_expr(&mut string_expr);
assert_eq!(control_flow, ControlFlow::Continue(()));
assert_eq!(
string_expr,
Expr::Interval(Interval {
value: Box::new(Expr::Value(Value::SingleQuotedString(
"36993906000 milliseconds".to_string()
))),
leading_field: None,
leading_precision: None,
last_field: None,
fractional_seconds_precision: None,
})
);
}
#[test]
fn test_visit_expr_when_interval_is_binary_op() {
let interval_transformation_rule = ExpandIntervalTransformRule {};
let binary_op = Box::new(Expr::BinaryOp {
left: single_quoted_string_expr("2d".to_string()),
op: BinaryOperator::Minus,
right: Box::new(create_interval(single_quoted_string_expr("1d".to_string()))),
});
let mut binary_op_expr = create_interval(binary_op);
let control_flow = interval_transformation_rule.visit_expr(&mut binary_op_expr);
assert_eq!(control_flow, ControlFlow::Continue(()));
assert_eq!(
binary_op_expr,
Expr::Interval(Interval {
value: Box::new(Expr::BinaryOp {
left: single_quoted_string_expr("2 days".to_string()),
op: BinaryOperator::Minus,
right: Box::new(Expr::Interval(Interval {
value: single_quoted_string_expr("1d".to_string()),
leading_field: None,
leading_precision: None,
last_field: None,
fractional_seconds_precision: None,
})),
}),
leading_field: None,
leading_precision: None,
last_field: None,
fractional_seconds_precision: None,
})
);
}
#[test]
fn test_visit_expr_when_cast_expr() {
let interval_transformation_rule = ExpandIntervalTransformRule {};
let mut cast_to_interval_expr = Expr::Cast {
expr: single_quoted_string_expr("3y2mon".to_string()),
data_type: DataType::Interval,
format: None,
};
let control_flow = interval_transformation_rule.visit_expr(&mut cast_to_interval_expr);
assert_eq!(control_flow, ControlFlow::Continue(()));
assert_eq!(
cast_to_interval_expr,
Expr::Interval(Interval {
value: Box::new(Expr::Value(Value::SingleQuotedString(
"3 years 2 months".to_string()
))),
leading_field: None,
leading_precision: None,
last_field: None,
fractional_seconds_precision: None,
})
);
let mut cast_to_i64_expr = Expr::Cast {
expr: single_quoted_string_expr("5".to_string()),
data_type: DataType::Int64,
format: None,
};
let control_flow = interval_transformation_rule.visit_expr(&mut cast_to_i64_expr);
assert_eq!(control_flow, ControlFlow::Continue(()));
assert_eq!(
cast_to_i64_expr,
Expr::Cast {
expr: single_quoted_string_expr("5".to_string()),
data_type: DataType::Int64,
format: None,
}
);
}
}