use std::ops::{ControlFlow, Deref};
use crate::error::Error;
use sqlparser::ast::{Expr, Insert, Statement, VisitMut, VisitorMut};
use sqlparser::ast::{Query, SetExpr, TopQuantity, Value};
use sqlparser::dialect::Dialect;
use sqlparser::parser::Parser;
use std::ops::DerefMut;
pub fn normalize(dialect: &dyn Dialect, sql: &str) -> Result<Vec<String>, Error> {
Normalizer::normalize(dialect, sql, NormalizerOptions::new())
}
pub fn normalize_with_options(
dialect: &dyn Dialect,
sql: &str,
options: NormalizerOptions,
) -> Result<Vec<String>, Error> {
Normalizer::normalize(dialect, sql, options)
}
#[derive(Default, Clone)]
pub struct NormalizerOptions {
pub unify_in_list: bool,
pub unify_values: bool,
pub alphabetize_insert_columns: bool,
}
impl NormalizerOptions {
pub fn new() -> Self {
Self::default()
}
pub fn with_unify_in_list(mut self, unify_in_list: bool) -> Self {
self.unify_in_list = unify_in_list;
self
}
pub fn with_unify_values(mut self, unify_values: bool) -> Self {
self.unify_values = unify_values;
self
}
pub fn with_alphabetize_insert_columns(mut self, alphabetize_insert_columns: bool) -> Self {
self.alphabetize_insert_columns = alphabetize_insert_columns;
self
}
}
#[derive(Default)]
pub struct Normalizer {
pub options: NormalizerOptions,
}
impl VisitorMut for Normalizer {
type Break = ();
fn post_visit_query(&mut self, query: &mut Query) -> ControlFlow<Self::Break> {
normalize_top(query.body.deref_mut());
if let SetExpr::Values(values) = query.body.deref_mut() {
if self.options.unify_values {
let rows = &mut values.rows;
if rows.is_empty()
|| rows.iter().all(|row| {
row.is_empty() || row.iter().all(|expr| matches!(expr, Expr::Value(_)))
})
{
*rows = vec![vec![Expr::Value(
Value::Placeholder("...".into()).with_empty_span(),
)]];
}
}
}
ControlFlow::Continue(())
}
fn post_visit_statement(
&mut self,
stmt: &mut sqlparser::ast::Statement,
) -> ControlFlow<Self::Break> {
if self.options.alphabetize_insert_columns {
if let Statement::Insert(Insert {
columns,
after_columns,
source,
..
}) = stmt
{
if let Some(Query { body, .. }) = source.as_deref() {
if let SetExpr::Values(v) = body.deref() {
if v.rows
== vec![vec![Expr::Value(
Value::Placeholder("...".into()).with_empty_span(),
)]]
{
if columns.len() > 1 {
columns.sort_by_key(|s| s.value.to_lowercase());
}
if after_columns.len() > 1 {
after_columns.sort_by_key(|s| s.value.to_lowercase());
}
}
}
}
}
}
ControlFlow::Continue(())
}
fn pre_visit_expr(&mut self, expr: &mut Expr) -> ControlFlow<Self::Break> {
if let Expr::UnaryOp { op: _, expr: child } = expr {
if Self::is_unary_chain_over_value(child) {
*expr = Expr::Value(Value::Placeholder("?".into()).with_empty_span());
}
}
ControlFlow::Continue(())
}
fn pre_visit_value(&mut self, value: &mut Value) -> ControlFlow<Self::Break> {
*value = Value::Placeholder("?".into());
ControlFlow::Continue(())
}
fn post_visit_expr(&mut self, expr: &mut Expr) -> ControlFlow<Self::Break> {
match expr {
Expr::InList { list, .. }
if self.options.unify_in_list
&& list.iter().all(Self::contains_only_tuples_of_values) =>
{
*list = vec![Expr::Value(
Value::Placeholder("...".into()).with_empty_span(),
)];
}
_ => {}
}
ControlFlow::Continue(())
}
}
impl Normalizer {
pub fn new() -> Self {
Self::default()
}
pub fn with_options(mut self, options: NormalizerOptions) -> Self {
self.options = options;
self
}
pub fn normalize(
dialect: &dyn Dialect,
sql: &str,
options: NormalizerOptions,
) -> Result<Vec<String>, Error> {
let mut statements = Parser::parse_sql(dialect, sql)?;
let _ = statements.visit(&mut Self::new().with_options(options));
Ok(statements
.into_iter()
.map(|statement| statement.to_string())
.collect::<Vec<String>>())
}
fn is_unary_chain_over_value(expr: &Expr) -> bool {
match expr {
Expr::Value(_) => true,
Expr::UnaryOp { expr: child, .. } => Self::is_unary_chain_over_value(child),
_ => false,
}
}
fn contains_only_tuples_of_values(expr: &Expr) -> bool {
match expr {
Expr::Value(_) => true,
Expr::Tuple(v) => v.iter().all(Self::contains_only_tuples_of_values),
_ => false,
}
}
}
fn normalize_top(body: &mut SetExpr) {
match body {
SetExpr::Select(select) => {
if let Some(top) = &mut select.top {
if matches!(top.quantity, Some(TopQuantity::Constant(_))) {
top.quantity = Some(TopQuantity::Expr(Expr::Value(
Value::Placeholder("?".into()).with_empty_span(),
)));
}
}
}
SetExpr::SetOperation { left, right, .. } => {
normalize_top(left);
normalize_top(right);
}
_ => {}
}
}