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