1use crate::error::{MongrelError, Result};
16use crate::memtable::Value;
17use serde::{Deserialize, Serialize};
18use std::cmp::Ordering;
19use std::collections::HashMap;
20
21#[derive(Debug, Clone, Default, Serialize, Deserialize, PartialEq)]
23pub struct TableConstraints {
24 #[serde(default)]
25 pub uniques: Vec<UniqueConstraint>,
26 #[serde(default)]
27 pub foreign_keys: Vec<ForeignKey>,
28 #[serde(default)]
29 pub checks: Vec<CheckConstraint>,
30}
31
32impl TableConstraints {
33 pub fn is_empty(&self) -> bool {
34 self.uniques.is_empty() && self.foreign_keys.is_empty() && self.checks.is_empty()
35 }
36}
37
38#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
42pub struct UniqueConstraint {
43 pub id: u16,
44 pub name: String,
45 pub columns: Vec<u16>,
46}
47
48#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq, Default)]
50pub enum FkAction {
51 #[default]
53 Restrict,
54 Cascade,
56 SetNull,
58}
59
60#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
64pub struct ForeignKey {
65 pub id: u16,
66 pub name: String,
67 pub columns: Vec<u16>,
68 pub ref_table: String,
69 pub ref_columns: Vec<u16>,
70 #[serde(default)]
71 pub on_delete: FkAction,
72}
73
74#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
77pub struct CheckConstraint {
78 pub id: u16,
79 pub name: String,
80 pub expr: CheckExpr,
81}
82
83#[derive(Debug, Clone, Serialize, Deserialize)]
93pub enum CheckExpr {
94 True,
96 Col(u16),
98 Lit(Value),
100 IsNull(u16),
101 IsNotNull(u16),
102 Eq(Box<CheckExpr>, Box<CheckExpr>),
103 Ne(Box<CheckExpr>, Box<CheckExpr>),
104 Lt(Box<CheckExpr>, Box<CheckExpr>),
105 Le(Box<CheckExpr>, Box<CheckExpr>),
106 Gt(Box<CheckExpr>, Box<CheckExpr>),
107 Ge(Box<CheckExpr>, Box<CheckExpr>),
108 And(Box<CheckExpr>, Box<CheckExpr>),
109 Or(Box<CheckExpr>, Box<CheckExpr>),
110 Not(Box<CheckExpr>),
111 Regex {
118 col: u16,
119 pattern: String,
120 negated: bool,
121 case_insensitive: bool,
122 #[serde(skip)]
123 cached: std::sync::OnceLock<regex::Regex>,
124 },
125}
126
127impl PartialEq for CheckExpr {
128 fn eq(&self, other: &Self) -> bool {
129 match (self, other) {
130 (Self::True, Self::True) => true,
131 (Self::Col(a), Self::Col(b)) => a == b,
132 (Self::Lit(a), Self::Lit(b)) => a == b,
133 (Self::IsNull(a), Self::IsNull(b)) => a == b,
134 (Self::IsNotNull(a), Self::IsNotNull(b)) => a == b,
135 (Self::Eq(a1, a2), Self::Eq(b1, b2)) => a1 == b1 && a2 == b2,
136 (Self::Ne(a1, a2), Self::Ne(b1, b2)) => a1 == b1 && a2 == b2,
137 (Self::Lt(a1, a2), Self::Lt(b1, b2)) => a1 == b1 && a2 == b2,
138 (Self::Le(a1, a2), Self::Le(b1, b2)) => a1 == b1 && a2 == b2,
139 (Self::Gt(a1, a2), Self::Gt(b1, b2)) => a1 == b1 && a2 == b2,
140 (Self::Ge(a1, a2), Self::Ge(b1, b2)) => a1 == b1 && a2 == b2,
141 (Self::And(a1, a2), Self::And(b1, b2)) => a1 == b1 && a2 == b2,
142 (Self::Or(a1, a2), Self::Or(b1, b2)) => a1 == b1 && a2 == b2,
143 (Self::Not(a), Self::Not(b)) => a == b,
144 (
145 Self::Regex {
146 col: ac,
147 pattern: ap,
148 negated: an,
149 case_insensitive: ai,
150 ..
151 },
152 Self::Regex {
153 col: bc,
154 pattern: bp,
155 negated: bn,
156 case_insensitive: bi,
157 ..
158 },
159 ) => ac == bc && ap == bp && an == bn && ai == bi,
160 _ => false,
161 }
162 }
163}
164
165#[derive(Debug, Clone, Copy, PartialEq, Eq)]
167enum Tri {
168 True,
169 False,
170 Unknown,
171}
172
173impl CheckExpr {
174 pub fn satisfied(&self, cells: &HashMap<u16, Value>) -> bool {
177 !matches!(self.eval(cells), Tri::False)
178 }
179
180 pub fn validate(&self) -> Result<()> {
185 match self {
186 CheckExpr::Eq(a, b)
187 | CheckExpr::Ne(a, b)
188 | CheckExpr::Lt(a, b)
189 | CheckExpr::Le(a, b)
190 | CheckExpr::Gt(a, b)
191 | CheckExpr::Ge(a, b)
192 | CheckExpr::And(a, b)
193 | CheckExpr::Or(a, b) => {
194 a.validate()?;
195 b.validate()?;
196 }
197 CheckExpr::Not(a) => a.validate()?,
198 CheckExpr::Regex { pattern, .. } => {
199 regex::Regex::new(pattern).map_err(|e| {
200 MongrelError::InvalidArgument(format!("invalid regex pattern: {e}"))
201 })?;
202 }
203 _ => {}
204 }
205 Ok(())
206 }
207
208 fn eval(&self, cells: &HashMap<u16, Value>) -> Tri {
209 match self {
210 CheckExpr::True => Tri::True,
211 CheckExpr::Col(id) => match cells.get(id) {
212 None | Some(Value::Null) => Tri::Unknown,
213 Some(v) => Tri::from_truthy(v),
214 },
215 CheckExpr::Lit(v) => Tri::from_truthy(v),
216 CheckExpr::IsNull(id) => match cells.get(id) {
217 None | Some(Value::Null) => Tri::True,
218 Some(_) => Tri::False,
219 },
220 CheckExpr::IsNotNull(id) => match cells.get(id) {
221 None | Some(Value::Null) => Tri::False,
222 Some(_) => Tri::True,
223 },
224 CheckExpr::Eq(a, b) => compare(a.eval_term(cells), b.eval_term(cells), |o| {
225 o == Ordering::Equal
226 }),
227 CheckExpr::Ne(a, b) => compare(a.eval_term(cells), b.eval_term(cells), |o| {
228 o != Ordering::Equal
229 }),
230 CheckExpr::Lt(a, b) => compare(a.eval_term(cells), b.eval_term(cells), |o| {
231 o == Ordering::Less
232 }),
233 CheckExpr::Le(a, b) => compare(a.eval_term(cells), b.eval_term(cells), |o| {
234 o != Ordering::Greater
235 }),
236 CheckExpr::Gt(a, b) => compare(a.eval_term(cells), b.eval_term(cells), |o| {
237 o == Ordering::Greater
238 }),
239 CheckExpr::Ge(a, b) => compare(a.eval_term(cells), b.eval_term(cells), |o| {
240 o != Ordering::Less
241 }),
242 CheckExpr::And(a, b) => and3(a.eval(cells), b.eval(cells)),
243 CheckExpr::Or(a, b) => or3(a.eval(cells), b.eval(cells)),
244 CheckExpr::Not(a) => not3(a.eval(cells)),
245 CheckExpr::Regex {
246 col,
247 pattern,
248 negated,
249 case_insensitive,
250 cached,
251 } => match cells.get(col) {
252 None | Some(Value::Null) => Tri::Unknown,
253 Some(Value::Bytes(b)) => {
254 let re = cached.get_or_init(|| {
255 let mut builder = regex::RegexBuilder::new(pattern);
256 builder.case_insensitive(*case_insensitive);
257 builder
261 .build()
262 .unwrap_or_else(|_| regex::Regex::new("$^").unwrap())
263 });
264 let matched = re.is_match(std::str::from_utf8(b).unwrap_or(""));
265 match (*negated, matched) {
266 (false, true) | (true, false) => Tri::True,
267 (false, false) | (true, true) => Tri::False,
268 }
269 }
270 Some(_) => Tri::Unknown,
272 },
273 }
274 }
275
276 fn eval_term(&self, cells: &HashMap<u16, Value>) -> Value {
279 match self {
280 CheckExpr::Col(id) => cells.get(id).cloned().unwrap_or(Value::Null),
281 CheckExpr::Lit(v) => v.clone(),
282 other => match other.eval(cells) {
283 Tri::True => Value::Int64(1),
284 Tri::False | Tri::Unknown => Value::Null,
285 },
286 }
287 }
288}
289
290impl Tri {
291 fn from_truthy(v: &Value) -> Tri {
292 match v {
293 Value::Null => Tri::Unknown,
294 Value::Bool(b) => {
295 if *b {
296 Tri::True
297 } else {
298 Tri::False
299 }
300 }
301 Value::Int64(n) => {
302 if *n != 0 {
303 Tri::True
304 } else {
305 Tri::False
306 }
307 }
308 Value::Float64(f) => {
309 if *f != 0.0 {
310 Tri::True
311 } else {
312 Tri::False
313 }
314 }
315 Value::Bytes(b) => {
316 if b.is_empty() {
317 Tri::False
318 } else {
319 Tri::True
320 }
321 }
322 Value::Embedding(v) => {
323 if v.is_empty() {
324 Tri::False
325 } else {
326 Tri::True
327 }
328 }
329 Value::Interval { .. } => Tri::Unknown,
330 Value::Uuid(_) | Value::Json(_) => Tri::Unknown,
331 Value::Decimal(d) => {
332 if *d != 0 {
333 Tri::True
334 } else {
335 Tri::False
336 }
337 }
338 }
339 }
340}
341
342fn and3(a: Tri, b: Tri) -> Tri {
343 match (a, b) {
344 (Tri::False, _) | (_, Tri::False) => Tri::False,
345 (Tri::Unknown, _) | (_, Tri::Unknown) => Tri::Unknown,
346 _ => Tri::True,
347 }
348}
349
350fn or3(a: Tri, b: Tri) -> Tri {
351 match (a, b) {
352 (Tri::True, _) | (_, Tri::True) => Tri::True,
353 (Tri::Unknown, _) | (_, Tri::Unknown) => Tri::Unknown,
354 _ => Tri::False,
355 }
356}
357
358fn not3(a: Tri) -> Tri {
359 match a {
360 Tri::True => Tri::False,
361 Tri::False => Tri::True,
362 Tri::Unknown => Tri::Unknown,
363 }
364}
365
366fn compare(a: Value, b: Value, pred: impl Fn(Ordering) -> bool) -> Tri {
370 if matches!(a, Value::Null) || matches!(b, Value::Null) {
371 return Tri::Unknown;
372 }
373 match value_cmp(&a, &b) {
374 Some(o) => {
375 if pred(o) {
376 Tri::True
377 } else {
378 Tri::False
379 }
380 }
381 None => Tri::Unknown,
382 }
383}
384
385pub(crate) fn value_cmp(a: &Value, b: &Value) -> Option<Ordering> {
389 match (a, b) {
390 (Value::Null, Value::Null) => Some(Ordering::Equal),
391 (Value::Bool(x), Value::Bool(y)) => Some((*x as u8).cmp(&(*y as u8))),
392 (Value::Int64(x), Value::Int64(y)) => Some(x.cmp(y)),
393 (Value::Float64(x), Value::Float64(y)) => x.partial_cmp(y),
394 (Value::Int64(x), Value::Float64(y)) => (*x as f64).partial_cmp(y),
395 (Value::Float64(x), Value::Int64(y)) => x.partial_cmp(&(*y as f64)),
396 (Value::Bytes(x), Value::Bytes(y)) => Some(x.cmp(y)),
397 (Value::Embedding(x), Value::Embedding(y)) => {
398 for (a, b) in x.iter().zip(y.iter()) {
399 match a.partial_cmp(b)? {
400 Ordering::Equal => continue,
401 non_eq => return Some(non_eq),
402 }
403 }
404 Some(x.len().cmp(&y.len()))
405 }
406 _ => None,
407 }
408}
409
410pub(crate) fn encode_composite_key(
415 columns: &[u16],
416 cells: &HashMap<u16, Value>,
417) -> Option<Vec<u8>> {
418 let mut out = Vec::new();
419 for cid in columns {
420 let v = cells.get(cid)?;
421 if matches!(v, Value::Null) {
422 return None;
423 }
424 let k = v.encode_key();
425 out.extend_from_slice(&(k.len() as u32).to_be_bytes());
427 out.extend_from_slice(&k);
428 }
429 Some(out)
430}
431
432pub(crate) fn validate_checks(
434 checks: &[CheckConstraint],
435 cells: &HashMap<u16, Value>,
436) -> Result<()> {
437 for c in checks {
438 if !c.expr.satisfied(cells) {
439 return Err(MongrelError::InvalidArgument(format!(
440 "CHECK constraint '{}' failed",
441 c.name
442 )));
443 }
444 }
445 Ok(())
446}
447
448#[cfg(test)]
449mod tests {
450 use super::*;
451
452 fn m(cols: &[(u16, Value)]) -> HashMap<u16, Value> {
453 cols.iter().cloned().collect()
454 }
455
456 #[test]
457 fn check_eq_literal() {
458 let e = CheckExpr::Eq(
460 Box::new(CheckExpr::Col(1)),
461 Box::new(CheckExpr::Lit(Value::Int64(5))),
462 );
463 assert!(e.satisfied(&m(&[(1, Value::Int64(5))])));
464 assert!(!e.satisfied(&m(&[(1, Value::Int64(6))])));
465 assert!(e.satisfied(&m(&[(1, Value::Null)])));
467 }
468
469 #[test]
470 fn check_range_and() {
471 let e = CheckExpr::And(
473 Box::new(CheckExpr::Ge(
474 Box::new(CheckExpr::Col(1)),
475 Box::new(CheckExpr::Lit(Value::Int64(0))),
476 )),
477 Box::new(CheckExpr::Le(
478 Box::new(CheckExpr::Col(1)),
479 Box::new(CheckExpr::Lit(Value::Int64(100))),
480 )),
481 );
482 assert!(e.satisfied(&m(&[(1, Value::Int64(50))])));
483 assert!(!e.satisfied(&m(&[(1, Value::Int64(101))])));
484 assert!(!e.satisfied(&m(&[(1, Value::Int64(-1))])));
485 }
486
487 #[test]
488 fn check_numeric_cross_type() {
489 let e = CheckExpr::Lt(
490 Box::new(CheckExpr::Col(1)),
491 Box::new(CheckExpr::Lit(Value::Float64(10.0))),
492 );
493 assert!(e.satisfied(&m(&[(1, Value::Int64(5))])));
494 assert!(!e.satisfied(&m(&[(1, Value::Int64(20))])));
495 }
496
497 #[test]
498 fn check_not_incomparable_is_unknown_passes() {
499 let e = CheckExpr::Eq(
501 Box::new(CheckExpr::Col(1)),
502 Box::new(CheckExpr::Lit(Value::Bytes(b"x".to_vec()))),
503 );
504 assert!(e.satisfied(&m(&[(1, Value::Int64(5))])));
505 }
506
507 #[test]
508 fn encode_composite_key_skips_null() {
509 let k = encode_composite_key(&[1, 2], &m(&[(1, Value::Int64(5)), (2, Value::Null)]));
510 assert!(k.is_none());
511 let k = encode_composite_key(&[1, 2], &m(&[(1, Value::Int64(5)), (2, Value::Int64(7))]));
512 assert!(k.is_some());
513 let k2 = encode_composite_key(&[1, 2], &m(&[(1, Value::Int64(6)), (2, Value::Int64(7))]));
515 assert_ne!(k.unwrap(), k2.unwrap());
516 }
517
518 fn regex_expr(col: u16, pattern: &str, negated: bool, ci: bool) -> CheckExpr {
519 CheckExpr::Regex {
520 col,
521 pattern: pattern.to_string(),
522 negated,
523 case_insensitive: ci,
524 cached: std::sync::OnceLock::new(),
525 }
526 }
527
528 #[test]
529 fn check_regex_match() {
530 let e = regex_expr(1, r"^a\w+z$", false, false);
531 assert!(e.satisfied(&m(&[(1, Value::Bytes(b"abz".to_vec()))])));
532 assert!(e.satisfied(&m(&[(1, Value::Bytes(b"alfredz".to_vec()))])));
533 assert!(!e.satisfied(&m(&[(1, Value::Bytes(b"abc".to_vec()))])));
534 }
535
536 #[test]
537 fn check_regex_null_is_unknown() {
538 let e = regex_expr(1, r"^\d+$", false, false);
539 assert!(e.satisfied(&m(&[(1, Value::Null)])));
540 assert!(e.satisfied(&m(&[]))); }
542
543 #[test]
544 fn check_regex_negated() {
545 let e = regex_expr(1, r"^\d+$", true, false);
546 assert!(e.satisfied(&m(&[(1, Value::Bytes(b"abc".to_vec()))])));
547 assert!(!e.satisfied(&m(&[(1, Value::Bytes(b"123".to_vec()))])));
548 }
549
550 #[test]
551 fn check_regex_case_insensitive() {
552 let e = regex_expr(1, r"^hello$", false, true);
553 assert!(e.satisfied(&m(&[(1, Value::Bytes(b"HELLO".to_vec()))])));
554 assert!(e.satisfied(&m(&[(1, Value::Bytes(b"Hello".to_vec()))])));
555 }
556
557 #[test]
558 fn check_regex_non_bytes_is_unknown() {
559 let e = regex_expr(1, r"\d+", false, false);
560 assert!(e.satisfied(&m(&[(1, Value::Int64(42))])));
562 }
563
564 #[test]
565 fn check_regex_validate_rejects_invalid_pattern() {
566 assert!(regex_expr(1, "[", false, false).validate().is_err());
567 assert!(regex_expr(1, r"^\d+$", false, false).validate().is_ok());
568 }
569
570 #[test]
571 fn check_regex_partial_eq_ignores_cache() {
572 let a = regex_expr(1, r"^\d+$", false, false);
573 let b = regex_expr(1, r"^\d+$", false, false);
574 assert_eq!(a, b);
575 a.satisfied(&m(&[(1, Value::Bytes(b"1".to_vec()))]));
577 assert_eq!(a, b);
578 }
579
580 #[test]
581 fn check_regex_serde_roundtrip() {
582 let e = regex_expr(1, r"^\w+@[\w.]+$", true, true);
583 let json = serde_json::to_string(&e).unwrap();
584 assert!(!json.contains("cached"));
586 let de: CheckExpr = serde_json::from_str(&json).unwrap();
587 assert_eq!(e, de);
588 assert!(de.satisfied(&m(&[(1, Value::Bytes(b"not-an-email".to_vec()))])));
590 assert!(!de.satisfied(&m(&[(1, Value::Bytes(b"a@b.com".to_vec()))])));
591 }
592
593 #[test]
594 fn check_regex_inside_logical_ops() {
595 let e = CheckExpr::And(
597 Box::new(regex_expr(1, r"^\d+$", false, false)),
598 Box::new(CheckExpr::IsNotNull(2)),
599 );
600 assert!(e.satisfied(&m(&[
601 (1, Value::Bytes(b"42".to_vec())),
602 (2, Value::Int64(1))
603 ])));
604 assert!(!e.satisfied(&m(&[(1, Value::Bytes(b"42".to_vec())), (2, Value::Null)])));
605 assert!(!e.satisfied(&m(&[
606 (1, Value::Bytes(b"abc".to_vec())),
607 (2, Value::Int64(1))
608 ])));
609 }
610}