use std::slice;
use arrow_array::RecordBatch;
use reifydb_codec::row::bytes::EncodedBytes;
use reifydb_core::{
interface::catalog::object::ObjectId,
internal_error,
key::{any::TaggedKey, partition::PartitionKey, row::PartitionedRowKey},
partition::partition_of,
value::batch::single_row,
};
use reifydb_rql::nodes::{AlterTableAction, AlterTableNode};
use reifydb_transaction::{
multi::RangeScope,
transaction::{Transaction, admin::AdminTransaction},
};
use reifydb_value::value::{Value, row_number::RowNumber, value_type::ValueType};
use crate::{
Result,
transaction::operation::{dictionary::DictionaryOperations, table::TableOperations},
vm::{
instruction::dml::primary_key::{PrimaryKeyIndex, get_primary_key},
services::Services,
},
};
pub(crate) fn execute_alter_table(
services: &Services,
txn: &mut AdminTransaction,
plan: AlterTableNode,
) -> Result<RecordBatch> {
let namespace_id = plan.namespace.def().id();
let namespace_name = plan.namespace.name().to_string();
let table_name = plan.table.text().to_string();
let table =
services.catalog.get_table_by_name(&mut Transaction::Admin(txn), namespace_id, plan.table.clone())?;
let (operation, details) = match plan.action {
AlterTableAction::AddColumn {
column,
} => {
let col_name = column.name.text().to_string();
services.catalog.add_table_column(txn, table.id, column, &namespace_name)?;
("ADD COLUMN", Value::Utf8(col_name))
}
AlterTableAction::DropColumn {
column,
} => {
let col_name = column.text().to_string();
services.catalog.drop_table_column(txn, table.id, column.text(), &namespace_name)?;
("DROP COLUMN", Value::Utf8(col_name))
}
AlterTableAction::RenameColumn {
old_name,
new_name,
} => {
let detail = format!("{} -> {}", old_name.text(), new_name.text());
services.catalog.rename_table_column(
txn,
table.id,
old_name.text(),
new_name.text(),
&namespace_name,
)?;
("RENAME COLUMN", Value::Utf8(detail))
}
AlterTableAction::DropPartition {
values,
remove_registry,
} => {
if table.partition_by.is_empty() {
return Err(internal_error!("table {} is not partitioned", table_name));
}
let mut part_values = Vec::with_capacity(table.partition_by.len());
let mut interned = true;
for col_name in &table.partition_by {
let Some((_, text)) = values.iter().find(|(c, _)| c == col_name) else {
return Err(internal_error!(
"DROP PARTITION must bind partition column {}",
col_name
));
};
let is_utf8 = table
.columns
.iter()
.find(|c| &c.name == col_name)
.map(|c| c.constraint.get_type()) == Some(ValueType::Utf8);
if !is_utf8 {
return Err(internal_error!(
"DROP PARTITION currently supports only Utf8 partition columns (column {})",
col_name
));
}
let value = Value::Utf8(text.clone());
match table.columns.iter().find(|c| &c.name == col_name).and_then(|c| c.dictionary_id) {
Some(dict_id) => {
let dictionary = services
.catalog
.find_dictionary(&mut Transaction::Admin(txn), dict_id)?
.ok_or_else(|| {
internal_error!(
"Dictionary {:?} not found for column {}",
dict_id,
col_name
)
})?;
match txn
.find_in_dictionary(&dictionary, slice::from_ref(&value))?
.remove(0)
{
Some(entry_id) => part_values.push(entry_id.to_value()),
None => interned = false,
}
}
None => part_values.push(value),
}
}
let partition =
interned.then(|| partition_of(&table.columns, &table.partition_by, &part_values));
let object = ObjectId::Table(table.id);
let mut ids: Vec<RowNumber> = Vec::new();
let mut rows: Vec<EncodedBytes> = Vec::new();
let mut last_key: Option<TaggedKey> = None;
while let Some(partition) = partition {
let batch: Vec<_> = txn
.range(
PartitionedRowKey::partition_scan_range(
table.id,
partition,
last_key.as_ref(),
),
RangeScope::All,
1024,
)?
.collect::<Result<Vec<_>>>()?;
if batch.is_empty() {
break;
}
let n = batch.len();
for entry in batch {
if let TaggedKey::PartitionedRow(pk) = &entry.key {
ids.push(pk.row);
rows.push(entry.bytes.clone());
}
last_key = Some(entry.key.clone());
}
if n < 1024 {
break;
}
}
let dropped = ids.len() as u64;
if let Some(partition) = partition
&& !ids.is_empty()
{
let mut tx = Transaction::Admin(txn);
if let Some(pk) = get_primary_key(&services.catalog, &mut tx, &table)? {
let mut index = PrimaryKeyIndex::new(&pk);
for row in &rows {
index.remove(&services.catalog, &mut tx, &table, row)?;
}
index.write(&mut tx)?;
}
let partitions = vec![partition; ids.len()];
txn.remove_from_table(&services.catalog, &table, &ids, &partitions, &rows)?;
}
if remove_registry {
if let Some(partition) = partition {
txn.remove(vec![(PartitionKey::new(object, partition).into(), None)])?;
}
("DROP PARTITION", Value::Uint8(dropped))
} else {
("TRUNCATE PARTITION", Value::Uint8(dropped))
}
}
};
single_row([
("operation", Value::Utf8(operation.to_string())),
("namespace", Value::Utf8(namespace_name)),
("table", Value::Utf8(table_name)),
("details", details),
])
}
#[cfg(test)]
mod tests {
use reifydb_test_harness::engine::create_test_admin_transaction;
use reifydb_transaction::transaction::admin::AdminTransaction;
use reifydb_value::{
params::Params,
value::{Value, column_view::ColumnView, frame::frame::Frame},
};
use crate::vm::{Admin, executor::Executor};
fn setup() -> (Executor, AdminTransaction) {
let instance = Executor::testing();
let mut txn = create_test_admin_transaction();
let r = instance.admin(
&mut txn,
Admin {
rql: "CREATE NAMESPACE app",
params: Params::default(),
},
);
if let Some(e) = r.error {
panic!("{e:?}");
}
let r = instance.admin(
&mut txn,
Admin {
rql: "CREATE TABLE app::users { id: Int4, name: Utf8 }",
params: Params::default(),
},
);
if let Some(e) = r.error {
panic!("{e:?}");
}
(instance, txn)
}
#[test]
fn test_alter_table_add_column() {
let (instance, mut txn) = setup();
let r = instance.admin(
&mut txn,
Admin {
rql: "ALTER TABLE app::users ADD COLUMN email: Utf8",
params: Params::default(),
},
);
if let Some(e) = r.error {
panic!("{e:?}");
}
let frame = &r[0];
assert_eq!(value_at(frame, 0), Value::Utf8("ADD COLUMN".to_string()));
assert_eq!(value_at(frame, 1), Value::Utf8("app".to_string()));
assert_eq!(value_at(frame, 2), Value::Utf8("users".to_string()));
assert_eq!(value_at(frame, 3), Value::Utf8("email".to_string()));
}
#[test]
fn test_alter_table_drop_column() {
let (instance, mut txn) = setup();
let r = instance.admin(
&mut txn,
Admin {
rql: "ALTER TABLE app::users DROP COLUMN name",
params: Params::default(),
},
);
if let Some(e) = r.error {
panic!("{e:?}");
}
let frame = &r[0];
assert_eq!(value_at(frame, 0), Value::Utf8("DROP COLUMN".to_string()));
assert_eq!(value_at(frame, 1), Value::Utf8("app".to_string()));
assert_eq!(value_at(frame, 2), Value::Utf8("users".to_string()));
assert_eq!(value_at(frame, 3), Value::Utf8("name".to_string()));
}
#[test]
fn test_alter_table_rename_column() {
let (instance, mut txn) = setup();
let r = instance.admin(
&mut txn,
Admin {
rql: "ALTER TABLE app::users RENAME COLUMN name TO full_name",
params: Params::default(),
},
);
if let Some(e) = r.error {
panic!("{e:?}");
}
let frame = &r[0];
assert_eq!(value_at(frame, 0), Value::Utf8("RENAME COLUMN".to_string()));
assert_eq!(value_at(frame, 1), Value::Utf8("app".to_string()));
assert_eq!(value_at(frame, 2), Value::Utf8("users".to_string()));
assert_eq!(value_at(frame, 3), Value::Utf8("name -> full_name".to_string()));
}
#[test]
fn test_alter_table_drop_nonexistent_column() {
let (instance, mut txn) = setup();
let r = instance.admin(
&mut txn,
Admin {
rql: "ALTER TABLE app::users DROP COLUMN nonexistent",
params: Params::default(),
},
);
assert!(r.is_err());
assert_eq!(r.error.unwrap().diagnostic().code, "CA_039");
}
#[test]
fn test_alter_table_rename_nonexistent_column() {
let (instance, mut txn) = setup();
let r = instance.admin(
&mut txn,
Admin {
rql: "ALTER TABLE app::users RENAME COLUMN nonexistent TO new_name",
params: Params::default(),
},
);
assert!(r.is_err());
assert_eq!(r.error.unwrap().diagnostic().code, "CA_039");
}
fn value_at(frame: &Frame, column: usize) -> Value {
ColumnView::try_from((frame.batch.column(column), frame.batch.schema_ref().field(column)))
.unwrap()
.get_value(0)
}
}