uqa-execution 0.5.0

//
// Unified Query Algebra
//
// Copyright (c) 2023-2026 Cognica, Inc.
//

//! Bind, persist, and publish declared columns while retaining the original column write lock.
use super::{materialize_metadata, resolve_table_name, table_not_found, SchemaPublicationContext};
use uqa_sql::ast::{
    ColumnDef, ColumnType, Expr, ForeignKey, GeneratedColumn, TableCheck, TableKeyConstraint,
};
use uqa_sql::schema::columns::publication::{self, ColumnProperty};
use uqa_sql::schema::dependencies::{regclass, registration};
use uqa_storage::{StorageBackendError, StorageBackendResult};

/// A retained column write guard; candidate persistence runs before its value is replaced.
pub trait ColumnSchemaWrite {
    fn columns(&self) -> &[ColumnDef];
    fn publish(&mut self, columns: Vec<ColumnDef>);
}

pub fn set_column_default(
    context: &SchemaPublicationContext<'_>,
    table: &str,
    column: &str,
    default: Option<Expr>,
) -> StorageBackendResult<bool> {
    set_column_property(context, table, column, ColumnProperty::Default(default))
}
pub fn set_column_generated(
    context: &SchemaPublicationContext<'_>,
    table: &str,
    column: &str,
    generated: Option<GeneratedColumn>,
) -> StorageBackendResult<bool> {
    set_column_property(context, table, column, ColumnProperty::Generated(generated))
}
pub fn set_column_type(
    context: &SchemaPublicationContext<'_>,
    table: &str,
    column: &str,
    ty: &ColumnType,
) -> StorageBackendResult<bool> {
    set_column_property(context, table, column, ColumnProperty::Type(ty))
}
pub fn set_column_auto_increment(
    context: &SchemaPublicationContext<'_>,
    table: &str,
    column: &str,
    provenance: Option<uqa_sql::ast::AutoIncrement>,
) -> StorageBackendResult<bool> {
    set_column_property(
        context,
        table,
        column,
        ColumnProperty::AutoIncrement(provenance),
    )
}
fn set_column_property(
    context: &SchemaPublicationContext<'_>,
    table: &str,
    column: &str,
    mut property: ColumnProperty<'_>,
) -> StorageBackendResult<bool> {
    let table_name = resolve_table_name(context.catalog, table)?;
    let state = context
        .catalog
        .table_state(&table_name)?
        .ok_or_else(|| table_not_found(&table_name))?;
    match &mut property {
        ColumnProperty::Default(Some(default)) => {
            regclass::bind_sequence_references_in_expr(context.bindings.references, default)
                .map_err(StorageBackendError::Other)?;
            registration::bind_default_routine_identities(
                &context.bindings,
                &table_name,
                column,
                default,
            )
            .map_err(StorageBackendError::Other)?;
        }
        ColumnProperty::Generated(Some(generated)) => {
            regclass::bind_sequence_references_in_expr(
                context.bindings.references,
                &mut generated.expression,
            )
            .map_err(StorageBackendError::Other)?;
        }
        _ => {}
    }
    // Reserving the OID of a stored expression reads a catalog snapshot, which includes this table's columns, so it precedes the columns' write guard. The relation lock of the altering statement keeps the column as the candidate saw it.
    let mut candidate = state.columns();
    publication::apply_property(&mut candidate, &table_name, column, property.clone())
        .map_err(StorageBackendError::Other)?;
    let definition = publication::column_mut(&mut candidate, &table_name, column)
        .map_err(StorageBackendError::Other)?;
    uqa_sql::schema::constraint_metadata::identity::materialize_default_oid(
        definition,
        &mut context.identity_allocator(),
    )
    .map_err(|error| StorageBackendError::backend("column default identity", error))?;
    let default_catalog_oid = definition.default_catalog_oid;
    // The reservation refreshes the catalog, so the columns are published into the table state it holds now.
    let state = super::current_table_state(context.catalog, &table_name, state.as_ref())?;
    let mut guard = state.write_columns();
    let mut next = guard.columns().to_vec();
    publication::apply_property(&mut next, &table_name, column, property)
        .map_err(StorageBackendError::Other)?;
    publication::column_mut(&mut next, &table_name, column)
        .map_err(StorageBackendError::Other)?
        .default_catalog_oid = default_catalog_oid;
    state.mark_statistics_dirty()?;
    state.persist_columns(&next)?;
    guard.publish(next);
    Ok(true)
}
pub fn set_column_not_null(
    context: &SchemaPublicationContext<'_>,
    table: &str,
    column: &str,
    not_null: bool,
) -> StorageBackendResult<bool> {
    let table_name = resolve_table_name(context.catalog, table)?;
    let state = context
        .catalog
        .table_state(&table_name)?
        .ok_or_else(|| table_not_found(&table_name))?;
    let mut next = state.columns();
    publication::set_not_null(&mut next, &table_name, column, not_null)
        .map_err(StorageBackendError::Other)?;
    let mut constraints = state.constraints();
    materialize_metadata(context, &table_name, &mut next, &mut constraints)?;
    let indexes = crate::schema::indexes::registry::prepare_constraint_indexes(
        &context.indexes,
        &table_name,
        state.object_id(),
        &mut next,
        &mut constraints,
    )?;
    let state = super::current_table_state(context.catalog, &table_name, state.as_ref())?;
    state.mark_statistics_dirty()?;
    state.persist_candidate(&next, &constraints)?;
    state.publish_constraints(next, constraints);
    indexes.publish(&context.indexes)?;
    Ok(true)
}
pub fn register_table_constraints(
    context: &SchemaPublicationContext<'_>,
    table: &str,
    checks: Vec<TableCheck>,
    mut foreign_keys: Vec<ForeignKey>,
    key_constraints: Vec<TableKeyConstraint>,
) -> StorageBackendResult<()> {
    let Some(table_name) = context.catalog.resolve_table_name(table)? else {
        return Err(StorageBackendError::Other(format!(
            "unknown table `{table}` while registering constraints"
        )));
    };
    let Some(state) = context.catalog.table_state(&table_name)? else {
        return Err(StorageBackendError::Other(format!(
            "unknown table `{table_name}` while registering constraints"
        )));
    };
    for foreign_key in &mut foreign_keys {
        foreign_key.ref_table = resolve_table_name(context.catalog, &foreign_key.ref_table)?;
    }
    let mut constraints = state.constraint_header();
    constraints.checks = checks;
    constraints.foreign_keys = foreign_keys;
    constraints.key_constraints = key_constraints;
    // Validate the stored relation identity before taking the declared column snapshot.
    let relation = uqa_core::RelationIdentity::from_legacy_name(&table_name)
        .map_err(StorageBackendError::Other)?;
    let mut columns = state.columns();
    registration::bind_table_schema_routine_identities(
        &context.bindings,
        &table_name,
        &mut columns,
        &mut constraints.checks,
    )
    .map_err(StorageBackendError::Other)?;
    let mut allocate = context.identity_allocator();
    let names = context.constraint_names().name_scope(&relation);
    uqa_sql::schema::constraint_metadata::materialize_constraint_metadata_with_names(
        &relation,
        &mut columns,
        &mut constraints,
        &mut allocate,
        &names,
    )
    .map_err(|error| StorageBackendError::backend("constraint identity", error))?;
    super::referenced_partitions::reconcile_derived_constraints(
        context,
        &mut columns,
        &mut constraints,
        &names.schema,
        &mut allocate,
    )?;
    let indexes = crate::schema::indexes::registry::prepare_constraint_indexes(
        &context.indexes,
        &table_name,
        state.object_id(),
        &mut columns,
        &mut constraints,
    )?;
    let state = super::current_table_state(context.catalog, &table_name, state.as_ref())?;
    state.persist_candidate(&columns, &constraints)?;
    state.publish_constraints(columns, constraints);
    indexes.publish(&context.indexes)?;
    Ok(())
}