1use serde::{Deserialize, Serialize};
2
3use crate::constraint::TableConstraints;
4use crate::error::{MongrelError, Result};
5use crate::memtable::Value;
6
7#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
10#[serde(tag = "kind", rename_all = "lowercase")]
11pub enum TypeId {
12 Bool,
13 Int8,
14 Int16,
15 Int32,
16 Int64,
17 UInt8,
18 UInt16,
19 UInt32,
20 UInt64,
21 Float32,
22 Float64,
23 TimestampNanos,
24 Date32,
25 Date64,
28 Time64,
30 Interval,
33 Uuid,
35 Json,
38 Array {
43 element_type: u8,
44 },
45 Bytes,
47 Embedding {
49 dim: u32,
50 },
51 Decimal128 {
54 precision: u8,
55 scale: i8,
56 },
57}
58
59impl TypeId {
60 pub const fn fixed_size(self) -> Option<usize> {
62 match self {
63 TypeId::Bool => Some(1),
64 TypeId::Int8 | TypeId::UInt8 => Some(1),
65 TypeId::Int16 | TypeId::UInt16 => Some(2),
66 TypeId::Int32 | TypeId::UInt32 | TypeId::Float32 | TypeId::Date32 => Some(4),
67 TypeId::Int64
68 | TypeId::UInt64
69 | TypeId::Float64
70 | TypeId::TimestampNanos
71 | TypeId::Date64
72 | TypeId::Time64 => Some(8),
73 TypeId::Bytes | TypeId::Embedding { .. } => None,
74 TypeId::Decimal128 { .. } => Some(16),
75 TypeId::Uuid => Some(16),
76 TypeId::Json | TypeId::Array { .. } => None,
77 TypeId::Interval => Some(20), }
79 }
80}
81
82#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, Deserialize)]
84pub struct ColumnFlags {
85 bits: u32,
86}
87
88impl ColumnFlags {
89 pub const NULLABLE: u32 = 1 << 0;
90 pub const PRIMARY_KEY: u32 = 1 << 1;
91 pub const ENCRYPTED: u32 = 1 << 2;
92 pub const ENCRYPTED_INDEXABLE: u32 = 1 << 3;
95 pub const EMBEDDING_BINARY_QUANTIZED: u32 = 1 << 4;
97 pub const AUTO_INCREMENT: u32 = 1 << 5;
104
105 #[inline]
106 pub const fn empty() -> Self {
107 Self { bits: 0 }
108 }
109
110 #[inline]
111 pub const fn with(mut self, flag: u32) -> Self {
112 self.bits |= flag;
113 self
114 }
115
116 #[inline]
117 pub const fn without(mut self, flag: u32) -> Self {
118 self.bits &= !flag;
119 self
120 }
121
122 #[inline]
123 pub const fn contains(&self, flag: u32) -> bool {
124 self.bits & flag != 0
125 }
126
127 #[inline]
128 pub const fn bits(&self) -> u32 {
129 self.bits
130 }
131}
132
133#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
135pub struct ColumnDef {
136 pub id: u16,
137 pub name: String,
138 pub ty: TypeId,
139 pub flags: ColumnFlags,
140}
141
142#[derive(Debug, Clone, Default, Serialize, Deserialize)]
144pub struct AlterColumn {
145 pub name: Option<String>,
146 pub ty: Option<TypeId>,
147 pub flags: Option<ColumnFlags>,
148}
149
150impl AlterColumn {
151 pub fn rename(name: impl Into<String>) -> Self {
152 Self {
153 name: Some(name.into()),
154 ty: None,
155 flags: None,
156 }
157 }
158
159 pub fn set_type(ty: TypeId) -> Self {
160 Self {
161 name: None,
162 ty: Some(ty),
163 flags: None,
164 }
165 }
166
167 pub fn set_flags(flags: ColumnFlags) -> Self {
168 Self {
169 name: None,
170 ty: None,
171 flags: Some(flags),
172 }
173 }
174}
175
176#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
179pub enum IndexKind {
180 Bitmap,
182 FmIndex,
184 Ann,
186 LearnedRange,
188 MinHash,
190 Sparse,
192}
193
194#[derive(Debug, Clone, Serialize, Deserialize)]
195pub struct IndexDef {
196 pub name: String,
197 pub column_id: u16,
198 pub kind: IndexKind,
199 #[serde(default, skip_serializing_if = "Option::is_none")]
203 pub predicate: Option<String>,
204}
205
206#[derive(Debug, Clone, Default, Serialize, Deserialize)]
207pub struct Schema {
208 pub schema_id: u64,
209 pub columns: Vec<ColumnDef>,
210 pub indexes: Vec<IndexDef>,
211 #[serde(default)]
216 pub colocation: Vec<Vec<u16>>,
217 #[serde(default)]
222 pub constraints: TableConstraints,
223 #[serde(default)]
226 pub clustered: bool,
227}
228
229impl Schema {
230 pub fn column(&self, name: &str) -> Option<&ColumnDef> {
231 self.columns.iter().find(|c| c.name == name)
232 }
233
234 pub fn primary_key(&self) -> Option<&ColumnDef> {
235 self.columns
236 .iter()
237 .find(|c| c.flags.contains(ColumnFlags::PRIMARY_KEY))
238 }
239
240 pub fn validate_not_null(&self, columns: &[(u16, Value)]) -> Result<()> {
245 let at = |id: u16| columns.iter().find(|(c, _)| *c == id).map(|(_, v)| v);
248 for col in &self.columns {
249 if col.flags.contains(ColumnFlags::NULLABLE) {
250 continue;
251 }
252 if col.flags.contains(ColumnFlags::AUTO_INCREMENT) {
255 match at(col.id) {
256 None | Some(Value::Null) => continue,
257 Some(_) => {}
258 }
259 }
260 match at(col.id) {
261 None => {
262 return Err(MongrelError::InvalidArgument(format!(
263 "column '{}' ({}) is NOT NULL but was omitted",
264 col.name, col.id
265 )));
266 }
267 Some(Value::Null) => {
268 return Err(MongrelError::InvalidArgument(format!(
269 "column '{}' ({}) is NOT NULL but got NULL",
270 col.name, col.id
271 )));
272 }
273 Some(_) => {}
274 }
275 }
276 Ok(())
277 }
278
279 pub fn validate_auto_increment(&self) -> Result<()> {
283 let mut seen: Option<&ColumnDef> = None;
284 for col in &self.columns {
285 if !col.flags.contains(ColumnFlags::AUTO_INCREMENT) {
286 continue;
287 }
288 if let Some(prev) = seen {
289 return Err(MongrelError::Schema(format!(
290 "AUTO_INCREMENT may be set on at most one column; '{}' and '{}' both carry it",
291 prev.name, col.name
292 )));
293 }
294 if col.ty != TypeId::Int64 {
295 return Err(MongrelError::Schema(format!(
296 "AUTO_INCREMENT column '{}' must be Int64, is {:?}",
297 col.name, col.ty
298 )));
299 }
300 if !col.flags.contains(ColumnFlags::PRIMARY_KEY) {
301 return Err(MongrelError::Schema(format!(
302 "AUTO_INCREMENT column '{}' must also be the primary key",
303 col.name
304 )));
305 }
306 if col.flags.contains(ColumnFlags::NULLABLE) {
307 return Err(MongrelError::Schema(format!(
308 "AUTO_INCREMENT column '{}' must not be nullable",
309 col.name
310 )));
311 }
312 seen = Some(col);
313 }
314 Ok(())
315 }
316
317 pub fn auto_increment_column(&self) -> Option<&ColumnDef> {
319 self.columns
320 .iter()
321 .find(|c| c.flags.contains(ColumnFlags::AUTO_INCREMENT))
322 }
323}
324
325#[cfg(test)]
326mod tests {
327 use super::*;
328
329 #[test]
330 fn flag_composition() {
331 let f = ColumnFlags::empty()
332 .with(ColumnFlags::PRIMARY_KEY)
333 .with(ColumnFlags::ENCRYPTED_INDEXABLE);
334 assert!(f.contains(ColumnFlags::PRIMARY_KEY));
335 assert!(f.contains(ColumnFlags::ENCRYPTED_INDEXABLE));
336 assert!(!f.contains(ColumnFlags::ENCRYPTED));
337 }
338
339 #[test]
340 fn fixed_size() {
341 assert_eq!(TypeId::Int64.fixed_size(), Some(8));
342 assert_eq!(TypeId::Bytes.fixed_size(), None);
343 assert_eq!(TypeId::Embedding { dim: 768 }.fixed_size(), None);
344 }
345
346 fn col(id: u16, name: &str, ty: TypeId, flags: ColumnFlags) -> ColumnDef {
347 ColumnDef {
348 id,
349 name: name.into(),
350 ty,
351 flags,
352 }
353 }
354
355 #[test]
356 fn auto_increment_validation_accepts_int64_pk() {
357 let s = Schema {
358 schema_id: 1,
359 columns: vec![col(
360 0,
361 "id",
362 TypeId::Int64,
363 ColumnFlags::empty().with(ColumnFlags::PRIMARY_KEY | ColumnFlags::AUTO_INCREMENT),
364 )],
365 indexes: vec![],
366 colocation: vec![],
367 constraints: Default::default(),
368 clustered: false,
369 };
370 assert!(s.validate_auto_increment().is_ok());
371 assert_eq!(s.auto_increment_column().unwrap().id, 0);
372 }
373
374 #[test]
375 fn auto_increment_validation_rejects_non_pk() {
376 let s = Schema {
377 schema_id: 1,
378 columns: vec![
379 col(
380 0,
381 "id",
382 TypeId::Int64,
383 ColumnFlags::empty().with(ColumnFlags::PRIMARY_KEY),
384 ),
385 col(
386 1,
387 "seq",
388 TypeId::Int64,
389 ColumnFlags::empty().with(ColumnFlags::AUTO_INCREMENT),
390 ),
391 ],
392 indexes: vec![],
393 colocation: vec![],
394 constraints: Default::default(),
395 clustered: false,
396 };
397 assert!(s.validate_auto_increment().is_err());
398 }
399
400 #[test]
401 fn auto_increment_validation_rejects_non_int64() {
402 let s = Schema {
403 schema_id: 1,
404 columns: vec![col(
405 0,
406 "id",
407 TypeId::Bytes,
408 ColumnFlags::empty().with(ColumnFlags::PRIMARY_KEY | ColumnFlags::AUTO_INCREMENT),
409 )],
410 indexes: vec![],
411 colocation: vec![],
412 constraints: Default::default(),
413 clustered: false,
414 };
415 assert!(s.validate_auto_increment().is_err());
416 }
417
418 #[test]
419 fn auto_increment_validation_rejects_two() {
420 let s = Schema {
421 schema_id: 1,
422 columns: vec![
423 col(
424 0,
425 "id",
426 TypeId::Int64,
427 ColumnFlags::empty()
428 .with(ColumnFlags::PRIMARY_KEY | ColumnFlags::AUTO_INCREMENT),
429 ),
430 col(
431 1,
432 "id2",
433 TypeId::Int64,
434 ColumnFlags::empty().with(ColumnFlags::AUTO_INCREMENT),
435 ),
436 ],
437 indexes: vec![],
438 colocation: vec![],
439 constraints: Default::default(),
440 clustered: false,
441 };
442 assert!(s.validate_auto_increment().is_err());
443 }
444
445 #[test]
446 fn auto_increment_exempt_from_not_null_when_omitted() {
447 let s = Schema {
448 schema_id: 1,
449 columns: vec![
450 col(
451 0,
452 "id",
453 TypeId::Int64,
454 ColumnFlags::empty()
455 .with(ColumnFlags::PRIMARY_KEY | ColumnFlags::AUTO_INCREMENT),
456 ),
457 col(1, "name", TypeId::Bytes, ColumnFlags::empty()),
458 ],
459 indexes: vec![],
460 colocation: vec![],
461 constraints: Default::default(),
462 clustered: false,
463 };
464 let cols = vec![(1u16, Value::Bytes(b"x".to_vec()))];
466 assert!(s.validate_not_null(&cols).is_ok());
467 }
468}