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::ColumnDefFrame {
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::IndexDefFrame {
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::ConstraintFrame {
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 frame = reddb_file::TableDefFrame {
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_frame(&frame)
192 }
193
194 pub fn from_bytes(data: &[u8]) -> Result<Self, TableDefError> {
196 let frame =
197 reddb_file::decode_table_def_frame(data).map_err(TableDefError::from_frame_codec)?;
198
199 let mut columns = Vec::with_capacity(frame.columns.len());
200 for col in frame.columns {
201 let data_type =
202 DataType::from_byte(col.data_type).ok_or(TableDefError::InvalidDataType)?;
203 let element_type = match col.element_type {
204 Some(byte) => {
205 Some(DataType::from_byte(byte).ok_or(TableDefError::InvalidDataType)?)
206 }
207 None => None,
208 };
209 columns.push(ColumnDef {
210 name: col.name,
211 data_type,
212 nullable: col.nullable,
213 default: col.default,
214 vector_dim: col.vector_dim,
215 compress: col.compress,
216 enum_variants: col.enum_variants,
217 decimal_precision: col.decimal_precision,
218 element_type,
219 metadata: col.metadata.into_iter().collect(),
220 });
221 }
222
223 let mut indexes = Vec::with_capacity(frame.indexes.len());
224 for idx in frame.indexes {
225 let index_type =
226 IndexType::from_byte(idx.index_type).ok_or(TableDefError::InvalidIndexType)?;
227 indexes.push(IndexDef {
228 name: idx.name,
229 columns: idx.columns,
230 index_type,
231 unique: idx.unique,
232 });
233 }
234
235 let mut constraints = Vec::with_capacity(frame.constraints.len());
236 for c in frame.constraints {
237 let constraint_type = ConstraintType::from_byte(c.constraint_type)
238 .ok_or(TableDefError::InvalidConstraintType)?;
239 constraints.push(Constraint {
240 name: c.name,
241 constraint_type,
242 columns: c.columns,
243 ref_table: c.ref_table,
244 ref_columns: c.ref_columns,
245 });
246 }
247
248 Ok(Self {
249 name: frame.name,
250 columns,
251 primary_key: frame.primary_key,
252 indexes,
253 constraints,
254 version: frame.version,
255 created_at: frame.created_at,
256 updated_at: frame.updated_at,
257 })
258 }
259}
260
261impl fmt::Display for TableDef {
262 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
263 writeln!(f, "TABLE {} (version {})", self.name, self.version)?;
264 writeln!(f, " Columns:")?;
265 for col in &self.columns {
266 writeln!(f, " {}", col)?;
267 }
268 if !self.primary_key.is_empty() {
269 writeln!(f, " Primary Key: ({})", self.primary_key.join(", "))?;
270 }
271 if !self.indexes.is_empty() {
272 writeln!(f, " Indexes:")?;
273 for idx in &self.indexes {
274 writeln!(f, " {}", idx)?;
275 }
276 }
277 Ok(())
278 }
279}
280
281#[derive(Debug, Clone)]
283pub struct ColumnDef {
284 pub name: String,
286 pub data_type: DataType,
288 pub nullable: bool,
290 pub default: Option<Vec<u8>>,
292 pub vector_dim: Option<u32>,
294 pub compress: bool,
296 pub enum_variants: Vec<String>,
298 pub decimal_precision: u8,
300 pub element_type: Option<DataType>,
302 pub metadata: HashMap<String, String>,
304}
305
306impl ColumnDef {
307 pub fn new(name: impl Into<String>, data_type: DataType) -> Self {
309 Self {
310 name: name.into(),
311 data_type,
312 nullable: true,
313 default: None,
314 vector_dim: None,
315 compress: false,
316 enum_variants: Vec::new(),
317 decimal_precision: 4,
318 element_type: None,
319 metadata: HashMap::new(),
320 }
321 }
322
323 pub fn not_null(mut self) -> Self {
325 self.nullable = false;
326 self
327 }
328
329 pub fn with_default(mut self, default: Vec<u8>) -> Self {
331 self.default = Some(default);
332 self
333 }
334
335 pub fn with_vector_dim(mut self, dim: u32) -> Self {
337 self.vector_dim = Some(dim);
338 self
339 }
340
341 pub fn with_metadata(mut self, key: impl Into<String>, value: impl Into<String>) -> Self {
343 self.metadata.insert(key.into(), value.into());
344 self
345 }
346
347 pub fn compressed(mut self) -> Self {
349 self.compress = true;
350 self
351 }
352
353 pub fn with_variants(mut self, variants: Vec<String>) -> Self {
355 self.enum_variants = variants;
356 self
357 }
358
359 pub fn with_precision(mut self, precision: u8) -> Self {
361 self.decimal_precision = precision;
362 self
363 }
364
365 pub fn with_element_type(mut self, dt: DataType) -> Self {
367 self.element_type = Some(dt);
368 self
369 }
370}
371
372impl fmt::Display for ColumnDef {
373 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
374 write!(f, "{} {}", self.name, self.data_type)?;
375 if let Some(dim) = self.vector_dim {
376 write!(f, "({})", dim)?;
377 }
378 if !self.nullable {
379 write!(f, " NOT NULL")?;
380 }
381 if self.default.is_some() {
382 write!(f, " DEFAULT <value>")?;
383 }
384 Ok(())
385 }
386}
387
388#[derive(Debug, Clone)]
390pub struct IndexDef {
391 pub name: String,
393 pub columns: Vec<String>,
395 pub index_type: IndexType,
397 pub unique: bool,
399}
400
401impl IndexDef {
402 pub fn new(name: impl Into<String>, columns: Vec<String>) -> Self {
404 Self {
405 name: name.into(),
406 columns,
407 index_type: IndexType::BTree,
408 unique: false,
409 }
410 }
411
412 pub fn unique(mut self) -> Self {
414 self.unique = true;
415 self
416 }
417
418 pub fn with_type(mut self, index_type: IndexType) -> Self {
420 self.index_type = index_type;
421 self
422 }
423}
424
425impl fmt::Display for IndexDef {
426 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
427 if self.unique {
428 write!(f, "UNIQUE ")?;
429 }
430 write!(
431 f,
432 "INDEX {} ({}) USING {:?}",
433 self.name,
434 self.columns.join(", "),
435 self.index_type
436 )
437 }
438}
439
440#[derive(Debug, Clone, Copy, PartialEq, Eq)]
442#[repr(u8)]
443pub enum IndexType {
444 BTree = 1,
446 Hash = 2,
448 IvfFlat = 3,
450 Hnsw = 4,
452}
453
454impl IndexType {
455 fn from_byte(b: u8) -> Option<Self> {
456 match b {
457 1 => Some(IndexType::BTree),
458 2 => Some(IndexType::Hash),
459 3 => Some(IndexType::IvfFlat),
460 4 => Some(IndexType::Hnsw),
461 _ => None,
462 }
463 }
464}
465
466#[derive(Debug, Clone)]
468pub struct Constraint {
469 pub name: String,
471 pub constraint_type: ConstraintType,
473 pub columns: Vec<String>,
475 pub ref_table: Option<String>,
477 pub ref_columns: Option<Vec<String>>,
479}
480
481impl Constraint {
482 pub fn new(name: impl Into<String>, constraint_type: ConstraintType) -> Self {
484 Self {
485 name: name.into(),
486 constraint_type,
487 columns: Vec::new(),
488 ref_table: None,
489 ref_columns: None,
490 }
491 }
492
493 pub fn on_columns(mut self, columns: Vec<String>) -> Self {
495 self.columns = columns;
496 self
497 }
498
499 pub fn references(mut self, table: String, columns: Vec<String>) -> Self {
501 self.ref_table = Some(table);
502 self.ref_columns = Some(columns);
503 self
504 }
505}
506
507#[derive(Debug, Clone, Copy, PartialEq, Eq)]
509#[repr(u8)]
510pub enum ConstraintType {
511 PrimaryKey = 1,
513 Unique = 2,
515 ForeignKey = 3,
517 Check = 4,
519 NotNull = 5,
521}
522
523impl ConstraintType {
524 fn from_byte(b: u8) -> Option<Self> {
525 match b {
526 1 => Some(ConstraintType::PrimaryKey),
527 2 => Some(ConstraintType::Unique),
528 3 => Some(ConstraintType::ForeignKey),
529 4 => Some(ConstraintType::Check),
530 5 => Some(ConstraintType::NotNull),
531 _ => None,
532 }
533 }
534}
535
536#[derive(Debug, Clone, PartialEq)]
538pub enum TableDefError {
539 EmptyTableName,
541 DuplicateColumn(String),
543 InvalidPrimaryKey(String),
545 InvalidIndexColumn(String),
547 InvalidConstraintColumn(String),
549 TruncatedData,
551 InvalidMagic,
553 InvalidDataType,
555 InvalidIndexType,
557 InvalidConstraintType,
559 VarintOverflow,
561}
562
563impl fmt::Display for TableDefError {
564 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
565 match self {
566 TableDefError::EmptyTableName => write!(f, "empty table name"),
567 TableDefError::DuplicateColumn(name) => write!(f, "duplicate column: {}", name),
568 TableDefError::InvalidPrimaryKey(name) => {
569 write!(f, "invalid primary key column: {}", name)
570 }
571 TableDefError::InvalidIndexColumn(name) => write!(f, "invalid index column: {}", name),
572 TableDefError::InvalidConstraintColumn(name) => {
573 write!(f, "invalid constraint column: {}", name)
574 }
575 TableDefError::TruncatedData => write!(f, "truncated data"),
576 TableDefError::InvalidMagic => write!(f, "invalid magic bytes"),
577 TableDefError::InvalidDataType => write!(f, "invalid data type"),
578 TableDefError::InvalidIndexType => write!(f, "invalid index type"),
579 TableDefError::InvalidConstraintType => write!(f, "invalid constraint type"),
580 TableDefError::VarintOverflow => write!(f, "varint overflow"),
581 }
582 }
583}
584
585impl std::error::Error for TableDefError {}
586
587impl TableDefError {
588 fn from_frame_codec(err: reddb_file::TableDefFrameError) -> Self {
592 match err {
593 reddb_file::TableDefFrameError::TruncatedData => TableDefError::TruncatedData,
594 reddb_file::TableDefFrameError::InvalidMagic => TableDefError::InvalidMagic,
595 reddb_file::TableDefFrameError::VarintOverflow => TableDefError::VarintOverflow,
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}