#![allow(clippy::new_without_default)]
use vibesql_ast::{
pretty_print::ToSql, ColumnConstraintKind, ColumnDef, Expression, TableConstraint,
TableConstraintKind,
};
use vibesql_catalog::{ColumnSchema, TableSchema};
use vibesql_types::DataType;
use crate::errors::ExecutorError;
pub struct ConstraintResult {
pub primary_key: Option<Vec<String>>,
pub unique_constraints: Vec<Vec<String>>,
pub check_constraints: Vec<(String, Expression)>,
pub not_null_columns: Vec<String>,
}
impl ConstraintResult {
pub fn new() -> Self {
Self {
primary_key: None,
unique_constraints: Vec::new(),
check_constraints: Vec::new(),
not_null_columns: Vec::new(),
}
}
}
pub struct ConstraintValidator;
impl ConstraintValidator {
pub fn process_constraints(
columns: &[ColumnDef],
table_constraints: &[TableConstraint],
) -> Result<ConstraintResult, ExecutorError> {
let mut result = ConstraintResult::new();
let mut has_column_level_pk = false;
for col_def in columns {
for constraint in &col_def.constraints {
match &constraint.kind {
ColumnConstraintKind::PrimaryKey { .. } => {
if has_column_level_pk {
return Err(ExecutorError::MultiplePrimaryKeys);
}
if result.primary_key.is_some() {
return Err(ExecutorError::MultiplePrimaryKeys);
}
result.primary_key = Some(vec![col_def.name.clone()]);
if col_def.data_type == DataType::Integer {
result.not_null_columns.push(col_def.name.clone());
}
has_column_level_pk = true;
}
ColumnConstraintKind::Unique { .. } => {
result.unique_constraints.push(vec![col_def.name.clone()]);
}
ColumnConstraintKind::Check(expr) => {
let constraint_name =
constraint.name.clone().unwrap_or_else(|| expr.to_sql());
result.check_constraints.push((constraint_name, (**expr).clone()));
}
ColumnConstraintKind::NotNull
| ColumnConstraintKind::NotNullWithConflict { .. } => {
result.not_null_columns.push(col_def.name.clone());
}
ColumnConstraintKind::References { .. } => {
}
ColumnConstraintKind::AutoIncrement => {
}
ColumnConstraintKind::Key => {
}
ColumnConstraintKind::Collate(_) => {
}
}
}
}
for table_constraint in table_constraints {
match &table_constraint.kind {
TableConstraintKind::PrimaryKey { columns: pk_cols, .. } => {
if result.primary_key.is_some() {
return Err(ExecutorError::MultiplePrimaryKeys);
}
let column_names: Vec<String> =
pk_cols.iter().map(|c| c.expect_column_name().to_string()).collect();
result.primary_key = Some(column_names.clone());
for col_name in &column_names {
if let Some(col_def) = columns.iter().find(|c| &c.name == col_name) {
if col_def.data_type == DataType::Integer
&& !result.not_null_columns.contains(col_name)
{
result.not_null_columns.push(col_name.to_string());
}
}
}
}
TableConstraintKind::Unique { columns, .. } => {
let column_names: Vec<String> =
columns.iter().map(|c| c.expect_column_name().to_string()).collect();
result.unique_constraints.push(column_names);
}
TableConstraintKind::Check { expr } => {
let constraint_name =
table_constraint.name.clone().unwrap_or_else(|| expr.to_sql());
result.check_constraints.push((constraint_name, (**expr).clone()));
}
TableConstraintKind::ForeignKey { .. } => {
}
TableConstraintKind::Fulltext { .. } => {
}
}
}
Ok(result)
}
pub fn apply_to_columns(columns: &mut [ColumnSchema], constraint_result: &ConstraintResult) {
for col_name in &constraint_result.not_null_columns {
if let Some(col) = columns.iter_mut().find(|c| c.name == *col_name) {
col.nullable = false;
}
}
}
pub fn apply_to_schema(table_schema: &mut TableSchema, constraint_result: &ConstraintResult) {
if let Some(pk) = &constraint_result.primary_key {
table_schema.primary_key = Some(pk.clone());
}
table_schema.unique_constraints = constraint_result.unique_constraints.clone();
table_schema.check_constraints = constraint_result.check_constraints.clone();
}
}
#[cfg(test)]
mod tests {
use vibesql_ast::ColumnConstraint;
use vibesql_types::DataType;
use super::*;
fn make_column_def(name: &str, constraint_kinds: Vec<ColumnConstraintKind>) -> ColumnDef {
make_column_def_with_type(name, DataType::Integer, constraint_kinds)
}
fn make_column_def_with_type(
name: &str,
data_type: DataType,
constraint_kinds: Vec<ColumnConstraintKind>,
) -> ColumnDef {
ColumnDef {
name: name.to_string(),
data_type,
nullable: true,
constraints: constraint_kinds
.into_iter()
.map(|kind| ColumnConstraint { name: None, kind })
.collect(),
default_value: None,
comment: None,
generated_expr: None,
is_exact_integer_type: false,
}
}
#[test]
fn test_column_level_primary_key() {
let columns = vec![make_column_def(
"id",
vec![ColumnConstraintKind::PrimaryKey { on_conflict: None }],
)];
let result = ConstraintValidator::process_constraints(&columns, &[]).unwrap();
assert_eq!(result.primary_key, Some(vec!["id".to_string()]));
assert!(result.not_null_columns.contains(&"id".to_string()));
}
#[test]
fn test_table_level_primary_key() {
let columns = vec![make_column_def("id", vec![]), make_column_def("tenant_id", vec![])];
let constraints = vec![TableConstraint {
name: None,
kind: TableConstraintKind::PrimaryKey {
columns: vec![
vibesql_ast::IndexColumn::Column {
column_name: "id".to_string(),
direction: vibesql_ast::OrderDirection::Asc,
prefix_length: None,
},
vibesql_ast::IndexColumn::Column {
column_name: "tenant_id".to_string(),
direction: vibesql_ast::OrderDirection::Asc,
prefix_length: None,
},
],
on_conflict: None,
},
}];
let result = ConstraintValidator::process_constraints(&columns, &constraints).unwrap();
assert_eq!(result.primary_key, Some(vec!["id".to_string(), "tenant_id".to_string()]));
assert!(result.not_null_columns.contains(&"id".to_string()));
assert!(result.not_null_columns.contains(&"tenant_id".to_string()));
}
#[test]
fn test_multiple_primary_keys_fails() {
let columns = vec![make_column_def(
"id",
vec![ColumnConstraintKind::PrimaryKey { on_conflict: None }],
)];
let constraints = vec![TableConstraint {
name: None,
kind: TableConstraintKind::PrimaryKey {
columns: vec![vibesql_ast::IndexColumn::Column {
column_name: "id".to_string(),
direction: vibesql_ast::OrderDirection::Asc,
prefix_length: None,
}],
on_conflict: None,
},
}];
let result = ConstraintValidator::process_constraints(&columns, &constraints);
assert!(matches!(result, Err(ExecutorError::MultiplePrimaryKeys)));
}
#[test]
fn test_unique_constraints() {
let columns = vec![
make_column_def("email", vec![ColumnConstraintKind::Unique { on_conflict: None }]),
make_column_def("username", vec![]),
];
let constraints = vec![TableConstraint {
name: None,
kind: TableConstraintKind::Unique {
columns: vec![vibesql_ast::IndexColumn::Column {
column_name: "username".to_string(),
direction: vibesql_ast::OrderDirection::Asc,
prefix_length: None,
}],
on_conflict: None,
},
}];
let result = ConstraintValidator::process_constraints(&columns, &constraints).unwrap();
assert_eq!(result.unique_constraints.len(), 2);
assert!(result.unique_constraints.contains(&vec!["email".to_string()]));
assert!(result.unique_constraints.contains(&vec!["username".to_string()]));
}
#[test]
fn test_check_constraints() {
use vibesql_types::SqlValue;
let check_expr = Expression::BinaryOp {
left: Box::new(Expression::ColumnRef(vibesql_ast::ColumnIdentifier::simple(
"age", false,
))),
op: vibesql_ast::BinaryOperator::GreaterThan,
right: Box::new(Expression::Literal(SqlValue::Integer(0))),
};
let columns = vec![make_column_def(
"age",
vec![ColumnConstraintKind::Check(Box::new(check_expr.clone()))],
)];
let result = ConstraintValidator::process_constraints(&columns, &[]).unwrap();
assert_eq!(result.check_constraints.len(), 1);
assert_eq!(result.check_constraints[0].1, check_expr);
}
#[test]
fn test_apply_to_columns() {
let mut columns = vec![
ColumnSchema::new("id".to_string(), DataType::Integer, true),
ColumnSchema::new(
"name".to_string(),
DataType::Varchar { max_length: Some(100) },
true,
),
];
let mut result = ConstraintResult::new();
result.not_null_columns.push("id".to_string());
ConstraintValidator::apply_to_columns(&mut columns, &result);
assert!(!columns[0].nullable); assert!(columns[1].nullable); }
#[test]
fn test_text_primary_key_allows_null() {
let columns = vec![make_column_def_with_type(
"name",
DataType::Varchar { max_length: None },
vec![ColumnConstraintKind::PrimaryKey { on_conflict: None }],
)];
let result = ConstraintValidator::process_constraints(&columns, &[]).unwrap();
assert_eq!(result.primary_key, Some(vec!["name".to_string()]));
assert!(!result.not_null_columns.contains(&"name".to_string()));
}
#[test]
fn test_typeless_primary_key_allows_null() {
let columns = vec![make_column_def_with_type(
"c",
DataType::Varchar { max_length: None },
vec![ColumnConstraintKind::PrimaryKey { on_conflict: None }],
)];
let result = ConstraintValidator::process_constraints(&columns, &[]).unwrap();
assert_eq!(result.primary_key, Some(vec!["c".to_string()]));
assert!(!result.not_null_columns.contains(&"c".to_string()));
}
#[test]
fn test_integer_primary_key_has_not_null() {
let columns = vec![make_column_def_with_type(
"id",
DataType::Integer,
vec![ColumnConstraintKind::PrimaryKey { on_conflict: None }],
)];
let result = ConstraintValidator::process_constraints(&columns, &[]).unwrap();
assert_eq!(result.primary_key, Some(vec!["id".to_string()]));
assert!(result.not_null_columns.contains(&"id".to_string()));
}
#[test]
fn test_table_level_pk_with_mixed_types() {
let columns = vec![
make_column_def_with_type("id", DataType::Integer, vec![]),
make_column_def_with_type("code", DataType::Varchar { max_length: None }, vec![]),
];
let constraints = vec![TableConstraint {
name: None,
kind: TableConstraintKind::PrimaryKey {
columns: vec![
vibesql_ast::IndexColumn::Column {
column_name: "id".to_string(),
direction: vibesql_ast::OrderDirection::Asc,
prefix_length: None,
},
vibesql_ast::IndexColumn::Column {
column_name: "code".to_string(),
direction: vibesql_ast::OrderDirection::Asc,
prefix_length: None,
},
],
on_conflict: None,
},
}];
let result = ConstraintValidator::process_constraints(&columns, &constraints).unwrap();
assert_eq!(result.primary_key, Some(vec!["id".to_string(), "code".to_string()]));
assert!(result.not_null_columns.contains(&"id".to_string()));
assert!(!result.not_null_columns.contains(&"code".to_string()));
}
#[test]
fn test_real_primary_key_allows_null() {
let columns = vec![make_column_def_with_type(
"value",
DataType::Real,
vec![ColumnConstraintKind::PrimaryKey { on_conflict: None }],
)];
let result = ConstraintValidator::process_constraints(&columns, &[]).unwrap();
assert_eq!(result.primary_key, Some(vec!["value".to_string()]));
assert!(!result.not_null_columns.contains(&"value".to_string()));
}
#[test]
fn test_bigint_primary_key_allows_null() {
let columns = vec![make_column_def_with_type(
"big_id",
DataType::Bigint,
vec![ColumnConstraintKind::PrimaryKey { on_conflict: None }],
)];
let result = ConstraintValidator::process_constraints(&columns, &[]).unwrap();
assert_eq!(result.primary_key, Some(vec!["big_id".to_string()]));
assert!(!result.not_null_columns.contains(&"big_id".to_string()));
}
}