1#![allow(clippy::new_without_default)]
7
8use vibesql_ast::{
9 pretty_print::ToSql, ColumnConstraintKind, ColumnDef, Expression, TableConstraint,
10 TableConstraintKind,
11};
12use vibesql_catalog::{ColumnSchema, TableSchema};
13use vibesql_types::DataType;
14
15use crate::errors::ExecutorError;
16
17pub struct ConstraintResult {
19 pub primary_key: Option<Vec<String>>,
21 pub unique_constraints: Vec<Vec<String>>,
23 pub check_constraints: Vec<(String, Expression)>,
25 pub not_null_columns: Vec<String>,
27}
28
29impl ConstraintResult {
30 pub fn new() -> Self {
32 Self {
33 primary_key: None,
34 unique_constraints: Vec::new(),
35 check_constraints: Vec::new(),
36 not_null_columns: Vec::new(),
37 }
38 }
39}
40
41pub struct ConstraintValidator;
43
44impl ConstraintValidator {
45 pub fn process_constraints(
60 columns: &[ColumnDef],
61 table_constraints: &[TableConstraint],
62 ) -> Result<ConstraintResult, ExecutorError> {
63 let mut result = ConstraintResult::new();
64
65 let mut has_column_level_pk = false;
67
68 for col_def in columns {
70 for constraint in &col_def.constraints {
71 match &constraint.kind {
72 ColumnConstraintKind::PrimaryKey { .. } => {
73 if has_column_level_pk {
74 return Err(ExecutorError::MultiplePrimaryKeys);
75 }
76 if result.primary_key.is_some() {
77 return Err(ExecutorError::MultiplePrimaryKeys);
78 }
79 result.primary_key = Some(vec![col_def.name.clone()]);
80 if col_def.data_type == DataType::Integer {
83 result.not_null_columns.push(col_def.name.clone());
84 }
85 has_column_level_pk = true;
86 }
87 ColumnConstraintKind::Unique { .. } => {
88 result.unique_constraints.push(vec![col_def.name.clone()]);
89 }
90 ColumnConstraintKind::Check(expr) => {
91 let constraint_name =
94 constraint.name.clone().unwrap_or_else(|| expr.to_sql());
95 result.check_constraints.push((constraint_name, (**expr).clone()));
96 }
97 ColumnConstraintKind::NotNull
98 | ColumnConstraintKind::NotNullWithConflict { .. } => {
99 result.not_null_columns.push(col_def.name.clone());
100 }
101 ColumnConstraintKind::References { .. } => {
102 }
105 ColumnConstraintKind::AutoIncrement => {
106 }
110 ColumnConstraintKind::Key => {
111 }
115 ColumnConstraintKind::Collate(_) => {
116 }
120 }
121 }
122 }
123
124 for table_constraint in table_constraints {
126 match &table_constraint.kind {
127 TableConstraintKind::PrimaryKey { columns: pk_cols, .. } => {
128 if result.primary_key.is_some() {
130 return Err(ExecutorError::MultiplePrimaryKeys);
131 }
132 let column_names: Vec<String> =
134 pk_cols.iter().map(|c| c.expect_column_name().to_string()).collect();
135 result.primary_key = Some(column_names.clone());
136 for col_name in &column_names {
139 if let Some(col_def) = columns.iter().find(|c| &c.name == col_name) {
140 if col_def.data_type == DataType::Integer
141 && !result.not_null_columns.contains(col_name)
142 {
143 result.not_null_columns.push(col_name.to_string());
144 }
145 }
146 }
147 }
148 TableConstraintKind::Unique { columns, .. } => {
149 let column_names: Vec<String> =
151 columns.iter().map(|c| c.expect_column_name().to_string()).collect();
152 result.unique_constraints.push(column_names);
153 }
154 TableConstraintKind::Check { expr } => {
155 let constraint_name =
158 table_constraint.name.clone().unwrap_or_else(|| expr.to_sql());
159 result.check_constraints.push((constraint_name, (**expr).clone()));
160 }
161 TableConstraintKind::ForeignKey { .. } => {
162 }
165 TableConstraintKind::Fulltext { .. } => {
166 }
170 }
171 }
172
173 Ok(result)
174 }
175
176 pub fn apply_to_columns(columns: &mut [ColumnSchema], constraint_result: &ConstraintResult) {
185 for col_name in &constraint_result.not_null_columns {
187 if let Some(col) = columns.iter_mut().find(|c| c.name == *col_name) {
188 col.nullable = false;
189 }
190 }
191 }
192
193 pub fn apply_to_schema(table_schema: &mut TableSchema, constraint_result: &ConstraintResult) {
202 if let Some(pk) = &constraint_result.primary_key {
204 table_schema.primary_key = Some(pk.clone());
205 }
206
207 table_schema.unique_constraints = constraint_result.unique_constraints.clone();
209
210 table_schema.check_constraints = constraint_result.check_constraints.clone();
212 }
213}
214
215#[cfg(test)]
216mod tests {
217 use vibesql_ast::ColumnConstraint;
218 use vibesql_types::DataType;
219
220 use super::*;
221
222 fn make_column_def(name: &str, constraint_kinds: Vec<ColumnConstraintKind>) -> ColumnDef {
223 make_column_def_with_type(name, DataType::Integer, constraint_kinds)
224 }
225
226 fn make_column_def_with_type(
227 name: &str,
228 data_type: DataType,
229 constraint_kinds: Vec<ColumnConstraintKind>,
230 ) -> ColumnDef {
231 ColumnDef {
232 name: name.to_string(),
233 data_type,
234 nullable: true,
235 constraints: constraint_kinds
236 .into_iter()
237 .map(|kind| ColumnConstraint { name: None, kind })
238 .collect(),
239 default_value: None,
240 comment: None,
241 generated_expr: None,
242 is_exact_integer_type: false,
243 }
244 }
245
246 #[test]
247 fn test_column_level_primary_key() {
248 let columns = vec![make_column_def(
249 "id",
250 vec![ColumnConstraintKind::PrimaryKey { on_conflict: None }],
251 )];
252 let result = ConstraintValidator::process_constraints(&columns, &[]).unwrap();
253
254 assert_eq!(result.primary_key, Some(vec!["id".to_string()]));
255 assert!(result.not_null_columns.contains(&"id".to_string()));
256 }
257
258 #[test]
259 fn test_table_level_primary_key() {
260 let columns = vec![make_column_def("id", vec![]), make_column_def("tenant_id", vec![])];
261 let constraints = vec![TableConstraint {
262 name: None,
263 kind: TableConstraintKind::PrimaryKey {
264 columns: vec![
265 vibesql_ast::IndexColumn::Column {
266 column_name: "id".to_string(),
267 direction: vibesql_ast::OrderDirection::Asc,
268 prefix_length: None,
269 },
270 vibesql_ast::IndexColumn::Column {
271 column_name: "tenant_id".to_string(),
272 direction: vibesql_ast::OrderDirection::Asc,
273 prefix_length: None,
274 },
275 ],
276 on_conflict: None,
277 },
278 }];
279
280 let result = ConstraintValidator::process_constraints(&columns, &constraints).unwrap();
281
282 assert_eq!(result.primary_key, Some(vec!["id".to_string(), "tenant_id".to_string()]));
283 assert!(result.not_null_columns.contains(&"id".to_string()));
284 assert!(result.not_null_columns.contains(&"tenant_id".to_string()));
285 }
286
287 #[test]
288 fn test_multiple_primary_keys_fails() {
289 let columns = vec![make_column_def(
290 "id",
291 vec![ColumnConstraintKind::PrimaryKey { on_conflict: None }],
292 )];
293 let constraints = vec![TableConstraint {
294 name: None,
295 kind: TableConstraintKind::PrimaryKey {
296 columns: vec![vibesql_ast::IndexColumn::Column {
297 column_name: "id".to_string(),
298 direction: vibesql_ast::OrderDirection::Asc,
299 prefix_length: None,
300 }],
301 on_conflict: None,
302 },
303 }];
304
305 let result = ConstraintValidator::process_constraints(&columns, &constraints);
306 assert!(matches!(result, Err(ExecutorError::MultiplePrimaryKeys)));
307 }
308
309 #[test]
310 fn test_unique_constraints() {
311 let columns = vec![
312 make_column_def("email", vec![ColumnConstraintKind::Unique { on_conflict: None }]),
313 make_column_def("username", vec![]),
314 ];
315 let constraints = vec![TableConstraint {
316 name: None,
317 kind: TableConstraintKind::Unique {
318 columns: vec![vibesql_ast::IndexColumn::Column {
319 column_name: "username".to_string(),
320 direction: vibesql_ast::OrderDirection::Asc,
321 prefix_length: None,
322 }],
323 on_conflict: None,
324 },
325 }];
326
327 let result = ConstraintValidator::process_constraints(&columns, &constraints).unwrap();
328
329 assert_eq!(result.unique_constraints.len(), 2);
330 assert!(result.unique_constraints.contains(&vec!["email".to_string()]));
331 assert!(result.unique_constraints.contains(&vec!["username".to_string()]));
332 }
333
334 #[test]
335 fn test_check_constraints() {
336 use vibesql_types::SqlValue;
337
338 let check_expr = Expression::BinaryOp {
339 left: Box::new(Expression::ColumnRef(vibesql_ast::ColumnIdentifier::simple(
340 "age", false,
341 ))),
342 op: vibesql_ast::BinaryOperator::GreaterThan,
343 right: Box::new(Expression::Literal(SqlValue::Integer(0))),
344 };
345
346 let columns = vec![make_column_def(
347 "age",
348 vec![ColumnConstraintKind::Check(Box::new(check_expr.clone()))],
349 )];
350
351 let result = ConstraintValidator::process_constraints(&columns, &[]).unwrap();
352
353 assert_eq!(result.check_constraints.len(), 1);
354 assert_eq!(result.check_constraints[0].1, check_expr);
355 }
356
357 #[test]
358 fn test_apply_to_columns() {
359 let mut columns = vec![
360 ColumnSchema::new("id".to_string(), DataType::Integer, true),
361 ColumnSchema::new(
362 "name".to_string(),
363 DataType::Varchar { max_length: Some(100) },
364 true,
365 ),
366 ];
367
368 let mut result = ConstraintResult::new();
369 result.not_null_columns.push("id".to_string());
370
371 ConstraintValidator::apply_to_columns(&mut columns, &result);
372
373 assert!(!columns[0].nullable); assert!(columns[1].nullable); }
376
377 #[test]
380 fn test_text_primary_key_allows_null() {
381 let columns = vec![make_column_def_with_type(
383 "name",
384 DataType::Varchar { max_length: None },
385 vec![ColumnConstraintKind::PrimaryKey { on_conflict: None }],
386 )];
387 let result = ConstraintValidator::process_constraints(&columns, &[]).unwrap();
388
389 assert_eq!(result.primary_key, Some(vec!["name".to_string()]));
390 assert!(!result.not_null_columns.contains(&"name".to_string()));
392 }
393
394 #[test]
395 fn test_typeless_primary_key_allows_null() {
396 let columns = vec![make_column_def_with_type(
398 "c",
399 DataType::Varchar { max_length: None },
400 vec![ColumnConstraintKind::PrimaryKey { on_conflict: None }],
401 )];
402 let result = ConstraintValidator::process_constraints(&columns, &[]).unwrap();
403
404 assert_eq!(result.primary_key, Some(vec!["c".to_string()]));
405 assert!(!result.not_null_columns.contains(&"c".to_string()));
406 }
407
408 #[test]
409 fn test_integer_primary_key_has_not_null() {
410 let columns = vec![make_column_def_with_type(
412 "id",
413 DataType::Integer,
414 vec![ColumnConstraintKind::PrimaryKey { on_conflict: None }],
415 )];
416 let result = ConstraintValidator::process_constraints(&columns, &[]).unwrap();
417
418 assert_eq!(result.primary_key, Some(vec!["id".to_string()]));
419 assert!(result.not_null_columns.contains(&"id".to_string()));
420 }
421
422 #[test]
423 fn test_table_level_pk_with_mixed_types() {
424 let columns = vec![
426 make_column_def_with_type("id", DataType::Integer, vec![]),
427 make_column_def_with_type("code", DataType::Varchar { max_length: None }, vec![]),
428 ];
429 let constraints = vec![TableConstraint {
430 name: None,
431 kind: TableConstraintKind::PrimaryKey {
432 columns: vec![
433 vibesql_ast::IndexColumn::Column {
434 column_name: "id".to_string(),
435 direction: vibesql_ast::OrderDirection::Asc,
436 prefix_length: None,
437 },
438 vibesql_ast::IndexColumn::Column {
439 column_name: "code".to_string(),
440 direction: vibesql_ast::OrderDirection::Asc,
441 prefix_length: None,
442 },
443 ],
444 on_conflict: None,
445 },
446 }];
447
448 let result = ConstraintValidator::process_constraints(&columns, &constraints).unwrap();
449
450 assert_eq!(result.primary_key, Some(vec!["id".to_string(), "code".to_string()]));
451 assert!(result.not_null_columns.contains(&"id".to_string()));
453 assert!(!result.not_null_columns.contains(&"code".to_string()));
454 }
455
456 #[test]
457 fn test_real_primary_key_allows_null() {
458 let columns = vec![make_column_def_with_type(
460 "value",
461 DataType::Real,
462 vec![ColumnConstraintKind::PrimaryKey { on_conflict: None }],
463 )];
464 let result = ConstraintValidator::process_constraints(&columns, &[]).unwrap();
465
466 assert_eq!(result.primary_key, Some(vec!["value".to_string()]));
467 assert!(!result.not_null_columns.contains(&"value".to_string()));
468 }
469
470 #[test]
471 fn test_bigint_primary_key_allows_null() {
472 let columns = vec![make_column_def_with_type(
474 "big_id",
475 DataType::Bigint,
476 vec![ColumnConstraintKind::PrimaryKey { on_conflict: None }],
477 )];
478 let result = ConstraintValidator::process_constraints(&columns, &[]).unwrap();
479
480 assert_eq!(result.primary_key, Some(vec!["big_id".to_string()]));
481 assert!(!result.not_null_columns.contains(&"big_id".to_string()));
483 }
484}