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, PartialEq)]
88pub enum CheckExpr {
89 True,
91 Col(u16),
93 Lit(Value),
95 IsNull(u16),
96 IsNotNull(u16),
97 Eq(Box<CheckExpr>, Box<CheckExpr>),
98 Ne(Box<CheckExpr>, Box<CheckExpr>),
99 Lt(Box<CheckExpr>, Box<CheckExpr>),
100 Le(Box<CheckExpr>, Box<CheckExpr>),
101 Gt(Box<CheckExpr>, Box<CheckExpr>),
102 Ge(Box<CheckExpr>, Box<CheckExpr>),
103 And(Box<CheckExpr>, Box<CheckExpr>),
104 Or(Box<CheckExpr>, Box<CheckExpr>),
105 Not(Box<CheckExpr>),
106}
107
108#[derive(Debug, Clone, Copy, PartialEq, Eq)]
110enum Tri {
111 True,
112 False,
113 Unknown,
114}
115
116impl CheckExpr {
117 pub fn satisfied(&self, cells: &HashMap<u16, Value>) -> bool {
120 !matches!(self.eval(cells), Tri::False)
121 }
122
123 fn eval(&self, cells: &HashMap<u16, Value>) -> Tri {
124 match self {
125 CheckExpr::True => Tri::True,
126 CheckExpr::Col(id) => match cells.get(id) {
127 None | Some(Value::Null) => Tri::Unknown,
128 Some(v) => Tri::from_truthy(v),
129 },
130 CheckExpr::Lit(v) => Tri::from_truthy(v),
131 CheckExpr::IsNull(id) => match cells.get(id) {
132 None | Some(Value::Null) => Tri::True,
133 Some(_) => Tri::False,
134 },
135 CheckExpr::IsNotNull(id) => match cells.get(id) {
136 None | Some(Value::Null) => Tri::False,
137 Some(_) => Tri::True,
138 },
139 CheckExpr::Eq(a, b) => compare(a.eval_term(cells), b.eval_term(cells), |o| {
140 o == Ordering::Equal
141 }),
142 CheckExpr::Ne(a, b) => compare(a.eval_term(cells), b.eval_term(cells), |o| {
143 o != Ordering::Equal
144 }),
145 CheckExpr::Lt(a, b) => compare(a.eval_term(cells), b.eval_term(cells), |o| {
146 o == Ordering::Less
147 }),
148 CheckExpr::Le(a, b) => compare(a.eval_term(cells), b.eval_term(cells), |o| {
149 o != Ordering::Greater
150 }),
151 CheckExpr::Gt(a, b) => compare(a.eval_term(cells), b.eval_term(cells), |o| {
152 o == Ordering::Greater
153 }),
154 CheckExpr::Ge(a, b) => compare(a.eval_term(cells), b.eval_term(cells), |o| {
155 o != Ordering::Less
156 }),
157 CheckExpr::And(a, b) => and3(a.eval(cells), b.eval(cells)),
158 CheckExpr::Or(a, b) => or3(a.eval(cells), b.eval(cells)),
159 CheckExpr::Not(a) => not3(a.eval(cells)),
160 }
161 }
162
163 fn eval_term(&self, cells: &HashMap<u16, Value>) -> Value {
166 match self {
167 CheckExpr::Col(id) => cells.get(id).cloned().unwrap_or(Value::Null),
168 CheckExpr::Lit(v) => v.clone(),
169 other => match other.eval(cells) {
170 Tri::True => Value::Int64(1),
171 Tri::False | Tri::Unknown => Value::Null,
172 },
173 }
174 }
175}
176
177impl Tri {
178 fn from_truthy(v: &Value) -> Tri {
179 match v {
180 Value::Null => Tri::Unknown,
181 Value::Bool(b) => {
182 if *b {
183 Tri::True
184 } else {
185 Tri::False
186 }
187 }
188 Value::Int64(n) => {
189 if *n != 0 {
190 Tri::True
191 } else {
192 Tri::False
193 }
194 }
195 Value::Float64(f) => {
196 if *f != 0.0 {
197 Tri::True
198 } else {
199 Tri::False
200 }
201 }
202 Value::Bytes(b) => {
203 if b.is_empty() {
204 Tri::False
205 } else {
206 Tri::True
207 }
208 }
209 Value::Embedding(v) => {
210 if v.is_empty() {
211 Tri::False
212 } else {
213 Tri::True
214 }
215 }
216 Value::Interval { .. } => Tri::Unknown,
217 Value::Uuid(_) | Value::Json(_) => Tri::Unknown,
218 Value::Decimal(d) => {
219 if *d != 0 {
220 Tri::True
221 } else {
222 Tri::False
223 }
224 }
225 }
226 }
227}
228
229fn and3(a: Tri, b: Tri) -> Tri {
230 match (a, b) {
231 (Tri::False, _) | (_, Tri::False) => Tri::False,
232 (Tri::Unknown, _) | (_, Tri::Unknown) => Tri::Unknown,
233 _ => Tri::True,
234 }
235}
236
237fn or3(a: Tri, b: Tri) -> Tri {
238 match (a, b) {
239 (Tri::True, _) | (_, Tri::True) => Tri::True,
240 (Tri::Unknown, _) | (_, Tri::Unknown) => Tri::Unknown,
241 _ => Tri::False,
242 }
243}
244
245fn not3(a: Tri) -> Tri {
246 match a {
247 Tri::True => Tri::False,
248 Tri::False => Tri::True,
249 Tri::Unknown => Tri::Unknown,
250 }
251}
252
253fn compare(a: Value, b: Value, pred: impl Fn(Ordering) -> bool) -> Tri {
257 if matches!(a, Value::Null) || matches!(b, Value::Null) {
258 return Tri::Unknown;
259 }
260 match value_cmp(&a, &b) {
261 Some(o) => {
262 if pred(o) {
263 Tri::True
264 } else {
265 Tri::False
266 }
267 }
268 None => Tri::Unknown,
269 }
270}
271
272pub(crate) fn value_cmp(a: &Value, b: &Value) -> Option<Ordering> {
276 match (a, b) {
277 (Value::Null, Value::Null) => Some(Ordering::Equal),
278 (Value::Bool(x), Value::Bool(y)) => Some((*x as u8).cmp(&(*y as u8))),
279 (Value::Int64(x), Value::Int64(y)) => Some(x.cmp(y)),
280 (Value::Float64(x), Value::Float64(y)) => x.partial_cmp(y),
281 (Value::Int64(x), Value::Float64(y)) => (*x as f64).partial_cmp(y),
282 (Value::Float64(x), Value::Int64(y)) => x.partial_cmp(&(*y as f64)),
283 (Value::Bytes(x), Value::Bytes(y)) => Some(x.cmp(y)),
284 (Value::Embedding(x), Value::Embedding(y)) => {
285 for (a, b) in x.iter().zip(y.iter()) {
286 match a.partial_cmp(b)? {
287 Ordering::Equal => continue,
288 non_eq => return Some(non_eq),
289 }
290 }
291 Some(x.len().cmp(&y.len()))
292 }
293 _ => None,
294 }
295}
296
297pub(crate) fn encode_composite_key(
302 columns: &[u16],
303 cells: &HashMap<u16, Value>,
304) -> Option<Vec<u8>> {
305 let mut out = Vec::new();
306 for cid in columns {
307 let v = cells.get(cid)?;
308 if matches!(v, Value::Null) {
309 return None;
310 }
311 let k = v.encode_key();
312 out.extend_from_slice(&(k.len() as u32).to_be_bytes());
314 out.extend_from_slice(&k);
315 }
316 Some(out)
317}
318
319pub(crate) fn validate_checks(
321 checks: &[CheckConstraint],
322 cells: &HashMap<u16, Value>,
323) -> Result<()> {
324 for c in checks {
325 if !c.expr.satisfied(cells) {
326 return Err(MongrelError::InvalidArgument(format!(
327 "CHECK constraint '{}' failed",
328 c.name
329 )));
330 }
331 }
332 Ok(())
333}
334
335#[cfg(test)]
336mod tests {
337 use super::*;
338
339 fn m(cols: &[(u16, Value)]) -> HashMap<u16, Value> {
340 cols.iter().cloned().collect()
341 }
342
343 #[test]
344 fn check_eq_literal() {
345 let e = CheckExpr::Eq(
347 Box::new(CheckExpr::Col(1)),
348 Box::new(CheckExpr::Lit(Value::Int64(5))),
349 );
350 assert!(e.satisfied(&m(&[(1, Value::Int64(5))])));
351 assert!(!e.satisfied(&m(&[(1, Value::Int64(6))])));
352 assert!(e.satisfied(&m(&[(1, Value::Null)])));
354 }
355
356 #[test]
357 fn check_range_and() {
358 let e = CheckExpr::And(
360 Box::new(CheckExpr::Ge(
361 Box::new(CheckExpr::Col(1)),
362 Box::new(CheckExpr::Lit(Value::Int64(0))),
363 )),
364 Box::new(CheckExpr::Le(
365 Box::new(CheckExpr::Col(1)),
366 Box::new(CheckExpr::Lit(Value::Int64(100))),
367 )),
368 );
369 assert!(e.satisfied(&m(&[(1, Value::Int64(50))])));
370 assert!(!e.satisfied(&m(&[(1, Value::Int64(101))])));
371 assert!(!e.satisfied(&m(&[(1, Value::Int64(-1))])));
372 }
373
374 #[test]
375 fn check_numeric_cross_type() {
376 let e = CheckExpr::Lt(
377 Box::new(CheckExpr::Col(1)),
378 Box::new(CheckExpr::Lit(Value::Float64(10.0))),
379 );
380 assert!(e.satisfied(&m(&[(1, Value::Int64(5))])));
381 assert!(!e.satisfied(&m(&[(1, Value::Int64(20))])));
382 }
383
384 #[test]
385 fn check_not_incomparable_is_unknown_passes() {
386 let e = CheckExpr::Eq(
388 Box::new(CheckExpr::Col(1)),
389 Box::new(CheckExpr::Lit(Value::Bytes(b"x".to_vec()))),
390 );
391 assert!(e.satisfied(&m(&[(1, Value::Int64(5))])));
392 }
393
394 #[test]
395 fn encode_composite_key_skips_null() {
396 let k = encode_composite_key(&[1, 2], &m(&[(1, Value::Int64(5)), (2, Value::Null)]));
397 assert!(k.is_none());
398 let k = encode_composite_key(&[1, 2], &m(&[(1, Value::Int64(5)), (2, Value::Int64(7))]));
399 assert!(k.is_some());
400 let k2 = encode_composite_key(&[1, 2], &m(&[(1, Value::Int64(6)), (2, Value::Int64(7))]));
402 assert_ne!(k.unwrap(), k2.unwrap());
403 }
404}