89 lines
2.1 KiB
Rust
89 lines
2.1 KiB
Rust
use sqlparser::{
|
|
ast::{Query, SetExpr, Statement},
|
|
dialect::PostgreSqlDialect,
|
|
parser::Parser,
|
|
};
|
|
|
|
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
|
pub enum SqlRoute {
|
|
Read,
|
|
Write,
|
|
}
|
|
|
|
pub fn route_sql(sql: &str) -> SqlRoute {
|
|
let trimmed = sql.trim_start();
|
|
|
|
if starts_with_hint(trimmed, "write") {
|
|
return SqlRoute::Write;
|
|
}
|
|
|
|
if starts_with_hint(trimmed, "read") {
|
|
return SqlRoute::Read;
|
|
}
|
|
|
|
if is_read_query_by_ast(sql) {
|
|
SqlRoute::Read
|
|
} else {
|
|
SqlRoute::Write
|
|
}
|
|
}
|
|
|
|
fn starts_with_hint(sql: &str, hint: &str) -> bool {
|
|
let expected = format!("/*+ {hint} */");
|
|
sql.starts_with(&expected)
|
|
}
|
|
|
|
fn is_read_query_by_ast(sql: &str) -> bool {
|
|
let dialect = PostgreSqlDialect {};
|
|
|
|
let Ok(statements) = Parser::parse_sql(&dialect, sql) else {
|
|
return false;
|
|
};
|
|
|
|
if statements.is_empty() {
|
|
return false;
|
|
}
|
|
|
|
statements.iter().all(is_read_statement)
|
|
}
|
|
|
|
fn is_read_statement(statement: &Statement) -> bool {
|
|
match statement {
|
|
Statement::Query(query) => is_read_query_ast(query),
|
|
Statement::ShowVariable { .. }
|
|
| Statement::ShowVariables { .. }
|
|
| Statement::ShowCreate { .. }
|
|
| Statement::ShowColumns { .. }
|
|
| Statement::ShowTables { .. } => true,
|
|
_ => false,
|
|
}
|
|
}
|
|
|
|
fn is_read_query_ast(query: &Query) -> bool {
|
|
if !query.locks.is_empty() {
|
|
return false;
|
|
}
|
|
|
|
match query.body.as_ref() {
|
|
SetExpr::Select(_) => true,
|
|
SetExpr::Query(inner) => is_read_query_ast(inner),
|
|
SetExpr::SetOperation { left, right, .. } => {
|
|
is_read_set_expr(left) && is_read_set_expr(right)
|
|
}
|
|
SetExpr::Values(_) => true,
|
|
_ => false,
|
|
}
|
|
}
|
|
|
|
fn is_read_set_expr(expr: &SetExpr) -> bool {
|
|
match expr {
|
|
SetExpr::Select(_) => true,
|
|
SetExpr::Query(query) => is_read_query_ast(query),
|
|
SetExpr::SetOperation { left, right, .. } => {
|
|
is_read_set_expr(left) && is_read_set_expr(right)
|
|
}
|
|
SetExpr::Values(_) => true,
|
|
_ => false,
|
|
}
|
|
}
|