use vibesql_parser::{Keyword, Lexer, Span, Token};
fn tokenize(sql: &str) -> Option<Vec<(Token, Span)>> {
let mut tokens = Lexer::new(sql).tokenize_with_spans().ok()?;
if matches!(tokens.last(), Some((Token::Eof, _))) {
tokens.pop();
}
Some(tokens)
}
fn column_list_parens(tokens: &[(Token, Span)]) -> Option<(usize, usize)> {
let open = tokens.iter().position(|(t, _)| matches!(t, Token::LParen))?;
let mut depth = 0usize;
for (i, (tok, _)) in tokens.iter().enumerate().skip(open) {
match tok {
Token::LParen => depth += 1,
Token::RParen => {
depth -= 1;
if depth == 0 {
return Some((open, i));
}
}
_ => {}
}
}
None
}
pub fn append_column(create_sql: &str, coldef: &str) -> Option<String> {
let coldef = coldef.trim();
if coldef.is_empty() {
return None;
}
let tokens = tokenize(create_sql)?;
let (_open, close) = column_list_parens(&tokens)?;
let insert_at = tokens[close].1.start;
let mut out = String::with_capacity(create_sql.len() + coldef.len() + 2);
out.push_str(&create_sql[..insert_at]);
out.push_str(", ");
out.push_str(coldef);
out.push_str(&create_sql[insert_at..]);
Some(out)
}
fn ident_matches(tok: &Token, name: &str) -> bool {
match tok {
Token::Identifier(s) | Token::DelimitedIdentifier(s) => s.eq_ignore_ascii_case(name),
_ => false,
}
}
pub fn rename_table(create_sql: &str, new_name: &str) -> Option<String> {
let tokens = tokenize(create_sql)?;
let name_idx = table_name_index(&tokens)?;
let span = tokens[name_idx].1;
Some(replace_span(create_sql, span, "e_ident(new_name)))
}
fn table_name_index(tokens: &[(Token, Span)]) -> Option<usize> {
let mut i = 0;
if !matches!(tokens.get(i), Some((Token::Keyword { keyword: Keyword::Create, .. }, _))) {
return None;
}
i += 1;
if matches!(
tokens.get(i),
Some((Token::Keyword { keyword: Keyword::Temp | Keyword::Temporary, .. }, _))
) {
i += 1;
}
if !matches!(tokens.get(i), Some((Token::Keyword { keyword: Keyword::Table, .. }, _))) {
return None;
}
i += 1;
if matches!(tokens.get(i), Some((Token::Keyword { keyword: Keyword::If, .. }, _)))
&& matches!(tokens.get(i + 1), Some((Token::Keyword { keyword: Keyword::Not, .. }, _)))
&& matches!(tokens.get(i + 2), Some((Token::Keyword { keyword: Keyword::Exists, .. }, _)))
{
i += 3;
}
match tokens.get(i) {
Some((Token::Identifier(_) | Token::DelimitedIdentifier(_), _)) => {}
_ => return None,
}
if matches!(tokens.get(i + 1), Some((Token::Symbol('.'), _))) {
i += 2; if !matches!(tokens.get(i), Some((Token::Identifier(_) | Token::DelimitedIdentifier(_), _)))
{
return None;
}
}
Some(i)
}
pub fn rename_column(create_sql: &str, old_col: &str, new_col: &str) -> Option<String> {
let tokens = tokenize(create_sql)?;
let (open, close) = column_list_parens(&tokens)?;
let mut depth = 0usize;
let mut at_def_start = false;
let mut target: Option<usize> = None;
for (i, (tok, _)) in tokens.iter().enumerate().take(close).skip(open) {
match tok {
Token::LParen => {
if depth == 1 {
at_def_start = false;
}
depth += 1;
if depth == 1 {
at_def_start = true; }
}
Token::RParen => depth -= 1,
Token::Comma if depth == 1 => at_def_start = true,
_ if depth == 1 => {
if at_def_start {
at_def_start = false;
if ident_matches(tok, old_col) {
if target.is_some() {
return None;
}
target = Some(i);
}
}
}
_ => {}
}
}
let idx = target?;
let span = tokens[idx].1;
let replacement =
if is_safe_bare_identifier(new_col) { new_col.to_string() } else { quote_ident(new_col) };
Some(replace_span(create_sql, span, &replacement))
}
struct DefStart {
tok_idx: usize,
is_column: bool,
}
fn definition_starts(tokens: &[(Token, Span)], open: usize, close: usize) -> Vec<DefStart> {
let mut defs = Vec::new();
let mut depth = 0usize;
let mut at_def_start = false;
for (i, (tok, _)) in tokens.iter().enumerate().take(close).skip(open) {
match tok {
Token::LParen => {
depth += 1;
if depth == 1 {
at_def_start = true; }
}
Token::RParen => depth -= 1,
Token::Comma if depth == 1 => at_def_start = true,
_ if depth == 1 && at_def_start => {
at_def_start = false;
let is_column = matches!(tok, Token::Identifier(_) | Token::DelimitedIdentifier(_));
defs.push(DefStart { tok_idx: i, is_column });
}
_ => {}
}
}
defs
}
pub fn drop_column(create_sql: &str, col: &str) -> Option<String> {
let tokens = tokenize(create_sql)?;
let (open, close) = column_list_parens(&tokens)?;
let defs = definition_starts(&tokens, open, close);
let mut target: Option<usize> = None;
for (di, def) in defs.iter().enumerate() {
if def.is_column && ident_matches(&tokens[def.tok_idx].0, col) {
if target.is_some() {
return None; }
target = Some(di);
}
}
let di = target?;
let column_count = defs.iter().filter(|d| d.is_column).count();
if column_count <= 1 {
return None;
}
let name_start = tokens[defs[di].tok_idx].1.start;
let next_col = defs[di + 1..].iter().find(|d| d.is_column);
let bytes = create_sql.as_bytes();
let (del_start, del_end) = if let Some(next) = next_col {
(name_start, tokens[next.tok_idx].1.start)
} else {
let end = match defs[di + 1..].first() {
Some(constraint) => {
let mut e = tokens[constraint.tok_idx].1.start;
while e > 0 && bytes[e - 1] != b',' {
e -= 1;
}
if e == 0 || bytes[e - 1] != b',' {
return None;
}
e - 1
}
None => tokens[close].1.start,
};
let mut start = name_start;
while start > 0 && bytes[start - 1] != b',' {
start -= 1;
}
if start == 0 || bytes[start - 1] != b',' {
return None;
}
(start - 1, end)
};
let mut out = String::with_capacity(create_sql.len());
out.push_str(&create_sql[..del_start]);
out.push_str(&create_sql[del_end..]);
Some(out)
}
pub fn extract_add_column_text(alter_sql: &str) -> Option<String> {
let tokens = tokenize(alter_sql)?;
let add_pos = tokens
.iter()
.position(|(t, _)| matches!(t, Token::Keyword { keyword: Keyword::Add, .. }))?;
let mut start_tok = add_pos + 1;
if matches!(tokens.get(start_tok), Some((Token::Keyword { keyword: Keyword::Column, .. }, _))) {
start_tok += 1;
}
let first = tokens.get(start_tok)?;
let start = first.1.start;
let mut end = first.1.end;
for (tok, span) in tokens.iter().skip(start_tok) {
if matches!(tok, Token::Semicolon) {
break;
}
end = span.end;
}
let text = alter_sql.get(start..end)?.trim();
if text.is_empty() {
None
} else {
Some(text.to_string())
}
}
fn replace_span(sql: &str, span: Span, replacement: &str) -> String {
let mut out = String::with_capacity(sql.len() + replacement.len());
out.push_str(&sql[..span.start]);
out.push_str(replacement);
out.push_str(&sql[span.end..]);
out
}
fn quote_ident(name: &str) -> String {
format!("\"{}\"", name.replace('"', "\"\""))
}
fn is_safe_bare_identifier(name: &str) -> bool {
let mut chars = name.chars();
match chars.next() {
Some(c) if c == '_' || c.is_ascii_alphabetic() => {}
_ => return false,
}
chars.all(|c| c == '_' || c.is_ascii_alphanumeric())
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn append_column_matches_sqlite() {
let sql = "CREATE TABLE t (\n a INTEGER PRIMARY KEY,\n b TEXT\n)";
let out = append_column(sql, "c INTEGER").unwrap();
assert_eq!(out, "CREATE TABLE t (\n a INTEGER PRIMARY KEY,\n b TEXT\n, c INTEGER)");
}
#[test]
fn append_column_preserves_default_text() {
let sql = "CREATE TABLE t (a INTEGER)";
let out = append_column(sql, "d TEXT DEFAULT 'x'").unwrap();
assert_eq!(out, "CREATE TABLE t (a INTEGER, d TEXT DEFAULT 'x')");
}
#[test]
fn append_column_no_column_list_returns_none() {
assert!(append_column("CREATE TABLE t AS SELECT 1", "c INTEGER").is_none());
}
#[test]
fn rename_table_quotes_new_name() {
let sql = "CREATE TABLE t (\n a INTEGER PRIMARY KEY,\n b TEXT\n)";
let out = rename_table(sql, "t2").unwrap();
assert_eq!(out, "CREATE TABLE \"t2\" (\n a INTEGER PRIMARY KEY,\n b TEXT\n)");
}
#[test]
fn rename_table_handles_quoted_original() {
let sql = "CREATE TABLE \"My Table\" (x int)";
let out = rename_table(sql, "t3").unwrap();
assert_eq!(out, "CREATE TABLE \"t3\" (x int)");
}
#[test]
fn rename_table_if_not_exists() {
let sql = "CREATE TABLE IF NOT EXISTS t (x int)";
let out = rename_table(sql, "t2").unwrap();
assert_eq!(out, "CREATE TABLE IF NOT EXISTS \"t2\" (x int)");
}
#[test]
fn rename_column_in_definition_position() {
let sql = "CREATE TABLE t (\n a INTEGER PRIMARY KEY,\n b TEXT\n)";
let out = rename_column(sql, "b", "bb").unwrap();
assert_eq!(out, "CREATE TABLE t (\n a INTEGER PRIMARY KEY,\n bb TEXT\n)");
}
#[test]
fn rename_column_first_column() {
let sql = "CREATE TABLE t (a INTEGER, b TEXT)";
let out = rename_column(sql, "a", "aa").unwrap();
assert_eq!(out, "CREATE TABLE t (aa INTEGER, b TEXT)");
}
#[test]
fn rename_column_quotes_unsafe_name() {
let sql = "CREATE TABLE t (a INTEGER, b TEXT)";
let out = rename_column(sql, "b", "new col").unwrap();
assert_eq!(out, "CREATE TABLE t (a INTEGER, \"new col\" TEXT)");
}
#[test]
fn rename_column_missing_returns_none() {
let sql = "CREATE TABLE t (a INTEGER, b TEXT)";
assert!(rename_column(sql, "zzz", "qqq").is_none());
}
#[test]
fn drop_column_last_multiline_matches_sqlite() {
let sql = "CREATE TABLE t (\n a INTEGER PRIMARY KEY,\n b TEXT,\n c INTEGER\n)";
let out = drop_column(sql, "c").unwrap();
assert_eq!(out, "CREATE TABLE t (\n a INTEGER PRIMARY KEY,\n b TEXT)");
}
#[test]
fn drop_column_first_matches_sqlite() {
let sql = "CREATE TABLE t (a INTEGER, b TEXT, c INTEGER)";
let out = drop_column(sql, "a").unwrap();
assert_eq!(out, "CREATE TABLE t (b TEXT, c INTEGER)");
}
#[test]
fn drop_column_middle_matches_sqlite() {
let sql = "CREATE TABLE t (a INTEGER, b TEXT, c INTEGER)";
let out = drop_column(sql, "b").unwrap();
assert_eq!(out, "CREATE TABLE t (a INTEGER, c INTEGER)");
}
#[test]
fn drop_column_middle_multiline_matches_sqlite() {
let sql = "CREATE TABLE t (\n a INTEGER,\n b TEXT,\n c INTEGER\n)";
let out = drop_column(sql, "b").unwrap();
assert_eq!(out, "CREATE TABLE t (\n a INTEGER,\n c INTEGER\n)");
}
#[test]
fn drop_column_first_with_spaces_matches_sqlite() {
let sql = "CREATE TABLE t (a INTEGER , b TEXT , c INTEGER)";
let out = drop_column(sql, "a").unwrap();
assert_eq!(out, "CREATE TABLE t (b TEXT , c INTEGER)");
}
#[test]
fn drop_column_last_with_spaces_matches_sqlite() {
let sql = "CREATE TABLE t (a INTEGER , b TEXT , c INTEGER)";
let out = drop_column(sql, "c").unwrap();
assert_eq!(out, "CREATE TABLE t (a INTEGER , b TEXT )");
}
#[test]
fn drop_column_before_table_constraint_matches_sqlite() {
let sql = "CREATE TABLE t (a INTEGER, b TEXT, c INTEGER, UNIQUE(a))";
let out = drop_column(sql, "c").unwrap();
assert_eq!(out, "CREATE TABLE t (a INTEGER, b TEXT, UNIQUE(a))");
}
#[test]
fn drop_column_middle_with_table_constraint_matches_sqlite() {
let sql = "CREATE TABLE t (a INTEGER, b TEXT, c INTEGER, UNIQUE(c))";
let out = drop_column(sql, "b").unwrap();
assert_eq!(out, "CREATE TABLE t (a INTEGER, c INTEGER, UNIQUE(c))");
}
#[test]
fn drop_column_missing_returns_none() {
let sql = "CREATE TABLE t (a INTEGER, b TEXT)";
assert!(drop_column(sql, "zzz").is_none());
}
#[test]
fn drop_column_only_column_returns_none() {
let sql = "CREATE TABLE t (a INTEGER)";
assert!(drop_column(sql, "a").is_none());
}
#[test]
fn drop_column_quoted_name_target() {
let sql = "CREATE TABLE t (\"a\" INTEGER, b TEXT)";
let out = drop_column(sql, "a").unwrap();
assert_eq!(out, "CREATE TABLE t (b TEXT)");
}
#[test]
fn rename_table_result_reloads_through_splitter() {
let sql = "CREATE TABLE t (a INTEGER, b TEXT)";
let out = rename_table(sql, "t2").unwrap();
assert_eq!(out, "CREATE TABLE \"t2\" (a INTEGER, b TEXT)");
}
#[test]
fn extract_add_column_text_with_column_keyword() {
assert_eq!(
extract_add_column_text("ALTER TABLE t ADD COLUMN c INTEGER").as_deref(),
Some("c INTEGER")
);
}
#[test]
fn extract_add_column_text_without_column_keyword() {
assert_eq!(
extract_add_column_text("ALTER TABLE t ADD c INTEGER").as_deref(),
Some("c INTEGER")
);
}
#[test]
fn extract_add_column_text_preserves_default_and_strips_semicolon() {
assert_eq!(
extract_add_column_text("ALTER TABLE t ADD COLUMN d TEXT DEFAULT 'x';").as_deref(),
Some("d TEXT DEFAULT 'x'")
);
}
#[test]
fn append_then_extract_matches_sqlite_end_to_end() {
let create = "CREATE TABLE t (\n a INTEGER PRIMARY KEY,\n b TEXT\n, c INTEGER)";
let coldef =
extract_add_column_text("ALTER TABLE t ADD COLUMN d TEXT DEFAULT 'x'").unwrap();
let out = append_column(create, &coldef).unwrap();
assert_eq!(
out,
"CREATE TABLE t (\n a INTEGER PRIMARY KEY,\n b TEXT\n, c INTEGER, d TEXT DEFAULT 'x')"
);
}
}