use super::{
table_next_id_metadata_key, table_not_found, Engine, RelationIdentity, StorageBackendError,
StorageBackendResult,
};
use crate::VectorIndexSpec;
impl Engine {
pub fn register_column(
&self,
table: &str,
column: uqa_sql::ast::ColumnDef,
) -> StorageBackendResult<()> {
self.try_register_column(table, column)
}
pub(crate) fn try_register_column(
&self,
table: &str,
column: uqa_sql::ast::ColumnDef,
) -> StorageBackendResult<()> {
self.with_implicit_storage_transaction(|engine| {
engine.try_register_column_inner(table, column, None)
})
}
pub(super) fn try_register_column_inner(
&self,
table: &str,
column: uqa_sql::ast::ColumnDef,
check_columns: Option<&[uqa_sql::ast::ColumnDef]>,
) -> StorageBackendResult<()> {
uqa_execution::schema::publication::register_column(
&self.schema_publication_context(),
table,
column,
check_columns,
&[],
)
}
pub fn drop_column(&self, table: &str, column: &str) -> StorageBackendResult<bool> {
self.with_implicit_storage_transaction(|engine| {
let mut rewritten = Vec::new();
if let Some(canonical) =
engine.resolve_table_ddl_target(table, "ALTER TABLE DROP COLUMN")?
{
if engine.table_has_column_in_execution(&canonical, column)? {
engine
.drop_column_routine_dependents(&canonical, column, false)
.map_err(|error| StorageBackendError::Other(error.to_string()))?;
rewritten = engine
.prepare_routine_column_alias_drop(
std::collections::BTreeSet::from([(canonical, column.to_string())]),
&[],
)
.map_err(|error| StorageBackendError::Other(error.to_string()))?;
}
}
let dropped = engine.try_drop_column_inner(table, column)?;
if dropped {
engine
.publish_stored_routine_body_rewrites(rewritten)
.map_err(|error| StorageBackendError::Other(error.to_string()))?;
engine
.refresh_stored_merge_target_plans()
.map_err(|error| StorageBackendError::Other(error.to_string()))?;
}
Ok(dropped)
})
}
pub(crate) fn try_drop_column(&self, table: &str, column: &str) -> StorageBackendResult<bool> {
self.with_implicit_storage_transaction(|engine| engine.try_drop_column_inner(table, column))
}
pub(crate) fn try_drop_column_cascade(
&self,
table: &str,
column: &str,
) -> StorageBackendResult<bool> {
self.with_implicit_storage_transaction(|engine| {
engine.try_drop_column_inner_with_sequence_cascade(table, column, true)
})
}
pub(crate) fn try_drop_column_inner(
&self,
table: &str,
column: &str,
) -> StorageBackendResult<bool> {
self.try_drop_column_inner_with_sequence_cascade(table, column, false)
}
fn try_drop_column_inner_with_sequence_cascade(
&self,
table: &str,
column: &str,
cascade: bool,
) -> StorageBackendResult<bool> {
uqa_execution::schema::publication::removal::drop_column(
&self.column_drop_publication_context(),
table,
column,
cascade,
)
}
pub(crate) fn try_drop_vector_indexes_for_column(
&self,
table: &str,
column: &str,
) -> StorageBackendResult<bool> {
self.with_implicit_storage_transaction(|engine| {
engine.try_drop_vector_indexes_for_column_inner(table, column)
})
}
pub(super) fn try_drop_vector_indexes_for_column_inner(
&self,
table: &str,
column: &str,
) -> StorageBackendResult<bool> {
let Some(table_name) = self.resolve_table_ddl_target(table, "ALTER TABLE ALTER COLUMN")?
else {
return Ok(false);
};
uqa_execution::catalog::index::vectors::remove_for_column_conversion(
&self.index_registry_context(),
&table_name,
column,
|| {
let state = self
.try_table(&table_name)?
.ok_or_else(|| table_not_found(&table_name))?;
let dimensions = state
.vector_indexes
.read()
.get(column)
.map(uqa_storage::VectorIndex::dimensions);
if let Some(dimensions) = dimensions {
self.prepare_vector_column_rewrite(&table_name, column, dimensions)?;
state.vector_indexes.write().live_mut()?.remove(column);
}
self.try_save_table_schema(&table_name, &state)
},
)?;
Ok(true)
}
pub(crate) fn prepare_vector_column_rewrite(
&self,
table: &str,
column: &str,
dimensions: u32,
) -> StorageBackendResult<()> {
let table_name = self
.try_resolve_table_name(table)?
.ok_or_else(|| table_not_found(table))?;
let state = self
.try_table(&table_name)?
.ok_or_else(|| table_not_found(&table_name))?;
uqa_execution::schema::indexes::diskann::retire_column(
&self.index_registry_context(),
&table_name,
column,
)?;
let canonical = self.build_vector_index_with_mode(
&table_name,
column,
dimensions,
VectorIndexSpec::BruteForce,
uqa_storage::VectorIndexOpenMode::Create,
)?;
uqa_execution::catalog::index::vectors::prepare_column_rewrite(
&mut state.vector_indexes.write(),
column,
canonical,
)?;
Ok(())
}
pub(crate) fn try_rebuild_vector_index_for_column(
&self,
table: &str,
column: &str,
dimensions: u32,
) -> StorageBackendResult<bool> {
self.with_implicit_storage_transaction(|engine| {
engine.try_rebuild_vector_index_for_column_inner(table, column, dimensions)
})
}
pub(super) fn try_rebuild_vector_index_for_column_inner(
&self,
table: &str,
column: &str,
dimensions: u32,
) -> StorageBackendResult<bool> {
let table_name = self
.try_resolve_table_name(table)?
.ok_or_else(|| table_not_found(table))?;
let spec = self
.vector_index_spec_for_column(&table_name, column, dimensions)?
.unwrap_or(VectorIndexSpec::BruteForce);
let t = self
.try_table(&table_name)?
.ok_or_else(|| table_not_found(&table_name))?;
self.try_save_table_schema(&table_name, &t)?;
if uqa_execution::schema::indexes::diskann::create_column(
&self.index_registry_context(),
&table_name,
column,
)? {
return Ok(true);
}
let rebuilt = self.rebuild_vector_field_with_spec(&table_name, column, dimensions, spec)?;
if !rebuilt {
return Err(StorageBackendError::Other(format!(
"failed to rebuild vector index for `{table_name}`.`{column}`"
)));
}
Ok(true)
}
pub fn rename_column(&self, table: &str, from: &str, to: &str) -> StorageBackendResult<bool> {
self.try_rename_column(table, from, to)
}
pub(crate) fn try_rename_column(
&self,
table: &str,
from: &str,
to: &str,
) -> StorageBackendResult<bool> {
self.with_implicit_storage_transaction(|engine| {
engine.try_rename_column_inner(table, from, to)
})
}
fn rename_column_acl(table: &super::TableState, from: &str, to: &str) {
if from == to {
return;
}
let mut security = table.security.write();
security.rename_column_acl(from, to);
}
fn rename_column_analyzer_assignments(&self, table_name: &str, from: &str, to: &str) {
let mut analyzers = self.durable.table_field_analyzers.write();
let mut moved = Vec::new();
analyzers.retain(|(table, field), value| {
if table == table_name && field == from {
moved.push(((table_name.to_string(), to.to_string()), value.clone()));
false
} else {
true
}
});
analyzers.extend(moved);
}
fn rename_document_fields(
&self,
table_name: &str,
table: &super::TableState,
from: &str,
to: &str,
) -> StorageBackendResult<()> {
let ids = table.document_store.read().doc_ids()?;
for doc_id in ids {
let Some(mut doc) = table.document_store.read().get(doc_id)? else {
continue;
};
if let Some(value) = doc.remove(from) {
doc.insert(to.to_string(), value);
self.rewrite_document_for_schema_change(table_name, doc_id, doc)
.map_err(|err| StorageBackendError::Other(err.to_string()))?;
}
}
Ok(())
}
fn bind_renamed_vector_field(
&self,
table: &str,
field: &str,
dimensions: u32,
restore: bool,
) -> StorageBackendResult<()> {
let Some(spec) = self.vector_index_spec_for_column(table, field, dimensions)? else {
return Ok(());
};
let bound = if restore {
self.restore_vector_field_index(table, field, dimensions, spec)?
} else {
self.rebuild_vector_field_with_spec(table, field, dimensions, spec)?
};
if !bound {
return Err(StorageBackendError::Other(format!(
"failed to bind renamed vector index for `{table}`.`{field}`"
)));
}
Ok(())
}
pub(super) fn try_rename_column_inner(
&self,
table: &str,
from: &str,
to: &str,
) -> StorageBackendResult<bool> {
let Some(table_name) = self.resolve_table_ddl_target(table, "ALTER TABLE RENAME COLUMN")?
else {
return Ok(false);
};
let Some(t) = self.try_table(table)? else {
return Ok(false);
};
{
let columns = t.columns.read();
if !columns.iter().any(|candidate| candidate.name == from) {
return Ok(false);
}
if from != to && columns.iter().any(|candidate| candidate.name == to) {
return Ok(false);
}
}
let analyzer_binding = if from != to && t.fts_fields().iter().any(|field| field == from) {
Some(
self.current_field_analyzer_binding(&table_name, &t, from)
.map_err(StorageBackendError::Other)?,
)
} else {
None
};
self.rewrite_column_rename_dependencies(&table_name, from, to)?;
Self::value_indexes_clear(&t);
{
let mut cols = t.columns.write();
for c in cols.iter_mut() {
if c.name == from {
c.name = to.to_string();
}
}
}
Self::rename_column_acl(&t, from, to);
for constraint in t.key_constraints.write().iter_mut() {
for column in &mut constraint.columns {
if column == from {
*column = to.to_string();
}
}
}
{
let mut fts = t.fts_fields.write();
for f in fts.iter_mut() {
if f == from {
*f = to.to_string();
}
}
}
if let Some(binding) = &analyzer_binding {
binding.install(to, t.inverted_index.write().as_mut())?;
}
self.rename_column_analyzer_assignments(&table_name, from, to);
let vector_rename = uqa_execution::catalog::index::vectors::detach_for_column_rename(
&mut t.vector_indexes.write(),
from,
)?;
let vector_dimensions = vector_rename.as_ref().map(|rename| rename.dimensions);
let retained_vector = vector_rename.and_then(|rename| rename.retained);
let restore_vector = retained_vector.is_some();
self.rename_document_fields(&table_name, &t, from, to)?;
if analyzer_binding.is_some() {
t.inverted_index
.write()
.remove_field_analyzers(from)
.map_err(StorageBackendError::Other)?;
}
if let Some(index) = retained_vector {
t.vector_indexes
.write()
.live_mut()?
.insert(to.into(), index);
} else if let Some(dimensions) = vector_dimensions {
self.create_vector_field(&table_name, to, dimensions)?;
}
self.rename_catalog_index_column_refs(&table_name, from, to)?;
if self.is_persistent() {
if let Some(catalog) = self.storage.catalog.as_ref() {
catalog.rename_column_data(&table_name, from, to)?;
}
self.try_save_table_schema(&table_name, &t)?;
if let Some(dimensions) = vector_dimensions {
self.bind_renamed_vector_field(&table_name, to, dimensions, restore_vector)?;
}
}
let relation = Self::resolved_relation_identity(&table_name)?;
self.rewrite_routine_column_references(&relation, from, to)
.map_err(|error| {
StorageBackendError::Other(format!("rewrite routine column references: {error}"))
})?;
self.event_lifecycle_context()
.rename_event_column_inner(&table_name, from, to)?;
self.mark_column_stats_dirty(&table_name, &t)?;
self.refresh_value_indexes_for_table(&table_name)?;
Ok(true)
}
pub fn rename_table(&self, from: &str, to: &str) -> StorageBackendResult<bool> {
self.try_rename_table(from, to)
}
pub(crate) fn try_rename_table(&self, from: &str, to: &str) -> StorageBackendResult<bool> {
self.with_implicit_storage_transaction(|engine| engine.try_rename_table_inner(from, to))
}
pub(super) fn try_rename_table_inner(
&self,
from: &str,
to: &str,
) -> StorageBackendResult<bool> {
let Some(from) = self.resolve_table_ddl_target(from, "ALTER TABLE RENAME")? else {
return Ok(false);
};
let from_relation = Self::resolved_relation_identity(&from)?;
let (target_schema, target_name) =
RelationIdentity::parse_reference(to).map_err(StorageBackendError::Other)?;
let to_relation = RelationIdentity::new(
target_schema.unwrap_or_else(|| from_relation.schema.clone()),
target_name,
);
if to_relation.schema != self.temporary_schema_name()
&& !self
.durable
.schemas
.read()
.contains_key(&to_relation.schema)
{
return Err(StorageBackendError::Other(format!(
"schema `{}` does not exist",
to_relation.schema
)));
}
let to = to_relation.qualified_name();
if let Some(kind) = self.relation_kind_at(&to)? {
return Err(StorageBackendError::Other(format!(
"relation `{to}` already exists as {kind}"
)));
}
self.relation_creation_context()
.reserve_row_type_name(&to)
.map_err(|error| StorageBackendError::backend("ALTER TABLE name", error))?;
let persist_catalog = {
let tables = self.storage.tables.read();
if !tables.contains_key(&from_relation) || tables.contains_key(&to_relation) {
return Ok(false);
}
self.is_persistent()
&& tables.get(&from_relation).is_some_and(|table| {
table.persistence != uqa_sql::ast::RelationPersistence::Temporary
})
};
self.rewrite_table_rename_dependencies(&from, &to)?;
if persist_catalog {
if let Some(catalog) = self.storage.catalog.as_ref() {
catalog.rename_table_data(&from, &to)?;
}
}
let mut tables = self.storage.tables.write();
if tables.contains_key(&to_relation) {
return Ok(false);
}
let Some(state) = tables.remove(&from_relation) else {
return Ok(false);
};
tables.insert(to_relation.clone(), state.clone());
drop(tables);
self.rewrite_relation_rename_dependents(&from_relation, &to_relation)?;
self.rename_catalog_index_table_refs(&from, &to);
{
let mut analyzers = self.durable.table_field_analyzers.write();
let mut moved = Vec::new();
analyzers.retain(|(table, field), value| {
if table == &from {
moved.push(((to.clone(), field.clone()), value.clone()));
false
} else {
true
}
});
analyzers.extend(moved);
}
if persist_catalog {
self.rebind_persistent_table_stores(&to, &state)?;
self.try_save_table_schema(&to, &state)?;
if state.columns.read().iter().any(|column| {
column
.auto_increment
.as_ref()
.is_some_and(uqa_sql::ast::AutoIncrement::is_legacy)
}) {
self.persist_next_id(&to)?;
}
if let Some(catalog) = self.storage.catalog.as_ref() {
catalog.set_metadata(&table_next_id_metadata_key(&from), "")?;
}
}
self.mark_column_stats_dirty(&to, &state)?;
self.refresh_value_indexes_for_table(&to)?;
self.rename_constraint_transaction_relation(&from_relation, &to_relation);
Ok(true)
}
}