uqa-sql 0.5.4

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

//! `PostgreSQL` rules for the values a statement writes to the columns whose values a table generates: identity columns and generated columns.

use crate::{
    assignment::columns::AssignmentColumnCatalog,
    ast::{AutoIncrementKind, OverridingKind},
    SQLError,
};

/// Stored expressions affected by an UPDATE's target columns. Explicitly assigning an unchanged input still selects its dependents; assigning a generated column to DEFAULT does not select it by itself. Callers with BEFORE ROW UPDATE triggers must recompute every stored expression instead.
#[must_use]
pub fn stored_columns_for_update(
    columns: &[crate::ast::ColumnDef],
    updated_columns: &[String],
) -> Vec<String> {
    columns
        .iter()
        .filter(|column| {
            column.generated.as_ref().is_some_and(|generated| {
                generated.kind == crate::ast::GeneratedColumnKind::Stored
                    && updated_columns.iter().any(|updated| {
                        crate::schema::dependencies::schema_expr_references_column(
                            &generated.expression,
                            updated,
                        )
                    })
            })
        })
        .map(|column| column.name.clone())
        .collect()
}

/// How a table generates a column's value.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum Generation {
    IdentityAlways,
    IdentityByDefault,
    Expression,
}

/// The identity and generated columns of one table, in column order, read once for a statement that writes it. `PostgreSQL` applies these rules when it rewrites the statement, column by column in table order, before it reads or writes a row, so a statement whose source yields no row is rejected as well.
pub struct GeneratedValueColumns {
    columns: Vec<(String, Generation)>,
}

impl GeneratedValueColumns {
    /// The identity and generated columns of `table`; none for a relation the catalog does not describe as a table.
    pub fn of(catalog: &dyn AssignmentColumnCatalog, table: &str) -> Result<Self, SQLError> {
        let columns = catalog
            .try_describe_table(table)
            .map_err(|error| SQLError::Internal(format!("read generated columns: {error}")))?
            .unwrap_or_default()
            .into_iter()
            .filter_map(|column| {
                let generation = match column.auto_increment.as_ref() {
                    Some(provenance) if provenance.kind == AutoIncrementKind::IdentityAlways => {
                        Generation::IdentityAlways
                    }
                    Some(provenance) if provenance.is_identity() => Generation::IdentityByDefault,
                    _ if column.generated.is_some() => Generation::Expression,
                    _ => return None,
                };
                Some((column.name, generation))
            })
            .collect();
        Ok(Self { columns })
    }

    /// Whether `column` is an identity column, whose value `OVERRIDING USER VALUE` leaves to its sequence.
    #[must_use]
    pub fn contains(&self, column: &str) -> bool {
        self.columns.iter().any(|(name, generation)| {
            name == column
                && matches!(
                    generation,
                    Generation::IdentityAlways | Generation::IdentityByDefault
                )
        })
    }

    /// Reject a value an `INSERT` supplies for a generated column, and for a `GENERATED ALWAYS` identity column without an `OVERRIDING` clause, which a generated column ignores. `targets` names each target column with whether some row supplies a value other than `DEFAULT` for it.
    pub fn validate_insert<'a>(
        &self,
        targets: impl IntoIterator<Item = (&'a str, bool)>,
        overriding: Option<OverridingKind>,
    ) -> Result<(), SQLError> {
        let supplied = targets
            .into_iter()
            .filter_map(|(column, supplied)| supplied.then_some(column))
            .collect::<std::collections::BTreeSet<_>>();
        for (column, generation) in &self.columns {
            if !supplied.contains(column.as_str()) {
                continue;
            }
            match generation {
                Generation::Expression => return Err(generated_column_insert_error(column)),
                Generation::IdentityAlways if overriding.is_none() => {
                    return Err(generated_always_insert_error(column));
                }
                Generation::IdentityAlways | Generation::IdentityByDefault => {}
            }
        }
        Ok(())
    }

    /// Reject an assignment other than `DEFAULT` to a generated column or a `GENERATED ALWAYS` identity column by an `UPDATE`, an `ON CONFLICT DO UPDATE` or a `MERGE` update. `assignments` names each assigned column with whether its value is `DEFAULT`.
    pub fn validate_update<'a>(
        &self,
        assignments: impl IntoIterator<Item = (&'a str, bool)>,
    ) -> Result<(), SQLError> {
        let assigned = assignments
            .into_iter()
            .filter_map(|(column, default)| (!default).then_some(column))
            .collect::<std::collections::BTreeSet<_>>();
        for (column, generation) in &self.columns {
            if !assigned.contains(column.as_str()) {
                continue;
            }
            match generation {
                Generation::Expression => return Err(generated_column_update_error(column)),
                Generation::IdentityAlways => return Err(generated_always_update_error(column)),
                Generation::IdentityByDefault => {}
            }
        }
        Ok(())
    }
}

/// The `428C9` error of a value an `INSERT` supplies for a `GENERATED ALWAYS` identity column.
#[must_use]
pub fn generated_always_insert_error(column: &str) -> SQLError {
    SQLError::Diagnostic {
        sqlstate: "428C9".into(),
        message: format!("cannot insert a non-DEFAULT value into column \"{column}\""),
        detail: Some(format!(
            "Column \"{column}\" is an identity column defined as GENERATED ALWAYS."
        )),
        hint: Some("Use OVERRIDING SYSTEM VALUE to override.".into()),
    }
}

/// The `428C9` error of an assignment other than `DEFAULT` to a `GENERATED ALWAYS` identity column.
#[must_use]
pub fn generated_always_update_error(column: &str) -> SQLError {
    SQLError::Diagnostic {
        sqlstate: "428C9".into(),
        message: format!("column \"{column}\" can only be updated to DEFAULT"),
        detail: Some(format!(
            "Column \"{column}\" is an identity column defined as GENERATED ALWAYS."
        )),
        hint: None,
    }
}

/// The `428C9` error of a value an `INSERT` supplies for a generated column.
#[must_use]
pub fn generated_column_insert_error(column: &str) -> SQLError {
    SQLError::Diagnostic {
        sqlstate: "428C9".into(),
        message: format!("cannot insert a non-DEFAULT value into column \"{column}\""),
        detail: Some(format!("Column \"{column}\" is a generated column.")),
        hint: None,
    }
}

/// The `428C9` error of an assignment other than `DEFAULT` to a generated column.
#[must_use]
pub fn generated_column_update_error(column: &str) -> SQLError {
    SQLError::Diagnostic {
        sqlstate: "428C9".into(),
        message: format!("column \"{column}\" can only be updated to DEFAULT"),
        detail: Some(format!("Column \"{column}\" is a generated column.")),
        hint: None,
    }
}

#[cfg(test)]
mod tests {
    #[test]
    fn stored_update_selection_follows_expression_inputs() {
        let crate::Statement::CreateTable(table) = crate::compile("CREATE TABLE t(a int,b int,unrelated int,ga int GENERATED ALWAYS AS(coalesce(a,1)) STORED,gb int GENERATED ALWAYS AS(b+2) STORED,c int GENERATED ALWAYS AS(42) STORED,v int GENERATED ALWAYS AS(b+3) VIRTUAL)").unwrap().remove(0) else { panic!("table") };
        for (targets, expected) in [
            (vec!["a"], vec!["ga"]),
            (vec!["b"], vec!["gb"]),
            (vec!["a", "b"], vec!["ga", "gb"]),
            (vec!["unrelated"], vec![]),
            (vec!["ga", "c", "v"], vec![]),
        ] {
            let targets = targets.into_iter().map(str::to_owned).collect::<Vec<_>>();
            assert_eq!(
                super::stored_columns_for_update(&table.columns, &targets),
                expected
            );
        }
    }
}