1use super::types::DataType;
7use std::collections::HashMap;
8use std::fmt;
9
10#[derive(Debug, Clone)]
12pub struct TableDef {
13 pub name: String,
15 pub columns: Vec<ColumnDef>,
17 pub primary_key: Vec<String>,
19 pub indexes: Vec<IndexDef>,
21 pub constraints: Vec<Constraint>,
23 pub version: u32,
25 pub created_at: u64,
27 pub updated_at: u64,
29}
30
31impl TableDef {
32 pub fn new(name: impl Into<String>) -> Self {
34 let now = std::time::SystemTime::now()
35 .duration_since(std::time::UNIX_EPOCH)
36 .unwrap_or_default()
37 .as_secs();
38
39 Self {
40 name: name.into(),
41 columns: Vec::new(),
42 primary_key: Vec::new(),
43 indexes: Vec::new(),
44 constraints: Vec::new(),
45 version: 1,
46 created_at: now,
47 updated_at: now,
48 }
49 }
50
51 pub fn add_column(mut self, column: ColumnDef) -> Self {
53 self.columns.push(column);
54 self
55 }
56
57 pub fn primary_key(mut self, columns: Vec<String>) -> Self {
59 self.primary_key = columns;
60 self
61 }
62
63 pub fn add_index(mut self, index: IndexDef) -> Self {
65 self.indexes.push(index);
66 self
67 }
68
69 pub fn add_constraint(mut self, constraint: Constraint) -> Self {
71 self.constraints.push(constraint);
72 self
73 }
74
75 pub fn get_column(&self, name: &str) -> Option<&ColumnDef> {
77 self.columns.iter().find(|c| c.name == name)
78 }
79
80 pub fn column_index(&self, name: &str) -> Option<usize> {
82 self.columns.iter().position(|c| c.name == name)
83 }
84
85 pub fn is_primary_key_column(&self, name: &str) -> bool {
87 self.primary_key.iter().any(|pk| pk == name)
88 }
89
90 pub fn validate(&self) -> Result<(), TableDefError> {
92 if self.name.is_empty() {
94 return Err(TableDefError::EmptyTableName);
95 }
96
97 let mut seen = HashMap::new();
99 for col in &self.columns {
100 if seen.insert(&col.name, true).is_some() {
101 return Err(TableDefError::DuplicateColumn(col.name.clone()));
102 }
103 }
104
105 for pk_col in &self.primary_key {
107 if self.get_column(pk_col).is_none() {
108 return Err(TableDefError::InvalidPrimaryKey(pk_col.clone()));
109 }
110 }
111
112 for index in &self.indexes {
114 for col in &index.columns {
115 if self.get_column(col).is_none() {
116 return Err(TableDefError::InvalidIndexColumn(col.clone()));
117 }
118 }
119 }
120
121 for constraint in &self.constraints {
123 for col in &constraint.columns {
124 if self.get_column(col).is_none() {
125 return Err(TableDefError::InvalidConstraintColumn(col.clone()));
126 }
127 }
128 }
129
130 Ok(())
131 }
132
133 pub fn to_bytes(&self) -> Vec<u8> {
140 let columns = self
141 .columns
142 .iter()
143 .map(|col| reddb_file::ColumnLayout {
144 name: col.name.clone(),
145 data_type: col.data_type.to_byte(),
146 nullable: col.nullable,
147 default: col.default.clone(),
148 vector_dim: col.vector_dim,
149 compress: col.compress,
150 enum_variants: col.enum_variants.clone(),
151 decimal_precision: col.decimal_precision,
152 element_type: col.element_type.map(|et| et.to_byte()),
153 metadata: col
154 .metadata
155 .iter()
156 .map(|(k, v)| (k.clone(), v.clone()))
157 .collect(),
158 })
159 .collect();
160 let indexes = self
161 .indexes
162 .iter()
163 .map(|idx| reddb_file::IndexLayout {
164 name: idx.name.clone(),
165 index_type: idx.index_type as u8,
166 unique: idx.unique,
167 columns: idx.columns.clone(),
168 })
169 .collect();
170 let constraints = self
171 .constraints
172 .iter()
173 .map(|c| reddb_file::ConstraintLayout {
174 name: c.name.clone(),
175 constraint_type: c.constraint_type as u8,
176 columns: c.columns.clone(),
177 ref_table: c.ref_table.clone(),
178 ref_columns: c.ref_columns.clone(),
179 })
180 .collect();
181 let layout = reddb_file::TableDefLayout {
182 version: self.version,
183 name: self.name.clone(),
184 created_at: self.created_at,
185 updated_at: self.updated_at,
186 columns,
187 primary_key: self.primary_key.clone(),
188 indexes,
189 constraints,
190 };
191 reddb_file::encode_table_def(&layout)
192 }
193
194 pub fn from_bytes(data: &[u8]) -> Result<Self, TableDefError> {
196 let layout = reddb_file::decode_table_def(data).map_err(TableDefError::from_codec)?;
197
198 let mut columns = Vec::with_capacity(layout.columns.len());
199 for col in layout.columns {
200 let data_type =
201 DataType::from_byte(col.data_type).ok_or(TableDefError::InvalidDataType)?;
202 let element_type = match col.element_type {
203 Some(byte) => {
204 Some(DataType::from_byte(byte).ok_or(TableDefError::InvalidDataType)?)
205 }
206 None => None,
207 };
208 columns.push(ColumnDef {
209 name: col.name,
210 data_type,
211 nullable: col.nullable,
212 default: col.default,
213 vector_dim: col.vector_dim,
214 compress: col.compress,
215 enum_variants: col.enum_variants,
216 decimal_precision: col.decimal_precision,
217 element_type,
218 metadata: col.metadata.into_iter().collect(),
219 });
220 }
221
222 let mut indexes = Vec::with_capacity(layout.indexes.len());
223 for idx in layout.indexes {
224 let index_type =
225 IndexType::from_byte(idx.index_type).ok_or(TableDefError::InvalidIndexType)?;
226 indexes.push(IndexDef {
227 name: idx.name,
228 columns: idx.columns,
229 index_type,
230 unique: idx.unique,
231 });
232 }
233
234 let mut constraints = Vec::with_capacity(layout.constraints.len());
235 for c in layout.constraints {
236 let constraint_type = ConstraintType::from_byte(c.constraint_type)
237 .ok_or(TableDefError::InvalidConstraintType)?;
238 constraints.push(Constraint {
239 name: c.name,
240 constraint_type,
241 columns: c.columns,
242 ref_table: c.ref_table,
243 ref_columns: c.ref_columns,
244 });
245 }
246
247 Ok(Self {
248 name: layout.name,
249 columns,
250 primary_key: layout.primary_key,
251 indexes,
252 constraints,
253 version: layout.version,
254 created_at: layout.created_at,
255 updated_at: layout.updated_at,
256 })
257 }
258}
259
260impl fmt::Display for TableDef {
261 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
262 writeln!(f, "TABLE {} (version {})", self.name, self.version)?;
263 writeln!(f, " Columns:")?;
264 for col in &self.columns {
265 writeln!(f, " {}", col)?;
266 }
267 if !self.primary_key.is_empty() {
268 writeln!(f, " Primary Key: ({})", self.primary_key.join(", "))?;
269 }
270 if !self.indexes.is_empty() {
271 writeln!(f, " Indexes:")?;
272 for idx in &self.indexes {
273 writeln!(f, " {}", idx)?;
274 }
275 }
276 Ok(())
277 }
278}
279
280#[derive(Debug, Clone)]
282pub struct ColumnDef {
283 pub name: String,
285 pub data_type: DataType,
287 pub nullable: bool,
289 pub default: Option<Vec<u8>>,
291 pub vector_dim: Option<u32>,
293 pub compress: bool,
295 pub enum_variants: Vec<String>,
297 pub decimal_precision: u8,
299 pub element_type: Option<DataType>,
301 pub metadata: HashMap<String, String>,
303}
304
305impl ColumnDef {
306 pub fn new(name: impl Into<String>, data_type: DataType) -> Self {
308 Self {
309 name: name.into(),
310 data_type,
311 nullable: true,
312 default: None,
313 vector_dim: None,
314 compress: false,
315 enum_variants: Vec::new(),
316 decimal_precision: 4,
317 element_type: None,
318 metadata: HashMap::new(),
319 }
320 }
321
322 pub fn not_null(mut self) -> Self {
324 self.nullable = false;
325 self
326 }
327
328 pub fn with_default(mut self, default: Vec<u8>) -> Self {
330 self.default = Some(default);
331 self
332 }
333
334 pub fn with_vector_dim(mut self, dim: u32) -> Self {
336 self.vector_dim = Some(dim);
337 self
338 }
339
340 pub fn with_metadata(mut self, key: impl Into<String>, value: impl Into<String>) -> Self {
342 self.metadata.insert(key.into(), value.into());
343 self
344 }
345
346 pub fn compressed(mut self) -> Self {
348 self.compress = true;
349 self
350 }
351
352 pub fn with_variants(mut self, variants: Vec<String>) -> Self {
354 self.enum_variants = variants;
355 self
356 }
357
358 pub fn with_precision(mut self, precision: u8) -> Self {
360 self.decimal_precision = precision;
361 self
362 }
363
364 pub fn with_element_type(mut self, dt: DataType) -> Self {
366 self.element_type = Some(dt);
367 self
368 }
369}
370
371impl fmt::Display for ColumnDef {
372 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
373 write!(f, "{} {}", self.name, self.data_type)?;
374 if let Some(dim) = self.vector_dim {
375 write!(f, "({})", dim)?;
376 }
377 if !self.nullable {
378 write!(f, " NOT NULL")?;
379 }
380 if self.default.is_some() {
381 write!(f, " DEFAULT <value>")?;
382 }
383 Ok(())
384 }
385}
386
387#[derive(Debug, Clone)]
389pub struct IndexDef {
390 pub name: String,
392 pub columns: Vec<String>,
394 pub index_type: IndexType,
396 pub unique: bool,
398}
399
400impl IndexDef {
401 pub fn new(name: impl Into<String>, columns: Vec<String>) -> Self {
403 Self {
404 name: name.into(),
405 columns,
406 index_type: IndexType::BTree,
407 unique: false,
408 }
409 }
410
411 pub fn unique(mut self) -> Self {
413 self.unique = true;
414 self
415 }
416
417 pub fn with_type(mut self, index_type: IndexType) -> Self {
419 self.index_type = index_type;
420 self
421 }
422}
423
424impl fmt::Display for IndexDef {
425 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
426 if self.unique {
427 write!(f, "UNIQUE ")?;
428 }
429 write!(
430 f,
431 "INDEX {} ({}) USING {:?}",
432 self.name,
433 self.columns.join(", "),
434 self.index_type
435 )
436 }
437}
438
439#[derive(Debug, Clone, Copy, PartialEq, Eq)]
441#[repr(u8)]
442pub enum IndexType {
443 BTree = 1,
445 Hash = 2,
447 IvfFlat = 3,
449 Hnsw = 4,
451}
452
453impl IndexType {
454 fn from_byte(b: u8) -> Option<Self> {
455 match b {
456 1 => Some(IndexType::BTree),
457 2 => Some(IndexType::Hash),
458 3 => Some(IndexType::IvfFlat),
459 4 => Some(IndexType::Hnsw),
460 _ => None,
461 }
462 }
463}
464
465#[derive(Debug, Clone)]
467pub struct Constraint {
468 pub name: String,
470 pub constraint_type: ConstraintType,
472 pub columns: Vec<String>,
474 pub ref_table: Option<String>,
476 pub ref_columns: Option<Vec<String>>,
478}
479
480impl Constraint {
481 pub fn new(name: impl Into<String>, constraint_type: ConstraintType) -> Self {
483 Self {
484 name: name.into(),
485 constraint_type,
486 columns: Vec::new(),
487 ref_table: None,
488 ref_columns: None,
489 }
490 }
491
492 pub fn on_columns(mut self, columns: Vec<String>) -> Self {
494 self.columns = columns;
495 self
496 }
497
498 pub fn references(mut self, table: String, columns: Vec<String>) -> Self {
500 self.ref_table = Some(table);
501 self.ref_columns = Some(columns);
502 self
503 }
504}
505
506#[derive(Debug, Clone, Copy, PartialEq, Eq)]
508#[repr(u8)]
509pub enum ConstraintType {
510 PrimaryKey = 1,
512 Unique = 2,
514 ForeignKey = 3,
516 Check = 4,
518 NotNull = 5,
520}
521
522impl ConstraintType {
523 fn from_byte(b: u8) -> Option<Self> {
524 match b {
525 1 => Some(ConstraintType::PrimaryKey),
526 2 => Some(ConstraintType::Unique),
527 3 => Some(ConstraintType::ForeignKey),
528 4 => Some(ConstraintType::Check),
529 5 => Some(ConstraintType::NotNull),
530 _ => None,
531 }
532 }
533}
534
535#[derive(Debug, Clone, PartialEq)]
537pub enum TableDefError {
538 EmptyTableName,
540 DuplicateColumn(String),
542 InvalidPrimaryKey(String),
544 InvalidIndexColumn(String),
546 InvalidConstraintColumn(String),
548 TruncatedData,
550 InvalidMagic,
552 InvalidDataType,
554 InvalidIndexType,
556 InvalidConstraintType,
558 VarintOverflow,
560}
561
562impl fmt::Display for TableDefError {
563 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
564 match self {
565 TableDefError::EmptyTableName => write!(f, "empty table name"),
566 TableDefError::DuplicateColumn(name) => write!(f, "duplicate column: {}", name),
567 TableDefError::InvalidPrimaryKey(name) => {
568 write!(f, "invalid primary key column: {}", name)
569 }
570 TableDefError::InvalidIndexColumn(name) => write!(f, "invalid index column: {}", name),
571 TableDefError::InvalidConstraintColumn(name) => {
572 write!(f, "invalid constraint column: {}", name)
573 }
574 TableDefError::TruncatedData => write!(f, "truncated data"),
575 TableDefError::InvalidMagic => write!(f, "invalid magic bytes"),
576 TableDefError::InvalidDataType => write!(f, "invalid data type"),
577 TableDefError::InvalidIndexType => write!(f, "invalid index type"),
578 TableDefError::InvalidConstraintType => write!(f, "invalid constraint type"),
579 TableDefError::VarintOverflow => write!(f, "varint overflow"),
580 }
581 }
582}
583
584impl std::error::Error for TableDefError {}
585
586impl TableDefError {
587 fn from_codec(err: reddb_file::TableDefCodecError) -> Self {
591 match err {
592 reddb_file::TableDefCodecError::TruncatedData => TableDefError::TruncatedData,
593 reddb_file::TableDefCodecError::InvalidMagic => TableDefError::InvalidMagic,
594 reddb_file::TableDefCodecError::VarintOverflow => TableDefError::VarintOverflow,
595 reddb_file::TableDefCodecError::InvalidUtf8 => TableDefError::TruncatedData,
596 }
597 }
598}
599
600#[cfg(test)]
601mod tests {
602 use super::*;
603
604 #[test]
605 fn test_table_def_basic() {
606 let table = TableDef::new("port_scans")
607 .add_column(ColumnDef::new("id", DataType::UnsignedInteger).not_null())
608 .add_column(ColumnDef::new("ip", DataType::IpAddr).not_null())
609 .add_column(ColumnDef::new("port", DataType::UnsignedInteger).not_null())
610 .add_column(ColumnDef::new("status", DataType::Text))
611 .add_column(ColumnDef::new("timestamp", DataType::Timestamp).not_null())
612 .primary_key(vec!["id".to_string()]);
613
614 assert_eq!(table.name, "port_scans");
615 assert_eq!(table.columns.len(), 5);
616 assert_eq!(table.primary_key, vec!["id"]);
617 assert!(table.validate().is_ok());
618 }
619
620 #[test]
621 fn test_table_def_with_indexes() {
622 let table = TableDef::new("subdomains")
623 .add_column(ColumnDef::new("id", DataType::UnsignedInteger).not_null())
624 .add_column(ColumnDef::new("domain", DataType::Text).not_null())
625 .add_column(ColumnDef::new("subdomain", DataType::Text).not_null())
626 .add_column(ColumnDef::new("ip", DataType::IpAddr))
627 .primary_key(vec!["id".to_string()])
628 .add_index(IndexDef::new("idx_domain", vec!["domain".to_string()]))
629 .add_index(IndexDef::new("idx_subdomain", vec!["subdomain".to_string()]).unique());
630
631 assert_eq!(table.indexes.len(), 2);
632 assert!(table.indexes[1].unique);
633 assert!(table.validate().is_ok());
634 }
635
636 #[test]
637 fn test_table_def_with_vector() {
638 let table = TableDef::new("embeddings")
639 .add_column(ColumnDef::new("id", DataType::UnsignedInteger).not_null())
640 .add_column(
641 ColumnDef::new("embedding", DataType::Vector)
642 .not_null()
643 .with_vector_dim(384),
644 )
645 .add_column(ColumnDef::new("text", DataType::Text))
646 .primary_key(vec!["id".to_string()])
647 .add_index(
648 IndexDef::new("idx_embedding", vec!["embedding".to_string()])
649 .with_type(IndexType::IvfFlat),
650 );
651
652 let col = table.get_column("embedding").unwrap();
653 assert_eq!(col.vector_dim, Some(384));
654 assert!(table.validate().is_ok());
655 }
656
657 #[test]
658 fn test_table_def_validation_duplicate_column() {
659 let table = TableDef::new("test")
660 .add_column(ColumnDef::new("id", DataType::Integer))
661 .add_column(ColumnDef::new("id", DataType::Text)); assert!(matches!(
664 table.validate(),
665 Err(TableDefError::DuplicateColumn(_))
666 ));
667 }
668
669 #[test]
670 fn test_table_def_validation_invalid_pk() {
671 let table = TableDef::new("test")
672 .add_column(ColumnDef::new("id", DataType::Integer))
673 .primary_key(vec!["nonexistent".to_string()]);
674
675 assert!(matches!(
676 table.validate(),
677 Err(TableDefError::InvalidPrimaryKey(_))
678 ));
679 }
680
681 #[test]
682 fn test_table_def_roundtrip() {
683 let table = TableDef::new("hosts")
684 .add_column(ColumnDef::new("id", DataType::UnsignedInteger).not_null())
685 .add_column(ColumnDef::new("ip", DataType::IpAddr).not_null())
686 .add_column(ColumnDef::new("hostname", DataType::Text))
687 .add_column(ColumnDef::new("last_seen", DataType::Timestamp))
688 .add_column(ColumnDef::new("fingerprint", DataType::Vector).with_vector_dim(128))
689 .primary_key(vec!["id".to_string()])
690 .add_index(IndexDef::new("idx_ip", vec!["ip".to_string()]).unique())
691 .add_index(
692 IndexDef::new("idx_fingerprint", vec!["fingerprint".to_string()])
693 .with_type(IndexType::IvfFlat),
694 );
695
696 let bytes = table.to_bytes();
697 let recovered = TableDef::from_bytes(&bytes).unwrap();
698
699 assert_eq!(table.name, recovered.name);
700 assert_eq!(table.columns.len(), recovered.columns.len());
701 assert_eq!(table.primary_key, recovered.primary_key);
702 assert_eq!(table.indexes.len(), recovered.indexes.len());
703
704 for (orig, rec) in table.columns.iter().zip(recovered.columns.iter()) {
705 assert_eq!(orig.name, rec.name);
706 assert_eq!(orig.data_type, rec.data_type);
707 assert_eq!(orig.nullable, rec.nullable);
708 assert_eq!(orig.vector_dim, rec.vector_dim);
709 }
710
711 for (orig, rec) in table.indexes.iter().zip(recovered.indexes.iter()) {
712 assert_eq!(orig.name, rec.name);
713 assert_eq!(orig.columns, rec.columns);
714 assert_eq!(orig.unique, rec.unique);
715 assert_eq!(orig.index_type, rec.index_type);
716 }
717 }
718
719 #[test]
720 fn test_column_def_metadata() {
721 let col = ColumnDef::new("ip", DataType::IpAddr)
722 .not_null()
723 .with_metadata("description", "Target IP address")
724 .with_metadata("indexed", "true");
725
726 assert_eq!(
727 col.metadata.get("description"),
728 Some(&"Target IP address".to_string())
729 );
730 assert_eq!(col.metadata.get("indexed"), Some(&"true".to_string()));
731 }
732
733 #[test]
734 fn test_constraint_foreign_key() {
735 let constraint = Constraint::new("fk_host", ConstraintType::ForeignKey)
736 .on_columns(vec!["host_id".to_string()])
737 .references("hosts".to_string(), vec!["id".to_string()]);
738
739 assert_eq!(constraint.constraint_type, ConstraintType::ForeignKey);
740 assert_eq!(constraint.columns, vec!["host_id"]);
741 assert_eq!(constraint.ref_table, Some("hosts".to_string()));
742 assert_eq!(constraint.ref_columns, Some(vec!["id".to_string()]));
743 }
744
745 #[test]
746 fn test_table_display() {
747 let table = TableDef::new("test")
748 .add_column(ColumnDef::new("id", DataType::Integer).not_null())
749 .add_column(ColumnDef::new("name", DataType::Text))
750 .primary_key(vec!["id".to_string()]);
751
752 let display = format!("{}", table);
753 assert!(display.contains("TABLE test"));
754 assert!(display.contains("id INTEGER NOT NULL"));
755 assert!(display.contains("name TEXT"));
756 assert!(display.contains("Primary Key: (id)"));
757 }
758}