1use serde::{Deserialize, Serialize};
2use std::sync::Arc;
3
4use crate::constraint::TableConstraints;
5use crate::error::{MongrelError, Result};
6use crate::memtable::Value;
7
8#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
11#[serde(tag = "kind", rename_all = "lowercase")]
12pub enum TypeId {
13 Bool,
14 Int8,
15 Int16,
16 Int32,
17 Int64,
18 UInt8,
19 UInt16,
20 UInt32,
21 UInt64,
22 Float32,
23 Float64,
24 TimestampNanos,
25 Date32,
26 Date64,
29 Time64,
31 Interval,
34 Uuid,
36 Json,
39 Array {
44 element_type: u8,
45 },
46 Bytes,
48 Embedding {
50 dim: u32,
51 },
52 Decimal128 {
55 precision: u8,
56 scale: i8,
57 },
58 Enum {
64 variants: Arc<[String]>,
65 },
66}
67
68impl TypeId {
69 pub fn fixed_size(&self) -> Option<usize> {
71 match self {
72 TypeId::Bool => Some(1),
73 TypeId::Int8 | TypeId::UInt8 => Some(1),
74 TypeId::Int16 | TypeId::UInt16 => Some(2),
75 TypeId::Int32 | TypeId::UInt32 | TypeId::Float32 | TypeId::Date32 => Some(4),
76 TypeId::Int64
77 | TypeId::UInt64
78 | TypeId::Float64
79 | TypeId::TimestampNanos
80 | TypeId::Date64
81 | TypeId::Time64 => Some(8),
82 TypeId::Bytes | TypeId::Embedding { .. } | TypeId::Enum { .. } => None,
83 TypeId::Decimal128 { .. } => Some(16),
84 TypeId::Uuid => Some(16),
85 TypeId::Json | TypeId::Array { .. } => None,
86 TypeId::Interval => Some(20), }
88 }
89}
90
91#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, Deserialize)]
93pub struct ColumnFlags {
94 bits: u32,
95}
96
97impl ColumnFlags {
98 pub const NULLABLE: u32 = 1 << 0;
99 pub const PRIMARY_KEY: u32 = 1 << 1;
100 pub const ENCRYPTED: u32 = 1 << 2;
101 pub const ENCRYPTED_INDEXABLE: u32 = 1 << 3;
104 pub const EMBEDDING_BINARY_QUANTIZED: u32 = 1 << 4;
106 pub const AUTO_INCREMENT: u32 = 1 << 5;
113
114 #[inline]
115 pub const fn empty() -> Self {
116 Self { bits: 0 }
117 }
118
119 #[inline]
120 pub const fn with(mut self, flag: u32) -> Self {
121 self.bits |= flag;
122 self
123 }
124
125 #[inline]
126 pub const fn without(mut self, flag: u32) -> Self {
127 self.bits &= !flag;
128 self
129 }
130
131 #[inline]
132 pub const fn contains(&self, flag: u32) -> bool {
133 self.bits & flag != 0
134 }
135
136 #[inline]
137 pub const fn bits(&self) -> u32 {
138 self.bits
139 }
140}
141
142#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
147pub enum DefaultExpr {
148 Static(Value),
150 Now,
153 Uuid,
155}
156
157#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
159pub struct ColumnDef {
160 pub id: u16,
161 pub name: String,
162 pub ty: TypeId,
163 pub flags: ColumnFlags,
164 #[serde(default, skip_serializing_if = "Option::is_none")]
168 pub default_value: Option<DefaultExpr>,
169}
170
171#[derive(Debug, Clone, Default, Serialize, Deserialize)]
173pub struct AlterColumn {
174 pub name: Option<String>,
175 pub ty: Option<TypeId>,
176 pub flags: Option<ColumnFlags>,
177 #[serde(default, skip_serializing_if = "Option::is_none")]
180 pub default_value: Option<Option<DefaultExpr>>,
181}
182
183impl AlterColumn {
184 pub fn rename(name: impl Into<String>) -> Self {
185 Self {
186 name: Some(name.into()),
187 ty: None,
188 flags: None,
189 default_value: None,
190 }
191 }
192
193 pub fn set_type(ty: TypeId) -> Self {
194 Self {
195 name: None,
196 ty: Some(ty),
197 flags: None,
198 default_value: None,
199 }
200 }
201
202 pub fn set_flags(flags: ColumnFlags) -> Self {
203 Self {
204 name: None,
205 ty: None,
206 flags: Some(flags),
207 default_value: None,
208 }
209 }
210
211 pub fn set_default(expr: DefaultExpr) -> Self {
212 Self {
213 name: None,
214 ty: None,
215 flags: None,
216 default_value: Some(Some(expr)),
217 }
218 }
219
220 pub fn drop_default() -> Self {
221 Self {
222 name: None,
223 ty: None,
224 flags: None,
225 default_value: Some(None),
226 }
227 }
228}
229
230#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
233pub enum IndexKind {
234 Bitmap,
236 FmIndex,
238 Ann,
240 LearnedRange,
242 MinHash,
244 Sparse,
246}
247
248#[derive(Debug, Clone, Serialize, Deserialize)]
249pub struct IndexDef {
250 pub name: String,
251 pub column_id: u16,
252 pub kind: IndexKind,
253 #[serde(default, skip_serializing_if = "Option::is_none")]
257 pub predicate: Option<String>,
258}
259
260#[derive(Debug, Clone, Default, Serialize, Deserialize)]
261pub struct Schema {
262 pub schema_id: u64,
263 pub columns: Vec<ColumnDef>,
264 pub indexes: Vec<IndexDef>,
265 #[serde(default)]
270 pub colocation: Vec<Vec<u16>>,
271 #[serde(default)]
276 pub constraints: TableConstraints,
277 #[serde(default)]
280 pub clustered: bool,
281}
282
283impl Schema {
284 pub fn column(&self, name: &str) -> Option<&ColumnDef> {
285 self.columns.iter().find(|c| c.name == name)
286 }
287
288 pub fn primary_key(&self) -> Option<&ColumnDef> {
289 self.columns
290 .iter()
291 .find(|c| c.flags.contains(ColumnFlags::PRIMARY_KEY))
292 }
293
294 pub fn validate_not_null(&self, columns: &[(u16, Value)]) -> Result<()> {
299 let at = |id: u16| columns.iter().find(|(c, _)| *c == id).map(|(_, v)| v);
302 for col in &self.columns {
303 if col.flags.contains(ColumnFlags::NULLABLE) {
304 continue;
305 }
306 if col.flags.contains(ColumnFlags::AUTO_INCREMENT) {
309 match at(col.id) {
310 None | Some(Value::Null) => continue,
311 Some(_) => {}
312 }
313 }
314 match at(col.id) {
315 None => {
316 return Err(MongrelError::InvalidArgument(format!(
317 "column '{}' ({}) is NOT NULL but was omitted",
318 col.name, col.id
319 )));
320 }
321 Some(Value::Null) => {
322 return Err(MongrelError::InvalidArgument(format!(
323 "column '{}' ({}) is NOT NULL but got NULL",
324 col.name, col.id
325 )));
326 }
327 Some(_) => {}
328 }
329 }
330 Ok(())
331 }
332
333 pub fn validate_auto_increment(&self) -> Result<()> {
337 let mut seen: Option<&ColumnDef> = None;
338 for col in &self.columns {
339 if !col.flags.contains(ColumnFlags::AUTO_INCREMENT) {
340 continue;
341 }
342 if let Some(prev) = seen {
343 return Err(MongrelError::Schema(format!(
344 "AUTO_INCREMENT may be set on at most one column; '{}' and '{}' both carry it",
345 prev.name, col.name
346 )));
347 }
348 if col.ty != TypeId::Int64 {
349 return Err(MongrelError::Schema(format!(
350 "AUTO_INCREMENT column '{}' must be Int64, is {:?}",
351 col.name, col.ty
352 )));
353 }
354 if !col.flags.contains(ColumnFlags::PRIMARY_KEY) {
355 return Err(MongrelError::Schema(format!(
356 "AUTO_INCREMENT column '{}' must also be the primary key",
357 col.name
358 )));
359 }
360 if col.flags.contains(ColumnFlags::NULLABLE) {
361 return Err(MongrelError::Schema(format!(
362 "AUTO_INCREMENT column '{}' must not be nullable",
363 col.name
364 )));
365 }
366 seen = Some(col);
367 }
368 Ok(())
369 }
370
371 pub fn auto_increment_column(&self) -> Option<&ColumnDef> {
373 self.columns
374 .iter()
375 .find(|c| c.flags.contains(ColumnFlags::AUTO_INCREMENT))
376 }
377
378 pub fn validate_defaults(&self) -> Result<()> {
382 for col in &self.columns {
383 let Some(expr) = &col.default_value else {
384 continue;
385 };
386 match expr {
387 DefaultExpr::Static(v) => {
388 if !value_matches_type(v, col.ty.clone()) {
389 return Err(MongrelError::Schema(format!(
390 "DEFAULT value for column '{}' ({:?}) does not match type {:?}",
391 col.name, v, col.ty
392 )));
393 }
394 }
395 DefaultExpr::Now => {
396 if !matches!(
397 col.ty,
398 TypeId::Bytes | TypeId::TimestampNanos | TypeId::Date64
399 ) {
400 return Err(MongrelError::Schema(format!(
401 "DEFAULT NOW() on column '{}' requires Bytes/TimestampNanos/Date64, is {:?}",
402 col.name, col.ty
403 )));
404 }
405 }
406 DefaultExpr::Uuid => {
407 if !matches!(col.ty, TypeId::Uuid | TypeId::Bytes) {
408 return Err(MongrelError::Schema(format!(
409 "DEFAULT UUID() on column '{}' requires Uuid/Bytes, is {:?}",
410 col.name, col.ty
411 )));
412 }
413 }
414 }
415 }
416 Ok(())
417 }
418}
419
420pub(crate) fn value_matches_type(v: &Value, ty: TypeId) -> bool {
424 matches!(
425 (v, ty),
426 (Value::Null, _)
427 | (Value::Bool(_), TypeId::Bool)
428 | (
429 Value::Int64(_),
430 TypeId::Int8 | TypeId::Int16 | TypeId::Int32 | TypeId::Int64
431 )
432 | (Value::Float64(_), TypeId::Float32 | TypeId::Float64)
433 | (
434 Value::Bytes(_),
435 TypeId::Bytes
436 | TypeId::Json
437 | TypeId::Uuid
438 | TypeId::Date64
439 | TypeId::Time64
440 | TypeId::Enum { .. }
441 )
442 | (
443 Value::Int64(_),
444 TypeId::TimestampNanos | TypeId::Date32 | TypeId::Date64 | TypeId::Time64
445 )
446 | (Value::Uuid(_), TypeId::Uuid)
447 | (Value::Decimal(_), TypeId::Decimal128 { .. })
448 | (Value::Json(_), TypeId::Json)
449 | (Value::Embedding(_), TypeId::Embedding { .. })
450 | (Value::Interval { .. }, TypeId::Interval)
451 )
452}
453
454#[cfg(test)]
455mod tests {
456 use super::*;
457
458 #[test]
459 fn flag_composition() {
460 let f = ColumnFlags::empty()
461 .with(ColumnFlags::PRIMARY_KEY)
462 .with(ColumnFlags::ENCRYPTED_INDEXABLE);
463 assert!(f.contains(ColumnFlags::PRIMARY_KEY));
464 assert!(f.contains(ColumnFlags::ENCRYPTED_INDEXABLE));
465 assert!(!f.contains(ColumnFlags::ENCRYPTED));
466 }
467
468 #[test]
469 fn fixed_size() {
470 assert_eq!(TypeId::Int64.fixed_size(), Some(8));
471 assert_eq!(TypeId::Bytes.fixed_size(), None);
472 assert_eq!(TypeId::Embedding { dim: 768 }.fixed_size(), None);
473 }
474
475 fn col(id: u16, name: &str, ty: TypeId, flags: ColumnFlags) -> ColumnDef {
476 ColumnDef {
477 id,
478 name: name.into(),
479 ty,
480 flags,
481 default_value: None,
482 }
483 }
484
485 #[test]
486 fn auto_increment_validation_accepts_int64_pk() {
487 let s = Schema {
488 schema_id: 1,
489 columns: vec![col(
490 0,
491 "id",
492 TypeId::Int64,
493 ColumnFlags::empty().with(ColumnFlags::PRIMARY_KEY | ColumnFlags::AUTO_INCREMENT),
494 )],
495 indexes: vec![],
496 colocation: vec![],
497 constraints: Default::default(),
498 clustered: false,
499 };
500 assert!(s.validate_auto_increment().is_ok());
501 assert_eq!(s.auto_increment_column().unwrap().id, 0);
502 }
503
504 #[test]
505 fn auto_increment_validation_rejects_non_pk() {
506 let s = Schema {
507 schema_id: 1,
508 columns: vec![
509 col(
510 0,
511 "id",
512 TypeId::Int64,
513 ColumnFlags::empty().with(ColumnFlags::PRIMARY_KEY),
514 ),
515 col(
516 1,
517 "seq",
518 TypeId::Int64,
519 ColumnFlags::empty().with(ColumnFlags::AUTO_INCREMENT),
520 ),
521 ],
522 indexes: vec![],
523 colocation: vec![],
524 constraints: Default::default(),
525 clustered: false,
526 };
527 assert!(s.validate_auto_increment().is_err());
528 }
529
530 #[test]
531 fn auto_increment_validation_rejects_non_int64() {
532 let s = Schema {
533 schema_id: 1,
534 columns: vec![col(
535 0,
536 "id",
537 TypeId::Bytes,
538 ColumnFlags::empty().with(ColumnFlags::PRIMARY_KEY | ColumnFlags::AUTO_INCREMENT),
539 )],
540 indexes: vec![],
541 colocation: vec![],
542 constraints: Default::default(),
543 clustered: false,
544 };
545 assert!(s.validate_auto_increment().is_err());
546 }
547
548 #[test]
549 fn auto_increment_validation_rejects_two() {
550 let s = Schema {
551 schema_id: 1,
552 columns: vec![
553 col(
554 0,
555 "id",
556 TypeId::Int64,
557 ColumnFlags::empty()
558 .with(ColumnFlags::PRIMARY_KEY | ColumnFlags::AUTO_INCREMENT),
559 ),
560 col(
561 1,
562 "id2",
563 TypeId::Int64,
564 ColumnFlags::empty().with(ColumnFlags::AUTO_INCREMENT),
565 ),
566 ],
567 indexes: vec![],
568 colocation: vec![],
569 constraints: Default::default(),
570 clustered: false,
571 };
572 assert!(s.validate_auto_increment().is_err());
573 }
574
575 #[test]
576 fn auto_increment_exempt_from_not_null_when_omitted() {
577 let s = Schema {
578 schema_id: 1,
579 columns: vec![
580 col(
581 0,
582 "id",
583 TypeId::Int64,
584 ColumnFlags::empty()
585 .with(ColumnFlags::PRIMARY_KEY | ColumnFlags::AUTO_INCREMENT),
586 ),
587 col(1, "name", TypeId::Bytes, ColumnFlags::empty()),
588 ],
589 indexes: vec![],
590 colocation: vec![],
591 constraints: Default::default(),
592 clustered: false,
593 };
594 let cols = vec![(1u16, Value::Bytes(b"x".to_vec()))];
596 assert!(s.validate_not_null(&cols).is_ok());
597 }
598
599 fn col_with_default(
600 id: u16,
601 name: &str,
602 ty: TypeId,
603 flags: ColumnFlags,
604 dv: DefaultExpr,
605 ) -> ColumnDef {
606 ColumnDef {
607 id,
608 name: name.into(),
609 ty,
610 flags,
611 default_value: Some(dv),
612 }
613 }
614
615 #[test]
616 fn validate_defaults_accepts_matching_static() {
617 let s = Schema {
618 schema_id: 1,
619 columns: vec![col_with_default(
620 0,
621 "active",
622 TypeId::Bool,
623 ColumnFlags::empty(),
624 DefaultExpr::Static(Value::Bool(true)),
625 )],
626 indexes: vec![],
627 colocation: vec![],
628 constraints: Default::default(),
629 clustered: false,
630 };
631 assert!(s.validate_defaults().is_ok());
632 }
633
634 #[test]
635 fn validate_defaults_rejects_mismatched_static() {
636 let s = Schema {
637 schema_id: 1,
638 columns: vec![col_with_default(
639 0,
640 "count",
641 TypeId::Int64,
642 ColumnFlags::empty(),
643 DefaultExpr::Static(Value::Bytes(b"oops".to_vec())),
644 )],
645 indexes: vec![],
646 colocation: vec![],
647 constraints: Default::default(),
648 clustered: false,
649 };
650 assert!(s.validate_defaults().is_err());
651 }
652
653 #[test]
654 fn validate_defaults_now_requires_temporal_or_bytes() {
655 let ok = Schema {
656 schema_id: 1,
657 columns: vec![col_with_default(
658 0,
659 "ts",
660 TypeId::Bytes,
661 ColumnFlags::empty(),
662 DefaultExpr::Now,
663 )],
664 indexes: vec![],
665 colocation: vec![],
666 constraints: Default::default(),
667 clustered: false,
668 };
669 assert!(ok.validate_defaults().is_ok());
670
671 let bad = Schema {
672 schema_id: 1,
673 columns: vec![col_with_default(
674 0,
675 "ts",
676 TypeId::Int64,
677 ColumnFlags::empty(),
678 DefaultExpr::Now,
679 )],
680 indexes: vec![],
681 colocation: vec![],
682 constraints: Default::default(),
683 clustered: false,
684 };
685 assert!(bad.validate_defaults().is_err());
686 }
687
688 #[test]
689 fn validate_defaults_uuid_requires_uuid_or_bytes() {
690 let ok = Schema {
691 schema_id: 1,
692 columns: vec![col_with_default(
693 0,
694 "id",
695 TypeId::Uuid,
696 ColumnFlags::empty(),
697 DefaultExpr::Uuid,
698 )],
699 indexes: vec![],
700 colocation: vec![],
701 constraints: Default::default(),
702 clustered: false,
703 };
704 assert!(ok.validate_defaults().is_ok());
705
706 let bad = Schema {
707 schema_id: 1,
708 columns: vec![col_with_default(
709 0,
710 "id",
711 TypeId::Bool,
712 ColumnFlags::empty(),
713 DefaultExpr::Uuid,
714 )],
715 indexes: vec![],
716 colocation: vec![],
717 constraints: Default::default(),
718 clustered: false,
719 };
720 assert!(bad.validate_defaults().is_err());
721 }
722
723 #[test]
724 fn serde_roundtrip_column_def_with_default() {
725 let c = col_with_default(
726 0,
727 "x",
728 TypeId::Bytes,
729 ColumnFlags::empty(),
730 DefaultExpr::Static(Value::Bytes(b"hello".to_vec())),
731 );
732 let json = serde_json::to_string(&c).unwrap();
733 let de: ColumnDef = serde_json::from_str(&json).unwrap();
734 assert_eq!(c, de);
735 let old_json = r#"{"id":0,"name":"y","ty":{"kind":"bytes"},"flags":{"bits":0}}"#;
737 let old: ColumnDef = serde_json::from_str(old_json).unwrap();
738 assert!(old.default_value.is_none());
739 }
740}