1use serde::{Deserialize, Serialize};
2use std::collections::HashMap;
3
4use crate::constraint::TableConstraints;
5use crate::error::{MongrelError, Result};
6use crate::memtable::Value;
7
8#[derive(Debug, Clone, Copy, 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 Bytes,
28 Embedding {
30 dim: u32,
31 },
32}
33
34impl TypeId {
35 pub const fn fixed_size(self) -> Option<usize> {
37 match self {
38 TypeId::Bool => Some(1),
39 TypeId::Int8 | TypeId::UInt8 => Some(1),
40 TypeId::Int16 | TypeId::UInt16 => Some(2),
41 TypeId::Int32 | TypeId::UInt32 | TypeId::Float32 | TypeId::Date32 => Some(4),
42 TypeId::Int64 | TypeId::UInt64 | TypeId::Float64 | TypeId::TimestampNanos => Some(8),
43 TypeId::Bytes | TypeId::Embedding { .. } => None,
44 }
45 }
46}
47
48#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, Deserialize)]
50pub struct ColumnFlags {
51 bits: u32,
52}
53
54impl ColumnFlags {
55 pub const NULLABLE: u32 = 1 << 0;
56 pub const PRIMARY_KEY: u32 = 1 << 1;
57 pub const ENCRYPTED: u32 = 1 << 2;
58 pub const ENCRYPTED_INDEXABLE: u32 = 1 << 3;
61 pub const EMBEDDING_BINARY_QUANTIZED: u32 = 1 << 4;
63 pub const AUTO_INCREMENT: u32 = 1 << 5;
70
71 #[inline]
72 pub const fn empty() -> Self {
73 Self { bits: 0 }
74 }
75
76 #[inline]
77 pub const fn with(mut self, flag: u32) -> Self {
78 self.bits |= flag;
79 self
80 }
81
82 #[inline]
83 pub const fn without(mut self, flag: u32) -> Self {
84 self.bits &= !flag;
85 self
86 }
87
88 #[inline]
89 pub const fn contains(&self, flag: u32) -> bool {
90 self.bits & flag != 0
91 }
92
93 #[inline]
94 pub const fn bits(&self) -> u32 {
95 self.bits
96 }
97}
98
99#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
101pub struct ColumnDef {
102 pub id: u16,
103 pub name: String,
104 pub ty: TypeId,
105 pub flags: ColumnFlags,
106}
107
108#[derive(Debug, Clone, Default, Serialize, Deserialize)]
110pub struct AlterColumn {
111 pub name: Option<String>,
112 pub ty: Option<TypeId>,
113 pub flags: Option<ColumnFlags>,
114}
115
116impl AlterColumn {
117 pub fn rename(name: impl Into<String>) -> Self {
118 Self {
119 name: Some(name.into()),
120 ty: None,
121 flags: None,
122 }
123 }
124
125 pub fn set_type(ty: TypeId) -> Self {
126 Self {
127 name: None,
128 ty: Some(ty),
129 flags: None,
130 }
131 }
132
133 pub fn set_flags(flags: ColumnFlags) -> Self {
134 Self {
135 name: None,
136 ty: None,
137 flags: Some(flags),
138 }
139 }
140}
141
142#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
145pub enum IndexKind {
146 Bitmap,
148 FmIndex,
150 Ann,
152 LearnedRange,
154 MinHash,
156 Sparse,
158}
159
160#[derive(Debug, Clone, Serialize, Deserialize)]
161pub struct IndexDef {
162 pub name: String,
163 pub column_id: u16,
164 pub kind: IndexKind,
165}
166
167#[derive(Debug, Clone, Default, Serialize, Deserialize)]
168pub struct Schema {
169 pub schema_id: u64,
170 pub columns: Vec<ColumnDef>,
171 pub indexes: Vec<IndexDef>,
172 #[serde(default)]
177 pub colocation: Vec<Vec<u16>>,
178 #[serde(default)]
183 pub constraints: TableConstraints,
184}
185
186impl Schema {
187 pub fn column(&self, name: &str) -> Option<&ColumnDef> {
188 self.columns.iter().find(|c| c.name == name)
189 }
190
191 pub fn primary_key(&self) -> Option<&ColumnDef> {
192 self.columns
193 .iter()
194 .find(|c| c.flags.contains(ColumnFlags::PRIMARY_KEY))
195 }
196
197 pub fn validate_not_null(&self, columns: &HashMap<u16, Value>) -> Result<()> {
202 for col in &self.columns {
203 if col.flags.contains(ColumnFlags::NULLABLE) {
204 continue;
205 }
206 if col.flags.contains(ColumnFlags::AUTO_INCREMENT) {
209 match columns.get(&col.id) {
210 None | Some(Value::Null) => continue,
211 Some(_) => {}
212 }
213 }
214 match columns.get(&col.id) {
215 None => {
216 return Err(MongrelError::InvalidArgument(format!(
217 "column '{}' ({}) is NOT NULL but was omitted",
218 col.name, col.id
219 )));
220 }
221 Some(Value::Null) => {
222 return Err(MongrelError::InvalidArgument(format!(
223 "column '{}' ({}) is NOT NULL but got NULL",
224 col.name, col.id
225 )));
226 }
227 Some(_) => {}
228 }
229 }
230 Ok(())
231 }
232
233 pub fn validate_auto_increment(&self) -> Result<()> {
237 let mut seen: Option<&ColumnDef> = None;
238 for col in &self.columns {
239 if !col.flags.contains(ColumnFlags::AUTO_INCREMENT) {
240 continue;
241 }
242 if let Some(prev) = seen {
243 return Err(MongrelError::Schema(format!(
244 "AUTO_INCREMENT may be set on at most one column; '{}' and '{}' both carry it",
245 prev.name, col.name
246 )));
247 }
248 if col.ty != TypeId::Int64 {
249 return Err(MongrelError::Schema(format!(
250 "AUTO_INCREMENT column '{}' must be Int64, is {:?}",
251 col.name, col.ty
252 )));
253 }
254 if !col.flags.contains(ColumnFlags::PRIMARY_KEY) {
255 return Err(MongrelError::Schema(format!(
256 "AUTO_INCREMENT column '{}' must also be the primary key",
257 col.name
258 )));
259 }
260 if col.flags.contains(ColumnFlags::NULLABLE) {
261 return Err(MongrelError::Schema(format!(
262 "AUTO_INCREMENT column '{}' must not be nullable",
263 col.name
264 )));
265 }
266 seen = Some(col);
267 }
268 Ok(())
269 }
270
271 pub fn auto_increment_column(&self) -> Option<&ColumnDef> {
273 self.columns
274 .iter()
275 .find(|c| c.flags.contains(ColumnFlags::AUTO_INCREMENT))
276 }
277}
278
279#[cfg(test)]
280mod tests {
281 use super::*;
282
283 #[test]
284 fn flag_composition() {
285 let f = ColumnFlags::empty()
286 .with(ColumnFlags::PRIMARY_KEY)
287 .with(ColumnFlags::ENCRYPTED_INDEXABLE);
288 assert!(f.contains(ColumnFlags::PRIMARY_KEY));
289 assert!(f.contains(ColumnFlags::ENCRYPTED_INDEXABLE));
290 assert!(!f.contains(ColumnFlags::ENCRYPTED));
291 }
292
293 #[test]
294 fn fixed_size() {
295 assert_eq!(TypeId::Int64.fixed_size(), Some(8));
296 assert_eq!(TypeId::Bytes.fixed_size(), None);
297 assert_eq!(TypeId::Embedding { dim: 768 }.fixed_size(), None);
298 }
299
300 fn col(id: u16, name: &str, ty: TypeId, flags: ColumnFlags) -> ColumnDef {
301 ColumnDef {
302 id,
303 name: name.into(),
304 ty,
305 flags,
306 }
307 }
308
309 #[test]
310 fn auto_increment_validation_accepts_int64_pk() {
311 let s = Schema {
312 schema_id: 1,
313 columns: vec![col(
314 0,
315 "id",
316 TypeId::Int64,
317 ColumnFlags::empty().with(ColumnFlags::PRIMARY_KEY | ColumnFlags::AUTO_INCREMENT),
318 )],
319 indexes: vec![],
320 colocation: vec![],
321 constraints: Default::default(),
322 };
323 assert!(s.validate_auto_increment().is_ok());
324 assert_eq!(s.auto_increment_column().unwrap().id, 0);
325 }
326
327 #[test]
328 fn auto_increment_validation_rejects_non_pk() {
329 let s = Schema {
330 schema_id: 1,
331 columns: vec![
332 col(
333 0,
334 "id",
335 TypeId::Int64,
336 ColumnFlags::empty().with(ColumnFlags::PRIMARY_KEY),
337 ),
338 col(
339 1,
340 "seq",
341 TypeId::Int64,
342 ColumnFlags::empty().with(ColumnFlags::AUTO_INCREMENT),
343 ),
344 ],
345 indexes: vec![],
346 colocation: vec![],
347 constraints: Default::default(),
348 };
349 assert!(s.validate_auto_increment().is_err());
350 }
351
352 #[test]
353 fn auto_increment_validation_rejects_non_int64() {
354 let s = Schema {
355 schema_id: 1,
356 columns: vec![col(
357 0,
358 "id",
359 TypeId::Bytes,
360 ColumnFlags::empty().with(ColumnFlags::PRIMARY_KEY | ColumnFlags::AUTO_INCREMENT),
361 )],
362 indexes: vec![],
363 colocation: vec![],
364 constraints: Default::default(),
365 };
366 assert!(s.validate_auto_increment().is_err());
367 }
368
369 #[test]
370 fn auto_increment_validation_rejects_two() {
371 let s = Schema {
372 schema_id: 1,
373 columns: vec![
374 col(
375 0,
376 "id",
377 TypeId::Int64,
378 ColumnFlags::empty()
379 .with(ColumnFlags::PRIMARY_KEY | ColumnFlags::AUTO_INCREMENT),
380 ),
381 col(
382 1,
383 "id2",
384 TypeId::Int64,
385 ColumnFlags::empty().with(ColumnFlags::AUTO_INCREMENT),
386 ),
387 ],
388 indexes: vec![],
389 colocation: vec![],
390 constraints: Default::default(),
391 };
392 assert!(s.validate_auto_increment().is_err());
393 }
394
395 #[test]
396 fn auto_increment_exempt_from_not_null_when_omitted() {
397 let s = Schema {
398 schema_id: 1,
399 columns: vec![
400 col(
401 0,
402 "id",
403 TypeId::Int64,
404 ColumnFlags::empty()
405 .with(ColumnFlags::PRIMARY_KEY | ColumnFlags::AUTO_INCREMENT),
406 ),
407 col(1, "name", TypeId::Bytes, ColumnFlags::empty()),
408 ],
409 indexes: vec![],
410 colocation: vec![],
411 constraints: Default::default(),
412 };
413 let mut cols = HashMap::new();
415 cols.insert(1u16, Value::Bytes(b"x".to_vec()));
416 assert!(s.validate_not_null(&cols).is_ok());
417 }
418}