This commit is contained in:
Levy A. 2025-06-03 22:56:54 -03:00
parent c2f25b6a1d
commit 49a6ddad97
8 changed files with 512 additions and 97 deletions

View file

@ -554,6 +554,21 @@ impl Display for MathFunc {
}
}
#[derive(Debug)]
pub enum AlterTableFunc {
RenameTable,
RenameColumn,
}
impl Display for AlterTableFunc {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
match self {
AlterTableFunc::RenameTable => write!(f, "limbo_rename_table"),
AlterTableFunc::RenameColumn => write!(f, "limbo_rename_column"),
}
}
}
#[derive(Debug)]
pub enum Func {
Agg(AggFunc),
@ -562,6 +577,7 @@ pub enum Func {
Vector(VectorFunc),
#[cfg(feature = "json")]
Json(JsonFunc),
AlterTable(AlterTableFunc),
External(Rc<ExternalFunc>),
}
@ -575,6 +591,7 @@ impl Display for Func {
#[cfg(feature = "json")]
Self::Json(json_func) => write!(f, "{}", json_func),
Self::External(generic_func) => write!(f, "{}", generic_func),
Self::AlterTable(alter_func) => write!(f, "{}", alter_func),
}
}
}
@ -595,6 +612,7 @@ impl Func {
#[cfg(feature = "json")]
Self::Json(json_func) => json_func.is_deterministic(),
Self::External(external_func) => external_func.is_deterministic(),
Self::AlterTable(alter_func) => true,
}
}
pub fn resolve_function(name: &str, arg_count: usize) -> Result<Self, LimboError> {

View file

@ -227,19 +227,6 @@ impl BTreeTable {
.find(|(_, column)| column.name.as_ref() == Some(&name))
}
/// Returns the column position and column for a given column name.
/// Returns None if the column name is not found.
/// E.g. if table is CREATE TABLE t(a, b, c)
/// then get_column("b") returns (1, &Column { .. })
pub fn get_column_mut(&mut self, name: &str) -> Option<(usize, &mut Column)> {
let name = normalize_ident(name);
self.columns
.iter_mut()
.enumerate()
.find(|(_, column)| column.name.as_ref() == Some(&name))
}
pub fn from_sql(sql: &str, root_page: usize) -> Result<BTreeTable> {
let mut parser = Parser::new(sql.as_bytes());
let cmd = parser.next()?;

View file

@ -6,7 +6,7 @@ use super::optimizer::Optimizable;
use super::plan::TableReferences;
#[cfg(feature = "json")]
use crate::function::JsonFunc;
use crate::function::{Func, FuncCtx, MathFuncArity, ScalarFunc, VectorFunc};
use crate::function::{AlterTableFunc, Func, FuncCtx, MathFuncArity, ScalarFunc, VectorFunc};
use crate::functions::datetime;
use crate::schema::{Affinity, Table, Type};
use crate::util::{exprs_are_equivalent, normalize_ident, parse_numeric_literal};
@ -1799,6 +1799,7 @@ pub fn translate_expr(
Ok(target_register)
}
},
Func::AlterTable(_) => unreachable!(),
}
}
ast::Expr::FunctionCallStar { .. } => todo!(),

View file

@ -31,6 +31,7 @@ pub(crate) mod update;
mod values;
use crate::fast_lock::SpinLock;
use crate::function::{AlterTableFunc, Func};
use crate::schema::{Column, Schema};
use crate::storage::pager::Pager;
use crate::storage::sqlite3_ondisk::DatabaseHeader;
@ -40,6 +41,7 @@ use crate::vdbe::builder::{ProgramBuilder, ProgramBuilderOpts, QueryMode};
use crate::vdbe::insn::{Insn, RegisterOrLiteral};
use crate::vdbe::Program;
use crate::{bail_parse_error, Connection, LimboError, Result, SymbolTable};
use emitter::TransactionMode;
use fallible_iterator::FallibleIterator as _;
use index::{translate_create_index, translate_drop_index};
use insert::translate_insert;
@ -109,7 +111,7 @@ pub fn translate_inner(
stmt: ast::Stmt,
syms: &SymbolTable,
query_mode: QueryMode,
program: ProgramBuilder,
mut program: ProgramBuilder,
) -> Result<ProgramBuilder> {
let program = match stmt {
ast::Stmt::AlterTable(a) => {
@ -177,7 +179,9 @@ pub fn translate_inner(
);
let mut parser = Parser::new(stmt.as_bytes());
let Some(ast::Cmd::Stmt(ast::Stmt::Update(mut update))) = parser.next().unwrap() else {
let Some(ast::Cmd::Stmt(ast::Stmt::Update(mut update))) =
parser.next().unwrap()
else {
unreachable!();
};
@ -313,43 +317,88 @@ pub fn translate_inner(
)?
}
ast::AlterTableBody::RenameColumn { old, new } => {
let ast::Name(old) = old;
let ast::Name(new) = new;
let ast::Name(rename_from) = old;
let ast::Name(rename_to) = new;
let Some((_, column)) = btree.get_column_mut(&old) else {
return Err(LimboError::ParseError(format!("no such column: \"{old}\"")));
if btree.get_column(&rename_from).is_none() {
return Err(LimboError::ParseError(format!("no such column: \"{rename_from}\"")));
};
column.name = Some(new);
let sqlite_schema = schema
.get_btree_table(SQLITE_TABLEID)
.expect("sqlite_schema should be on schema");
let sql = btree.to_sql();
let stmt = format!(
r#"
UPDATE {SQLITE_TABLEID}
SET sql = '{sql}'
WHERE name = '{table_name}' COLLATE NOCASE AND type = 'table'
"#,
let cursor_id = program.alloc_cursor_id(
crate::vdbe::builder::CursorType::BTreeTable(sqlite_schema.clone()),
);
let mut parser = Parser::new(stmt.as_bytes());
let Some(ast::Cmd::Stmt(ast::Stmt::Update(mut update))) = parser.next().unwrap() else {
unreachable!();
};
program.emit_insn(Insn::OpenWrite {
cursor_id,
root_page: RegisterOrLiteral::Literal(sqlite_schema.root_page),
name: sqlite_schema.name.clone(),
});
translate_update_with_after(
QueryMode::Normal,
schema,
&mut update,
syms,
program,
|program| {
program.emit_insn(Insn::ParseSchema {
db: usize::MAX, // TODO: This value is unused, change when we do something with it
where_clause: None,
});
},
)?
program.cursor_loop(cursor_id, |program| {
let rowid = program.alloc_register();
program.emit_insn(Insn::RowId {
cursor_id,
dest: rowid,
});
let first_column = program.alloc_registers(5);
for i in 0..5 {
program.emit_column(cursor_id, i, first_column + i);
}
program.emit_string8_new_reg(table_name.clone());
program.mark_last_insn_constant();
program.emit_string8_new_reg(rename_from.clone());
program.mark_last_insn_constant();
program.emit_string8_new_reg(rename_to.clone());
program.mark_last_insn_constant();
let out = program.alloc_registers(5);
program.emit_insn(Insn::Function {
constant_mask: 0,
start_reg: first_column,
dest: out,
func: crate::function::FuncCtx {
func: Func::AlterTable(AlterTableFunc::RenameColumn),
arg_count: 8,
},
});
let record = program.alloc_register();
program.emit_insn(Insn::MakeRecord {
start_reg: out,
count: 5,
dest_reg: record,
index_name: None,
});
program.emit_insn(Insn::Insert {
cursor: cursor_id,
key_reg: rowid,
record_reg: record,
flag: 0,
table_name: table_name.clone(),
});
});
program.emit_insn(Insn::ParseSchema {
db: usize::MAX, // TODO: This value is unused, change when we do something with it
where_clause: None,
});
program.epilogue(TransactionMode::Write);
program
}
ast::AlterTableBody::RenameTo(new_name) => {
let ast::Name(new_name) = new_name;
@ -360,39 +409,78 @@ pub fn translate_inner(
)));
};
btree.name = new_name;
let sqlite_schema = schema
.get_btree_table(SQLITE_TABLEID)
.expect("sqlite_schema should be on schema");
let sql = btree.to_sql();
let stmt = format!(
r#"
UPDATE {SQLITE_TABLEID}
SET name = CASE WHEN type = 'table' THEN '{new_name}' ELSE name END
, tbl_name = '{new_name}'
, sql = CASE WHEN type = 'table' THEN '{sql}' ELSE sql END
WHERE tbl_name = '{table_name}' COLLATE NOCASE
"#,
new_name = &btree.name,
let cursor_id = program.alloc_cursor_id(
crate::vdbe::builder::CursorType::BTreeTable(sqlite_schema.clone()),
);
let mut parser = Parser::new(stmt.as_bytes());
let Some(ast::Cmd::Stmt(ast::Stmt::Update(mut update))) = parser.next().unwrap() else {
unreachable!();
};
program.emit_insn(Insn::OpenWrite {
cursor_id,
root_page: RegisterOrLiteral::Literal(sqlite_schema.root_page),
name: sqlite_schema.name.clone(),
});
translate_update_with_after(
QueryMode::Normal,
schema,
&mut update,
syms,
program,
|program| {
program.emit_insn(Insn::ParseSchema {
db: usize::MAX, // TODO: This value is unused, change when we do something with it
where_clause: None,
});
},
)?
program.cursor_loop(cursor_id, |program| {
let rowid = program.alloc_register();
program.emit_insn(Insn::RowId {
cursor_id,
dest: rowid,
});
let first_column = program.alloc_registers(5);
for i in 0..5 {
program.emit_column(cursor_id, i, first_column + i);
}
program.emit_string8_new_reg(table_name.clone());
program.mark_last_insn_constant();
program.emit_string8_new_reg(new_name.clone());
program.mark_last_insn_constant();
let out = program.alloc_registers(5);
program.emit_insn(Insn::Function {
constant_mask: 0,
start_reg: first_column,
dest: out,
func: crate::function::FuncCtx {
func: Func::AlterTable(AlterTableFunc::RenameTable),
arg_count: 7,
},
});
let record = program.alloc_register();
program.emit_insn(Insn::MakeRecord {
start_reg: out,
count: 5,
dest_reg: record,
index_name: None,
});
program.emit_insn(Insn::Insert {
cursor: cursor_id,
key_reg: rowid,
record_reg: record,
flag: 0,
table_name: table_name.clone(),
});
});
program.emit_insn(Insn::ParseSchema {
db: usize::MAX, // TODO: This value is unused, change when we do something with it
where_clause: None,
});
program.epilogue(TransactionMode::Write);
program
}
}
}
@ -481,13 +569,9 @@ pub fn translate_inner(
)?
.program
}
ast::Stmt::Update(mut update) => translate_update(
query_mode,
schema,
&mut update,
syms,
program,
)?,
ast::Stmt::Update(mut update) => {
translate_update(query_mode, schema, &mut update, syms, program)?
}
ast::Stmt::Vacuum(_, _) => bail_parse_error!("VACUUM not supported yet"),
ast::Stmt::Insert(insert) => {
let Insert {

View file

@ -93,6 +93,12 @@ impl Text {
}
}
impl AsRef<str> for Text {
fn as_ref(&self) -> &str {
self.as_str()
}
}
impl From<String> for Text {
fn from(value: String) -> Self {
Text {

View file

@ -1,4 +1,5 @@
#![allow(unused_variables)]
use crate::function::AlterTableFunc;
use crate::numeric::{NullableInteger, Numeric};
use crate::schema::Schema;
use crate::storage::database::FileMemoryStorage;
@ -7,6 +8,7 @@ use crate::storage::pager::CreateBTreeFlags;
use crate::storage::wal::DummyWAL;
use crate::translate::collate::CollationSeq;
use crate::types::{ImmutableRecord, Text};
use crate::util::normalize_ident;
use crate::{
error::{LimboError, SQLITE_CONSTRAINT, SQLITE_CONSTRAINT_PRIMARYKEY},
ext::ExtValue,
@ -53,6 +55,10 @@ use super::{
insn::{Cookie, RegisterOrLiteral},
CommitState,
};
use fallible_iterator::FallibleIterator;
use limbo_sqlite3_parser::ast;
use limbo_sqlite3_parser::ast::fmt::ToTokens;
use limbo_sqlite3_parser::lexer::sql::Parser;
use parking_lot::RwLock;
use rand::thread_rng;
@ -3769,6 +3775,275 @@ pub fn op_function(
),
},
},
crate::function::Func::AlterTable(alter_func) => {
let r#type = &state.registers[*start_reg + 0].get_owned_value().clone();
let Value::Text(name) = &state.registers[*start_reg + 1].get_owned_value() else {
panic!("sqlite_schema.name should be TEXT")
};
let name = name.to_string();
let Value::Text(tbl_name) = &state.registers[*start_reg + 2].get_owned_value() else {
panic!("sqlite_schema.tbl_name should be TEXT")
};
let tbl_name = tbl_name.to_string();
let Value::Integer(root_page) =
&state.registers[*start_reg + 3].get_owned_value().clone()
else {
panic!("sqlite_schema.root_page should be INTEGER")
};
let sql = &state.registers[*start_reg + 4].get_owned_value().clone();
let (new_name, new_tbl_name, new_sql) = match alter_func {
AlterTableFunc::RenameTable => {
let rename_from = {
match &state.registers[*start_reg + 5].get_owned_value() {
Value::Text(rename_from) => normalize_ident(rename_from.as_str()),
_ => panic!("rename_from parameter should be TEXT"),
}
};
let rename_to = {
match &state.registers[*start_reg + 6].get_owned_value() {
Value::Text(rename_to) => normalize_ident(rename_to.as_str()),
_ => panic!("rename_to parameter should be TEXT"),
}
};
let new_name = if let Some(column) =
&name.strip_prefix(&format!("sqlite_autoindex_{rename_from}_"))
{
format!("sqlite_autoindex_{rename_to}_{column}")
} else if name == rename_from {
rename_to.clone()
} else {
name
};
let new_tbl_name = if tbl_name == rename_from {
rename_to.clone()
} else {
tbl_name
};
let new_sql = 'sql: {
let Value::Text(sql) = sql else {
break 'sql None;
};
let mut parser = Parser::new(sql.as_str().as_bytes());
let ast::Cmd::Stmt(stmt) = parser.next().unwrap().unwrap() else {
todo!()
};
match stmt {
ast::Stmt::CreateIndex {
unique,
if_not_exists,
idx_name,
tbl_name,
columns,
where_clause,
} => {
let table_name = normalize_ident(&tbl_name.0);
if rename_from != table_name {
break 'sql None;
}
Some(
ast::Stmt::CreateIndex {
unique,
if_not_exists,
idx_name,
tbl_name: ast::Name(rename_to),
columns,
where_clause,
}
.format()
.unwrap(),
)
}
ast::Stmt::CreateTable {
temporary,
if_not_exists,
tbl_name,
body,
} => {
let table_name = normalize_ident(&tbl_name.name.0);
if rename_from != table_name {
break 'sql None;
}
Some(
ast::Stmt::CreateTable {
temporary,
if_not_exists,
tbl_name: ast::QualifiedName {
db_name: None,
name: ast::Name(rename_to),
alias: None,
},
body,
}
.format()
.unwrap(),
)
}
_ => todo!(),
}
};
(new_name, new_tbl_name, new_sql)
}
AlterTableFunc::RenameColumn => {
let table = {
match &state.registers[*start_reg + 5].get_owned_value() {
Value::Text(rename_to) => normalize_ident(rename_to.as_str()),
_ => panic!("table parameter should be TEXT"),
}
};
let rename_from = {
match &state.registers[*start_reg + 6].get_owned_value() {
Value::Text(rename_from) => normalize_ident(rename_from.as_str()),
_ => panic!("rename_from parameter should be TEXT"),
}
};
let rename_to = {
match &state.registers[*start_reg + 7].get_owned_value() {
Value::Text(rename_to) => normalize_ident(rename_to.as_str()),
_ => panic!("rename_to parameter should be TEXT"),
}
};
let new_sql = 'sql: {
if table != tbl_name {
break 'sql None;
}
let Value::Text(sql) = sql else {
break 'sql None;
};
let mut parser = Parser::new(sql.as_str().as_bytes());
let ast::Cmd::Stmt(stmt) = parser.next().unwrap().unwrap() else {
todo!()
};
match stmt {
ast::Stmt::CreateIndex {
unique,
if_not_exists,
idx_name,
tbl_name,
mut columns,
where_clause,
} => {
if table != normalize_ident(&tbl_name.0) {
break 'sql None;
}
for column in &mut columns {
match &mut column.expr {
ast::Expr::Id(ast::Id(id))
if normalize_ident(&id) == rename_from =>
{
*id = rename_to.clone();
}
_ => {}
}
}
Some(
ast::Stmt::CreateIndex {
unique,
if_not_exists,
idx_name,
tbl_name,
columns,
where_clause,
}
.format()
.unwrap(),
)
}
ast::Stmt::CreateTable {
temporary,
if_not_exists,
tbl_name,
body,
} => {
if table != normalize_ident(&tbl_name.name.0) {
break 'sql None;
}
let ast::CreateTableBody::ColumnsAndConstraints {
mut columns,
constraints,
options,
} = *body
else {
todo!()
};
let column_index = columns
.get_index_of(&ast::Name(rename_from))
.expect("column being renamed should be present");
let mut column_definition =
columns.get_index(column_index).unwrap().1.clone();
column_definition.col_name = ast::Name(rename_to.clone());
assert!(columns
.insert(ast::Name(rename_to), column_definition.clone())
.is_none());
// Swaps indexes with the last one and pops the end, effectively
// replacing the entry.
columns.swap_remove_index(column_index).unwrap();
Some(
ast::Stmt::CreateTable {
temporary,
if_not_exists,
tbl_name,
body: Box::new(
ast::CreateTableBody::ColumnsAndConstraints {
columns,
constraints,
options,
},
),
}
.format()
.unwrap(),
)
}
_ => todo!(),
}
};
(name, tbl_name, new_sql)
}
};
state.registers[*dest + 0] = Register::Value(r#type.clone());
state.registers[*dest + 1] = Register::Value(Value::Text(Text::from(new_name)));
state.registers[*dest + 2] = Register::Value(Value::Text(Text::from(new_tbl_name)));
state.registers[*dest + 3] = Register::Value(Value::Integer(*root_page));
if let Some(new_sql) = new_sql {
state.registers[*dest + 4] = Register::Value(Value::Text(Text::from(new_sql)));
} else {
state.registers[*dest + 4] = Register::Value(sql.clone());
}
}
crate::function::Func::Agg(_) => {
unreachable!("Aggregate functions should not be handled here")
}

View file

@ -4,16 +4,20 @@ set testdir [file dirname $argv0]
source $testdir/tester.tcl
do_execsql_test_on_specific_db {:memory:} alter-table-rename-table {
CREATE TABLE t1(x INTEGER PRIMARY KEY);
CREATE TABLE t1(x INTEGER PRIMARY KEY, u UNIQUE);
ALTER TABLE t1 RENAME TO t2;
SELECT tbl_name FROM sqlite_schema;
SELECT name FROM sqlite_schema WHERE type = 'table';
} { "t2" }
do_execsql_test_on_specific_db {:memory:} alter-table-rename-column {
CREATE TABLE t(a);
CREATE INDEX i ON t(a);
ALTER TABLE t RENAME a TO b;
SELECT sql FROM sqlite_schema;
} { "CREATE TABLE t(b)" }
} {
"CREATE TABLE t(b)"
"CREATE INDEX i ON t(b)"
}
do_execsql_test_on_specific_db {:memory:} alter-table-add-column {
CREATE TABLE t(a);

View file

@ -46,6 +46,45 @@ impl TokenStream for FmtTokenStream<'_, '_> {
}
}
struct WriteTokenStream<'a, T: fmt::Write> {
write: &'a mut T,
spaced: bool,
}
impl<T: fmt::Write> TokenStream for WriteTokenStream<'_, T> {
type Error = fmt::Error;
fn append(&mut self, ty: TokenType, value: Option<&str>) -> fmt::Result {
if !self.spaced {
match ty {
TK_COMMA | TK_SEMI | TK_RP | TK_DOT => {}
_ => {
self.write.write_char(' ')?;
self.spaced = true;
}
};
}
if ty == TK_BLOB {
self.write.write_char('X')?;
self.write.write_char('\'')?;
if let Some(str) = value {
self.write.write_str(str)?;
}
return self.write.write_char('\'');
} else if let Some(str) = ty.as_str() {
self.write.write_str(str)?;
self.spaced = ty == TK_LP || ty == TK_DOT; // str should not be whitespace
}
if let Some(str) = value {
// trick for pretty-print
self.spaced = str.bytes().all(|b| b.is_ascii_whitespace());
self.write.write_str(str)
} else {
Ok(())
}
}
}
/// Stream of token
pub trait TokenStream {
/// Potential error raised
@ -63,6 +102,19 @@ pub trait ToTokens {
let mut s = FmtTokenStream { f, spaced: true };
self.to_tokens(&mut s)
}
// Format AST node to string
fn format(&self) -> Result<String, fmt::Error> {
let mut s = String::new();
let mut w = WriteTokenStream {
write: &mut s,
spaced: true,
};
self.to_tokens(&mut w)?;
Ok(s)
}
}
impl<T: ?Sized + ToTokens> ToTokens for &T {
@ -77,18 +129,6 @@ impl ToTokens for String {
}
}
/* FIXME: does not work, find why
impl Display for dyn ToTokens {
fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
let mut s = FmtTokenStream { f, spaced: true };
match self.to_tokens(&mut s) {
Err(_) => Err(fmt::Error),
Ok(()) => Ok(()),
}
}
}
*/
impl ToTokens for Cmd {
fn to_tokens<S: TokenStream>(&self, s: &mut S) -> Result<(), S::Error> {
match self {