1use crate::DataType;
7use std::collections::HashMap;
8use std::fmt;
9
10#[derive(Debug, Clone, PartialEq, Eq)]
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::new_at(name, now)
40 }
41
42 pub fn new_at(name: impl Into<String>, now: u64) -> Self {
44 Self {
45 name: name.into(),
46 columns: Vec::new(),
47 primary_key: Vec::new(),
48 indexes: Vec::new(),
49 constraints: Vec::new(),
50 version: 1,
51 created_at: now,
52 updated_at: now,
53 }
54 }
55
56 pub fn add_column(mut self, column: ColumnDef) -> Self {
58 self.columns.push(column);
59 self
60 }
61
62 pub fn primary_key(mut self, columns: Vec<String>) -> Self {
64 self.primary_key = columns;
65 self
66 }
67
68 pub fn add_index(mut self, index: IndexDef) -> Self {
70 self.indexes.push(index);
71 self
72 }
73
74 pub fn add_constraint(mut self, constraint: Constraint) -> Self {
76 self.constraints.push(constraint);
77 self
78 }
79
80 pub fn get_column(&self, name: &str) -> Option<&ColumnDef> {
82 self.columns.iter().find(|column| column.name == name)
83 }
84
85 pub fn column_index(&self, name: &str) -> Option<usize> {
87 self.columns.iter().position(|column| column.name == name)
88 }
89
90 pub fn is_primary_key_column(&self, name: &str) -> bool {
92 self.primary_key.iter().any(|column| column == name)
93 }
94
95 pub fn validate(&self) -> Result<(), TableDefError> {
97 if self.name.is_empty() {
98 return Err(TableDefError::EmptyTableName);
99 }
100
101 let mut seen = HashMap::new();
102 for column in &self.columns {
103 if seen.insert(&column.name, true).is_some() {
104 return Err(TableDefError::DuplicateColumn(column.name.clone()));
105 }
106 }
107
108 for primary_key_column in &self.primary_key {
109 if self.get_column(primary_key_column).is_none() {
110 return Err(TableDefError::InvalidPrimaryKey(primary_key_column.clone()));
111 }
112 }
113
114 for index in &self.indexes {
115 for column in &index.columns {
116 if self.get_column(column).is_none() {
117 return Err(TableDefError::InvalidIndexColumn(column.clone()));
118 }
119 }
120 }
121
122 for constraint in &self.constraints {
123 for column in &constraint.columns {
124 if self.get_column(column).is_none() {
125 return Err(TableDefError::InvalidConstraintColumn(column.clone()));
126 }
127 }
128 }
129
130 Ok(())
131 }
132}
133
134impl fmt::Display for TableDef {
135 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
136 writeln!(f, "TABLE {} (version {})", self.name, self.version)?;
137 writeln!(f, " Columns:")?;
138 for column in &self.columns {
139 writeln!(f, " {column}")?;
140 }
141 if !self.primary_key.is_empty() {
142 writeln!(f, " Primary Key: ({})", self.primary_key.join(", "))?;
143 }
144 if !self.indexes.is_empty() {
145 writeln!(f, " Indexes:")?;
146 for index in &self.indexes {
147 writeln!(f, " {index}")?;
148 }
149 }
150 Ok(())
151 }
152}
153
154#[derive(Debug, Clone, PartialEq, Eq)]
156pub struct ColumnDef {
157 pub name: String,
159 pub data_type: DataType,
161 pub nullable: bool,
163 pub default: Option<Vec<u8>>,
165 pub vector_dim: Option<u32>,
167 pub compress: bool,
169 pub enum_variants: Vec<String>,
171 pub decimal_precision: u8,
173 pub element_type: Option<DataType>,
175 pub metadata: HashMap<String, String>,
177}
178
179impl ColumnDef {
180 pub fn new(name: impl Into<String>, data_type: DataType) -> Self {
182 Self {
183 name: name.into(),
184 data_type,
185 nullable: true,
186 default: None,
187 vector_dim: None,
188 compress: false,
189 enum_variants: Vec::new(),
190 decimal_precision: 4,
191 element_type: None,
192 metadata: HashMap::new(),
193 }
194 }
195
196 pub fn not_null(mut self) -> Self {
198 self.nullable = false;
199 self
200 }
201
202 pub fn with_default(mut self, default: Vec<u8>) -> Self {
204 self.default = Some(default);
205 self
206 }
207
208 pub fn with_vector_dim(mut self, dimension: u32) -> Self {
210 self.vector_dim = Some(dimension);
211 self
212 }
213
214 pub fn with_metadata(mut self, key: impl Into<String>, value: impl Into<String>) -> Self {
216 self.metadata.insert(key.into(), value.into());
217 self
218 }
219
220 pub fn compressed(mut self) -> Self {
222 self.compress = true;
223 self
224 }
225
226 pub fn with_variants(mut self, variants: Vec<String>) -> Self {
228 self.enum_variants = variants;
229 self
230 }
231
232 pub fn with_precision(mut self, precision: u8) -> Self {
234 self.decimal_precision = precision;
235 self
236 }
237
238 pub fn with_element_type(mut self, data_type: DataType) -> Self {
240 self.element_type = Some(data_type);
241 self
242 }
243}
244
245impl fmt::Display for ColumnDef {
246 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
247 write!(f, "{} {}", self.name, self.data_type)?;
248 if let Some(dimension) = self.vector_dim {
249 write!(f, "({dimension})")?;
250 }
251 if !self.nullable {
252 write!(f, " NOT NULL")?;
253 }
254 if self.default.is_some() {
255 write!(f, " DEFAULT <value>")?;
256 }
257 Ok(())
258 }
259}
260
261#[derive(Debug, Clone, PartialEq, Eq)]
263pub struct IndexDef {
264 pub name: String,
265 pub columns: Vec<String>,
266 pub index_type: IndexType,
267 pub unique: bool,
268}
269
270impl IndexDef {
271 pub fn new(name: impl Into<String>, columns: Vec<String>) -> Self {
272 Self {
273 name: name.into(),
274 columns,
275 index_type: IndexType::BTree,
276 unique: false,
277 }
278 }
279
280 pub fn unique(mut self) -> Self {
281 self.unique = true;
282 self
283 }
284
285 pub fn with_type(mut self, index_type: IndexType) -> Self {
286 self.index_type = index_type;
287 self
288 }
289}
290
291impl fmt::Display for IndexDef {
292 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
293 if self.unique {
294 write!(f, "UNIQUE ")?;
295 }
296 write!(
297 f,
298 "INDEX {} ({}) USING {:?}",
299 self.name,
300 self.columns.join(", "),
301 self.index_type
302 )
303 }
304}
305
306#[derive(Debug, Clone, Copy, PartialEq, Eq)]
308#[repr(u8)]
309pub enum IndexType {
310 BTree = 1,
311 Hash = 2,
312 IvfFlat = 3,
313 Hnsw = 4,
314}
315
316impl IndexType {
317 pub fn from_byte(byte: u8) -> Option<Self> {
318 match byte {
319 1 => Some(Self::BTree),
320 2 => Some(Self::Hash),
321 3 => Some(Self::IvfFlat),
322 4 => Some(Self::Hnsw),
323 _ => None,
324 }
325 }
326}
327
328#[derive(Debug, Clone, PartialEq, Eq)]
330pub struct Constraint {
331 pub name: String,
332 pub constraint_type: ConstraintType,
333 pub columns: Vec<String>,
334 pub ref_table: Option<String>,
335 pub ref_columns: Option<Vec<String>>,
336}
337
338impl Constraint {
339 pub fn new(name: impl Into<String>, constraint_type: ConstraintType) -> Self {
340 Self {
341 name: name.into(),
342 constraint_type,
343 columns: Vec::new(),
344 ref_table: None,
345 ref_columns: None,
346 }
347 }
348
349 pub fn on_columns(mut self, columns: Vec<String>) -> Self {
350 self.columns = columns;
351 self
352 }
353
354 pub fn references(mut self, table: String, columns: Vec<String>) -> Self {
355 self.ref_table = Some(table);
356 self.ref_columns = Some(columns);
357 self
358 }
359}
360
361#[derive(Debug, Clone, Copy, PartialEq, Eq)]
363#[repr(u8)]
364pub enum ConstraintType {
365 PrimaryKey = 1,
366 Unique = 2,
367 ForeignKey = 3,
368 Check = 4,
369 NotNull = 5,
370}
371
372impl ConstraintType {
373 pub fn from_byte(byte: u8) -> Option<Self> {
374 match byte {
375 1 => Some(Self::PrimaryKey),
376 2 => Some(Self::Unique),
377 3 => Some(Self::ForeignKey),
378 4 => Some(Self::Check),
379 5 => Some(Self::NotNull),
380 _ => None,
381 }
382 }
383}
384
385#[derive(Debug, Clone, PartialEq, Eq)]
387pub enum TableDefError {
388 EmptyTableName,
389 DuplicateColumn(String),
390 InvalidPrimaryKey(String),
391 InvalidIndexColumn(String),
392 InvalidConstraintColumn(String),
393 TruncatedData,
394 InvalidMagic,
395 InvalidDataType,
396 InvalidIndexType,
397 InvalidConstraintType,
398 VarintOverflow,
399}
400
401impl fmt::Display for TableDefError {
402 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
403 match self {
404 Self::EmptyTableName => write!(f, "empty table name"),
405 Self::DuplicateColumn(name) => write!(f, "duplicate column: {name}"),
406 Self::InvalidPrimaryKey(name) => write!(f, "invalid primary key column: {name}"),
407 Self::InvalidIndexColumn(name) => write!(f, "invalid index column: {name}"),
408 Self::InvalidConstraintColumn(name) => {
409 write!(f, "invalid constraint column: {name}")
410 }
411 Self::TruncatedData => write!(f, "truncated data"),
412 Self::InvalidMagic => write!(f, "invalid magic bytes"),
413 Self::InvalidDataType => write!(f, "invalid data type"),
414 Self::InvalidIndexType => write!(f, "invalid index type"),
415 Self::InvalidConstraintType => write!(f, "invalid constraint type"),
416 Self::VarintOverflow => write!(f, "varint overflow"),
417 }
418 }
419}
420
421impl std::error::Error for TableDefError {}
422
423#[cfg(test)]
424mod tests {
425 use super::*;
426
427 #[test]
428 fn table_validation_accepts_valid_references() {
429 let table = TableDef::new("hosts")
430 .add_column(ColumnDef::new("id", DataType::UnsignedInteger).not_null())
431 .add_column(ColumnDef::new("ip", DataType::IpAddr))
432 .primary_key(vec!["id".into()])
433 .add_index(IndexDef::new("idx_ip", vec!["ip".into()]).unique());
434
435 assert!(table.validate().is_ok());
436 assert!(table.is_primary_key_column("id"));
437 assert_eq!(table.column_index("ip"), Some(1));
438 }
439
440 #[test]
441 fn table_validation_rejects_invalid_references() {
442 let duplicate = TableDef::new("hosts")
443 .add_column(ColumnDef::new("id", DataType::Integer))
444 .add_column(ColumnDef::new("id", DataType::Text));
445 assert_eq!(
446 duplicate.validate(),
447 Err(TableDefError::DuplicateColumn("id".into()))
448 );
449
450 let missing_key = TableDef::new("hosts").primary_key(vec!["id".into()]);
451 assert_eq!(
452 missing_key.validate(),
453 Err(TableDefError::InvalidPrimaryKey("id".into()))
454 );
455 }
456
457 #[test]
458 fn builders_preserve_column_and_constraint_details() {
459 let column = ColumnDef::new("embedding", DataType::Vector)
460 .not_null()
461 .with_default(vec![1, 2, 3])
462 .with_vector_dim(384)
463 .compressed()
464 .with_variants(vec!["a".into(), "b".into()])
465 .with_precision(6)
466 .with_element_type(DataType::Float)
467 .with_metadata("unit", "f32");
468 assert_eq!(column.vector_dim, Some(384));
469 assert_eq!(column.metadata.get("unit").map(String::as_str), Some("f32"));
470
471 let constraint = Constraint::new("fk_host", ConstraintType::ForeignKey)
472 .on_columns(vec!["host_id".into()])
473 .references("hosts".into(), vec!["id".into()]);
474 assert_eq!(constraint.ref_table.as_deref(), Some("hosts"));
475 }
476}