Skip to main content

mongreldb_kit_core/
schema.rs

1//! Language-neutral schema model for MongrelDB Kit.
2//!
3//! A [`Schema`] is a collection of [`Table`]s. Each table has [`Column`]s,
4//! indexes, unique constraints, foreign keys, and check constraints.
5
6use serde::{Deserialize, Serialize};
7use std::collections::{HashMap, HashSet};
8
9/// Storage/application types supported by Kit columns.
10#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
11#[serde(rename_all = "snake_case")]
12pub enum ColumnType {
13    Bool,
14    Int8,
15    Int16,
16    Int32,
17    Int64,
18    Float32,
19    Float64,
20    Text,
21    Bytes,
22    Json,
23    Date,
24    DateTime,
25    TimestampNanos,
26    Date64,
27    Time64,
28    Interval,
29    Decimal128,
30    /// RFC 4122 UUID (16 bytes, big-endian).
31    Uuid,
32    /// Native JSON value (stored as typed bytes).
33    JsonNative,
34    /// Variable-length array of homogeneous values.
35    Array,
36    /// A dense float32 vector for nearest-neighbour (ANN) search. The dimension
37    /// is carried on the column as `embedding_dim`.
38    Embedding,
39    /// A learned-sparse (SPLADE-style) weighted token vector, stored as a
40    /// `[[token_id, weight], ...]` list, for sparse retrieval.
41    Sparse,
42}
43
44/// How a default value is produced when a row omits a column.
45#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
46#[serde(rename_all = "snake_case")]
47pub enum DefaultKind {
48    /// A fixed JSON value written literally.
49    Static(serde_json::Value),
50    /// The current timestamp as an ISO-8601 string.
51    Now,
52    /// A fresh UUIDv4 string.
53    Uuid,
54    /// The next value from a named sequence.
55    Sequence(String),
56    /// A user-defined default registered by name (resolved at runtime).
57    CustomName(String),
58}
59
60/// Where dense embedding values for a column originate.
61///
62/// Mirrors `mongreldb_core::EmbeddingSource` in a language-neutral shape
63/// (string paths, serde-tagged). Omitting this on a column means
64/// application-supplied vectors (the engine default). Transactional generation
65/// uses a portable [`GeneratedEmbeddingSpec`] plus a process-local provider.
66#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, Default)]
67#[serde(rename_all = "snake_case")]
68pub enum EmbeddingSpecNormalization {
69    #[default]
70    None,
71    L2,
72}
73
74#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, Default)]
75#[serde(rename_all = "snake_case")]
76pub enum EmbeddingWriteFailurePolicy {
77    #[default]
78    AbortWrite,
79}
80
81#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
82pub struct GeneratedEmbeddingSpec {
83    pub provider_id: String,
84    pub model_id: String,
85    pub model_version: String,
86    pub source_columns: Vec<u32>,
87    pub input_template: String,
88    pub dimension: u32,
89    pub normalization: EmbeddingSpecNormalization,
90    pub failure_policy: EmbeddingWriteFailurePolicy,
91}
92
93#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
94#[serde(tag = "kind", rename_all = "snake_case")]
95pub enum EmbeddingSource {
96    /// Application writes float vectors directly (default).
97    SuppliedByApplication,
98    /// Local on-disk model; a provider registered under `model_id` runs inference.
99    LocalModel {
100        /// Filesystem path to model weights / tokenizer bundle.
101        model_path: String,
102        /// Stable model identity (registry key and ANN generation stamp).
103        model_id: String,
104    },
105    /// Portable model identity resolved from node-local provider configuration.
106    ConfiguredModel {
107        provider_id: String,
108        model_id: String,
109        model_version: String,
110    },
111    /// Named provider registered on the process (`provider` registry key).
112    GeneratedColumn {
113        /// Registry key of the provider.
114        provider: String,
115    },
116    /// Transactionally materialized embedding from source columns.
117    GeneratedColumnSpec { spec: GeneratedEmbeddingSpec },
118}
119
120/// A column definition.
121#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
122pub struct Column {
123    /// Stable column identifier. IDs must be unique within a table.
124    pub id: u32,
125    /// Logical column name.
126    pub name: String,
127    /// Physical storage type.
128    pub storage_type: ColumnType,
129    /// Application-facing type (often the same as `storage_type`).
130    pub application_type: ColumnType,
131    /// Whether the column may contain `null`.
132    pub nullable: bool,
133    /// Whether this column is part of the primary key.
134    pub primary_key: bool,
135    /// Optional default value generator.
136    pub default: Option<DefaultKind>,
137    /// Whether the value is generated on every mutation.
138    pub generated: bool,
139    /// Permitted string values, if any.
140    #[serde(default, skip_serializing_if = "Option::is_none")]
141    pub enum_values: Option<Vec<String>>,
142    /// Minimum numeric value.
143    #[serde(default, skip_serializing_if = "Option::is_none")]
144    pub min: Option<f64>,
145    /// Maximum numeric value.
146    #[serde(default, skip_serializing_if = "Option::is_none")]
147    pub max: Option<f64>,
148    /// Minimum string/bytes length.
149    #[serde(default, skip_serializing_if = "Option::is_none")]
150    pub min_length: Option<usize>,
151    /// Maximum string/bytes length.
152    #[serde(default, skip_serializing_if = "Option::is_none")]
153    pub max_length: Option<usize>,
154    /// Regular expression a `text` value must match, stored as its source pattern.
155    #[serde(default, skip_serializing_if = "Option::is_none")]
156    pub regex: Option<String>,
157    /// An optional check expression name for runtime evaluation.
158    #[serde(default, skip_serializing_if = "Option::is_none")]
159    pub check_expr: Option<String>,
160    /// Vector dimension for an `Embedding` column (required for ANN).
161    #[serde(default, skip_serializing_if = "Option::is_none")]
162    pub embedding_dim: Option<u32>,
163    /// How embedding values are produced. Only meaningful for
164    /// [`ColumnType::Embedding`]. `None` = application-supplied (engine default).
165    #[serde(default, skip_serializing_if = "Option::is_none")]
166    pub embedding_source: Option<EmbeddingSource>,
167    /// Encrypt this column's page payload at rest (requires an encrypted db).
168    #[serde(default, skip_serializing_if = "std::ops::Not::not")]
169    pub encrypted: bool,
170    /// Encrypt the column but keep it queryable via deterministic equality
171    /// tokens / order-preserving encoding (requires an encrypted db).
172    #[serde(default, skip_serializing_if = "std::ops::Not::not")]
173    pub encrypted_indexable: bool,
174}
175
176impl Column {
177    /// Convenience constructor for the common case.
178    pub fn new(id: u32, name: impl Into<String>, storage_type: ColumnType) -> Self {
179        Self {
180            id,
181            name: name.into(),
182            storage_type,
183            application_type: storage_type,
184            nullable: false,
185            primary_key: false,
186            default: None,
187            generated: false,
188            enum_values: None,
189            min: None,
190            max: None,
191            min_length: None,
192            max_length: None,
193            regex: None,
194            check_expr: None,
195            embedding_dim: None,
196            embedding_source: None,
197            encrypted: false,
198            encrypted_indexable: false,
199        }
200    }
201}
202
203/// The kind of secondary index the Kit declares on a column.
204#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, Serialize, Deserialize)]
205#[serde(rename_all = "snake_case")]
206pub enum IndexKind {
207    /// Equality / `IN` acceleration (the default).
208    #[default]
209    Bitmap,
210    /// FM-index substring search (`contains(col, needle)` pushes to `FmContains`).
211    Fm,
212    /// HNSW approximate-nearest-neighbour index for `Embedding` columns.
213    Ann,
214    /// SPLADE-style learned-sparse retrieval index for `Sparse` columns.
215    Sparse,
216    /// MinHash/LSH set-similarity index over a JSON-array set column
217    /// (accelerates `set_similarity`).
218    MinHash,
219    /// Learned zonemap (PGM) index for ordered range predicates on numeric /
220    /// timestamp columns. Accelerates `Range`/`RangeF64` conditions.
221    LearnedRange,
222}
223
224/// ANN representation and distance metric.
225#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, Serialize, Deserialize)]
226#[serde(rename_all = "snake_case")]
227pub enum AnnQuantization {
228    /// One sign bit per component, ranked by Hamming distance.
229    #[default]
230    BinarySign,
231    /// Full `f32` vectors, ranked by cosine distance.
232    Dense,
233    /// Product quantization: vectors split into `num_subvectors` groups, each
234    /// encoded to `bits`-bit codes against trained codebooks. Ranked by ADC
235    /// (asymmetric distance computation) with optional exact rerank.
236    Product { num_subvectors: u16, bits: u8 },
237}
238
239/// ANN graph/structure algorithm. Orthogonal to [`AnnQuantization`]: the
240/// algorithm chooses how search walks the index; the quantization chooses how
241/// vectors are represented.
242#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, Serialize, Deserialize)]
243pub enum AnnAlgorithm {
244    /// Hierarchical Navigable Small World (the default).
245    #[default]
246    #[serde(rename = "hnsw")]
247    Hnsw,
248    /// DiskANN / Vamana single-layer robust-pruned graph.
249    #[serde(rename = "diskann")]
250    DiskAnn,
251    /// Inverted file index (k-means centroids + inverted lists).
252    #[serde(rename = "ivf")]
253    Ivf,
254}
255
256/// An index on one or more columns.
257#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, Default)]
258pub struct Index {
259    pub name: String,
260    pub columns: Vec<String>,
261    pub unique: bool,
262    /// Index kind; defaults to `Bitmap` so pre-existing schemas deserialize
263    /// unchanged.
264    #[serde(default)]
265    pub kind: IndexKind,
266    /// ANN representation. Ignored for non-ANN indexes.
267    #[serde(default)]
268    pub ann_quantization: AnnQuantization,
269    /// ANN graph/structure algorithm (Phase 2). Defaults to HNSW. Ignored for
270    /// non-ANN indexes. Orthogonal to `ann_quantization`.
271    #[serde(default)]
272    pub ann_algorithm: AnnAlgorithm,
273    /// Optional SQL predicate for a partial index.
274    #[serde(default, skip_serializing_if = "Option::is_none")]
275    pub predicate: Option<String>,
276    /// HNSW graph degree. Engine default when omitted.
277    #[serde(default, skip_serializing_if = "Option::is_none")]
278    pub ann_m: Option<usize>,
279    /// HNSW construction search width. Engine default when omitted.
280    #[serde(default, skip_serializing_if = "Option::is_none")]
281    pub ann_ef_construction: Option<usize>,
282    /// HNSW query search width. Engine default when omitted.
283    #[serde(default, skip_serializing_if = "Option::is_none")]
284    pub ann_ef_search: Option<usize>,
285    /// DiskANN max graph degree R. Engine default when omitted.
286    #[serde(default, skip_serializing_if = "Option::is_none")]
287    pub ann_diskann_r: Option<usize>,
288    /// DiskANN build search-list size L. Engine default when omitted.
289    #[serde(default, skip_serializing_if = "Option::is_none")]
290    pub ann_diskann_l: Option<usize>,
291    /// DiskANN query beam width. Engine default when omitted.
292    #[serde(default, skip_serializing_if = "Option::is_none")]
293    pub ann_diskann_beam_width: Option<usize>,
294    /// DiskANN robust-prune alpha × 100 (120 = 1.2). Engine default when omitted.
295    #[serde(default, skip_serializing_if = "Option::is_none")]
296    pub ann_diskann_alpha: Option<u32>,
297    /// IVF inverted-list (centroid) count. Engine default when omitted.
298    #[serde(default, skip_serializing_if = "Option::is_none")]
299    pub ann_ivf_nlist: Option<usize>,
300    /// IVF probe count at query time. Engine default when omitted.
301    #[serde(default, skip_serializing_if = "Option::is_none")]
302    pub ann_ivf_nprobe: Option<usize>,
303    /// IVF k-means training sample cap. Engine default when omitted.
304    #[serde(default, skip_serializing_if = "Option::is_none")]
305    pub ann_ivf_training_samples: Option<usize>,
306    /// Product-quantizer training sample cap. Engine default when omitted.
307    #[serde(default, skip_serializing_if = "Option::is_none")]
308    pub ann_pq_training_samples: Option<usize>,
309    /// Product-quantizer deterministic training seed. Engine default when omitted.
310    #[serde(default, skip_serializing_if = "Option::is_none")]
311    pub ann_pq_seed: Option<u64>,
312    /// Product-quantizer exact-rerank factor (0 disables). Engine default when omitted.
313    #[serde(default, skip_serializing_if = "Option::is_none")]
314    pub ann_pq_rerank_factor: Option<usize>,
315    /// MinHash permutation count. Engine default when omitted.
316    #[serde(default, skip_serializing_if = "Option::is_none")]
317    pub minhash_permutations: Option<usize>,
318    /// MinHash LSH band count. Engine default when omitted.
319    #[serde(default, skip_serializing_if = "Option::is_none")]
320    pub minhash_bands: Option<usize>,
321    /// Learned-range error bound. Engine default when omitted.
322    #[serde(default, skip_serializing_if = "Option::is_none")]
323    pub learned_range_epsilon: Option<usize>,
324}
325
326/// A uniqueness constraint over one or more columns.
327#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
328pub struct UniqueConstraint {
329    pub name: String,
330    pub columns: Vec<String>,
331}
332
333/// A foreign-key reference from child columns to parent columns.
334#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
335pub struct ForeignKey {
336    pub name: String,
337    pub columns: Vec<String>,
338    pub references_table: String,
339    pub references_columns: Vec<String>,
340    #[serde(default)]
341    pub on_delete: ForeignKeyAction,
342}
343
344/// Action taken when a referenced parent row is deleted.
345#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, Serialize, Deserialize)]
346#[serde(rename_all = "snake_case")]
347pub enum ForeignKeyAction {
348    #[default]
349    Restrict,
350    Cascade,
351    SetNull,
352}
353
354/// A named table-level check constraint.
355#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
356pub struct CheckConstraint {
357    pub name: String,
358    pub expr: String,
359}
360
361/// A monotonic sequence allocator.
362#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
363pub struct Sequence {
364    pub name: String,
365    pub next_value: i64,
366}
367
368/// A table definition.
369#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
370pub struct Table {
371    /// Stable table identifier. IDs must be unique within a schema.
372    pub id: u32,
373    pub name: String,
374    pub columns: Vec<Column>,
375    pub primary_key: Vec<String>,
376    #[serde(default, skip_serializing_if = "Vec::is_empty")]
377    pub indexes: Vec<Index>,
378    #[serde(default, skip_serializing_if = "Vec::is_empty")]
379    pub foreign_keys: Vec<ForeignKey>,
380    #[serde(default, skip_serializing_if = "Vec::is_empty")]
381    pub unique_constraints: Vec<UniqueConstraint>,
382    #[serde(default, skip_serializing_if = "Vec::is_empty")]
383    pub check_constraints: Vec<CheckConstraint>,
384}
385
386impl Table {
387    /// Find a column by name.
388    pub fn column(&self, name: &str) -> Option<&Column> {
389        self.columns.iter().find(|c| c.name == name)
390    }
391
392    /// Whether the named column is part of the primary key.
393    pub fn is_pk_column(&self, name: &str) -> bool {
394        self.primary_key.iter().any(|c| c == name)
395    }
396}
397
398/// Errors that can occur while constructing a [`Schema`].
399#[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)]
400pub enum SchemaError {
401    #[error("duplicate table name \"{0}\"")]
402    DuplicateTableName(String),
403    #[error("duplicate table id {0}")]
404    DuplicateTableId(u32),
405    #[error("duplicate column name \"{1}\" in table \"{0}\"")]
406    DuplicateColumnName(String, String),
407    #[error("duplicate column id {1} in table \"{0}\"")]
408    DuplicateColumnId(String, u32),
409    #[error("primary key column \"{1}\" not found in table \"{0}\"")]
410    MissingPrimaryKeyColumn(String, String),
411    #[error("index \"{1}\" references unknown column \"{2}\" in table \"{0}\"")]
412    MissingIndexColumn(String, String, String),
413    #[error("unique constraint \"{1}\" references unknown column \"{2}\" in table \"{0}\"")]
414    MissingUniqueColumn(String, String, String),
415    #[error("foreign key \"{1}\" references unknown column \"{2}\" in table \"{0}\"")]
416    MissingForeignKeyColumn(String, String, String),
417    #[error("foreign key \"{1}\" references unknown table \"{2}\"")]
418    MissingReferencedTable(String, String, String),
419    #[error("foreign key \"{1}\" references unknown column \"{2}\" on table \"{3}\"")]
420    MissingReferencedColumn(String, String, String, String),
421    #[error(
422        "column \"{1}\" on table \"{0}\" sets embedding_source but is not an embedding column"
423    )]
424    EmbeddingSourceOnNonEmbedding(String, String),
425    #[error(
426        "embedding column \"{1}\" on table \"{0}\" with LocalModel/GeneratedColumn source requires embedding_dim > 0"
427    )]
428    EmbeddingSourceMissingDim(String, String),
429    #[error("generated embedding column \"{1}\" on table \"{0}\" is invalid: {2}")]
430    InvalidGeneratedEmbeddingSpec(String, String, String),
431}
432
433/// A validated collection of tables.
434#[derive(Debug, Clone, PartialEq, Serialize)]
435pub struct Schema {
436    pub tables: Vec<Table>,
437    by_name: HashMap<String, usize>,
438    by_id: HashMap<u32, usize>,
439}
440
441impl<'de> serde::Deserialize<'de> for Schema {
442    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
443    where
444        D: serde::Deserializer<'de>,
445    {
446        #[derive(serde::Deserialize)]
447        struct SchemaHelper {
448            tables: Vec<Table>,
449        }
450        let helper = SchemaHelper::deserialize(deserializer)?;
451        Schema::new(helper.tables).map_err(serde::de::Error::custom)
452    }
453}
454
455/// A unique index also enforces uniqueness (guard-backed), matching SQL where a
456/// UNIQUE index is a UNIQUE constraint. Synthesize a constraint for each unique
457/// index unless an existing (or already-synthesized) unique constraint already
458/// covers exactly the same columns. Mirrors the TypeScript kit's `table()`.
459fn synthesize_unique_from_indexes(table: &mut Table) {
460    let mut synthesized: Vec<UniqueConstraint> = Vec::new();
461    for idx in &table.indexes {
462        if !idx.unique {
463            continue;
464        }
465        let covered = table
466            .unique_constraints
467            .iter()
468            .chain(synthesized.iter())
469            .any(|u| u.columns == idx.columns);
470        if !covered {
471            synthesized.push(UniqueConstraint {
472                name: idx.name.clone(),
473                columns: idx.columns.clone(),
474            });
475        }
476    }
477    table.unique_constraints.extend(synthesized);
478}
479
480impl Schema {
481    /// Build and validate a schema from a list of tables.
482    pub fn new(mut tables: Vec<Table>) -> Result<Self, SchemaError> {
483        for table in &mut tables {
484            synthesize_unique_from_indexes(table);
485        }
486
487        let mut by_name = HashMap::with_capacity(tables.len());
488        let mut by_id = HashMap::with_capacity(tables.len());
489
490        for (idx, table) in tables.iter().enumerate() {
491            if by_name.contains_key(&table.name) {
492                return Err(SchemaError::DuplicateTableName(table.name.clone()));
493            }
494            if by_id.contains_key(&table.id) {
495                return Err(SchemaError::DuplicateTableId(table.id));
496            }
497            by_name.insert(table.name.clone(), idx);
498            by_id.insert(table.id, idx);
499        }
500
501        for table in &tables {
502            Self::validate_table(table, &by_name)?;
503        }
504
505        Ok(Self {
506            tables,
507            by_name,
508            by_id,
509        })
510    }
511
512    fn validate_table(
513        table: &Table,
514        table_names: &HashMap<String, usize>,
515    ) -> Result<(), SchemaError> {
516        let mut column_names = HashMap::with_capacity(table.columns.len());
517        let mut column_ids = HashMap::with_capacity(table.columns.len());
518
519        for col in &table.columns {
520            if column_names.contains_key(&col.name) {
521                return Err(SchemaError::DuplicateColumnName(
522                    table.name.clone(),
523                    col.name.clone(),
524                ));
525            }
526            if column_ids.contains_key(&col.id) {
527                return Err(SchemaError::DuplicateColumnId(table.name.clone(), col.id));
528            }
529            if col.embedding_source.is_some() && col.storage_type != ColumnType::Embedding {
530                return Err(SchemaError::EmbeddingSourceOnNonEmbedding(
531                    table.name.clone(),
532                    col.name.clone(),
533                ));
534            }
535            if matches!(
536                col.embedding_source,
537                Some(
538                    EmbeddingSource::LocalModel { .. }
539                        | EmbeddingSource::ConfiguredModel { .. }
540                        | EmbeddingSource::GeneratedColumn { .. }
541                        | EmbeddingSource::GeneratedColumnSpec { .. }
542                )
543            ) && col.embedding_dim.unwrap_or(0) == 0
544            {
545                return Err(SchemaError::EmbeddingSourceMissingDim(
546                    table.name.clone(),
547                    col.name.clone(),
548                ));
549            }
550            column_names.insert(col.name.clone(), col.id);
551            column_ids.insert(col.id, col.name.clone());
552        }
553
554        for col in &table.columns {
555            if let Some(EmbeddingSource::ConfiguredModel {
556                provider_id,
557                model_id,
558                model_version,
559            }) = col.embedding_source.as_ref()
560            {
561                if provider_id.is_empty() || model_id.is_empty() || model_version.is_empty() {
562                    return Err(SchemaError::InvalidGeneratedEmbeddingSpec(
563                        table.name.clone(),
564                        col.name.clone(),
565                        "provider, model, and version are required".into(),
566                    ));
567                }
568            }
569        }
570
571        for col in &table.columns {
572            let Some(EmbeddingSource::GeneratedColumnSpec { spec }) = col.embedding_source.as_ref()
573            else {
574                continue;
575            };
576            if spec.provider_id.is_empty()
577                || spec.model_id.is_empty()
578                || spec.model_version.is_empty()
579                || spec.source_columns.is_empty()
580            {
581                return Err(SchemaError::InvalidGeneratedEmbeddingSpec(
582                    table.name.clone(),
583                    col.name.clone(),
584                    "provider, model, version, and source columns are required".into(),
585                ));
586            }
587            if col.embedding_dim != Some(spec.dimension) {
588                return Err(SchemaError::InvalidGeneratedEmbeddingSpec(
589                    table.name.clone(),
590                    col.name.clone(),
591                    "spec dimension must match embedding_dim".into(),
592                ));
593            }
594            let mut seen = HashSet::new();
595            if spec.source_columns.iter().any(|source_id| {
596                *source_id == col.id
597                    || !seen.insert(*source_id)
598                    || !column_ids.contains_key(source_id)
599            }) {
600                return Err(SchemaError::InvalidGeneratedEmbeddingSpec(
601                    table.name.clone(),
602                    col.name.clone(),
603                    "source columns must exist, be unique, and exclude the target".into(),
604                ));
605            }
606        }
607
608        for pk in &table.primary_key {
609            if !column_names.contains_key(pk) {
610                return Err(SchemaError::MissingPrimaryKeyColumn(
611                    table.name.clone(),
612                    pk.clone(),
613                ));
614            }
615        }
616
617        for idx in &table.indexes {
618            for col in &idx.columns {
619                if !column_names.contains_key(col) {
620                    return Err(SchemaError::MissingIndexColumn(
621                        table.name.clone(),
622                        idx.name.clone(),
623                        col.clone(),
624                    ));
625                }
626            }
627        }
628
629        for uq in &table.unique_constraints {
630            for col in &uq.columns {
631                if !column_names.contains_key(col) {
632                    return Err(SchemaError::MissingUniqueColumn(
633                        table.name.clone(),
634                        uq.name.clone(),
635                        col.clone(),
636                    ));
637                }
638            }
639        }
640
641        for fk in &table.foreign_keys {
642            for col in &fk.columns {
643                if !column_names.contains_key(col) {
644                    return Err(SchemaError::MissingForeignKeyColumn(
645                        table.name.clone(),
646                        fk.name.clone(),
647                        col.clone(),
648                    ));
649                }
650            }
651            if !table_names.contains_key(&fk.references_table) {
652                return Err(SchemaError::MissingReferencedTable(
653                    table.name.clone(),
654                    fk.name.clone(),
655                    fk.references_table.clone(),
656                ));
657            }
658        }
659
660        Ok(())
661    }
662
663    /// Look up a table by name.
664    pub fn table(&self, name: &str) -> Option<&Table> {
665        self.by_name.get(name).map(|&idx| &self.tables[idx])
666    }
667
668    /// Look up a table by stable id.
669    pub fn table_by_id(&self, id: u32) -> Option<&Table> {
670        self.by_id.get(&id).map(|&idx| &self.tables[idx])
671    }
672
673    /// Whether the schema contains a table with the given name.
674    pub fn has_table(&self, name: &str) -> bool {
675        self.by_name.contains_key(name)
676    }
677
678    /// Rename a table in place, keeping the `by_name` index in sync. Returns
679    /// `false` if `from` is absent or `to` is already in use (no change made).
680    /// Does *not* retarget foreign keys — callers that need that should do it
681    /// before/after on the tables they own.
682    pub fn rename_table(&mut self, from: &str, to: &str) -> bool {
683        if from == to {
684            return self.has_table(from);
685        }
686        if !self.has_table(from) || self.has_table(to) {
687            return false;
688        }
689        let idx = *self.by_name.get(from).unwrap();
690        self.tables[idx].name = to.to_string();
691        self.by_name.remove(from);
692        self.by_name.insert(to.to_string(), idx);
693        true
694    }
695}
696
697#[cfg(test)]
698mod tests {
699    use super::*;
700
701    fn make_table(name: &str, id: u32) -> Table {
702        Table {
703            id,
704            name: name.into(),
705            columns: vec![Column::new(1, "id", ColumnType::Int64)],
706            primary_key: vec!["id".into()],
707            indexes: vec![],
708            foreign_keys: vec![],
709            unique_constraints: vec![],
710            check_constraints: vec![],
711        }
712    }
713
714    #[test]
715    fn schema_rejects_duplicate_table_name() {
716        let err = Schema::new(vec![make_table("a", 1), make_table("a", 2)]).unwrap_err();
717        assert!(matches!(err, SchemaError::DuplicateTableName(n) if n == "a"));
718    }
719
720    #[test]
721    fn schema_rejects_duplicate_table_id() {
722        let err = Schema::new(vec![make_table("a", 1), make_table("b", 1)]).unwrap_err();
723        assert!(matches!(err, SchemaError::DuplicateTableId(1)));
724    }
725
726    #[test]
727    fn schema_rejects_missing_pk_column() {
728        let t = Table {
729            id: 1,
730            name: "t".into(),
731            columns: vec![Column::new(1, "x", ColumnType::Text)],
732            primary_key: vec!["id".into()],
733            indexes: vec![],
734            foreign_keys: vec![],
735            unique_constraints: vec![],
736            check_constraints: vec![],
737        };
738        let err = Schema::new(vec![t]).unwrap_err();
739        assert!(matches!(err, SchemaError::MissingPrimaryKeyColumn(_, _)));
740    }
741
742    #[test]
743    fn unique_index_synthesizes_unique_constraint() {
744        let schema = Schema::new(vec![Table {
745            id: 1,
746            name: "users".into(),
747            columns: vec![
748                Column::new(1, "id", ColumnType::Int64),
749                Column::new(2, "email", ColumnType::Text),
750                Column::new(3, "handle", ColumnType::Text),
751            ],
752            primary_key: vec!["id".into()],
753            indexes: vec![
754                Index {
755                    name: "idx_email".into(),
756                    columns: vec!["email".into()],
757                    unique: true,
758                    kind: Default::default(),
759                    ann_quantization: Default::default(),
760                    ..Default::default()
761                },
762                // A non-unique index must NOT synthesize a constraint.
763                Index {
764                    name: "idx_handle".into(),
765                    columns: vec!["handle".into()],
766                    unique: false,
767                    kind: Default::default(),
768                    ann_quantization: Default::default(),
769                    ..Default::default()
770                },
771            ],
772            foreign_keys: vec![],
773            unique_constraints: vec![],
774            check_constraints: vec![],
775        }])
776        .unwrap();
777        let table = schema.table("users").unwrap();
778        assert_eq!(table.unique_constraints.len(), 1);
779        assert_eq!(
780            table.unique_constraints[0].columns,
781            vec!["email".to_string()]
782        );
783    }
784
785    #[test]
786    fn unique_index_does_not_duplicate_existing_constraint() {
787        let schema = Schema::new(vec![Table {
788            id: 1,
789            name: "users".into(),
790            columns: vec![
791                Column::new(1, "id", ColumnType::Int64),
792                Column::new(2, "email", ColumnType::Text),
793            ],
794            primary_key: vec!["id".into()],
795            indexes: vec![Index {
796                name: "idx_email".into(),
797                columns: vec!["email".into()],
798                unique: true,
799                kind: Default::default(),
800                ann_quantization: Default::default(),
801                ..Default::default()
802            }],
803            foreign_keys: vec![],
804            unique_constraints: vec![UniqueConstraint {
805                name: "uq_email".into(),
806                columns: vec!["email".into()],
807            }],
808            check_constraints: vec![],
809        }])
810        .unwrap();
811        // The pre-existing constraint already covers `email`; no synthesis.
812        let table = schema.table("users").unwrap();
813        assert_eq!(table.unique_constraints.len(), 1);
814        assert_eq!(table.unique_constraints[0].name, "uq_email");
815    }
816
817    #[test]
818    fn schema_roundtrips_json() {
819        let schema = Schema::new(vec![Table {
820            id: 1,
821            name: "users".into(),
822            columns: vec![
823                Column::new(1, "id", ColumnType::Int64),
824                Column {
825                    nullable: true,
826                    ..Column::new(2, "email", ColumnType::Text)
827                },
828            ],
829            primary_key: vec!["id".into()],
830            indexes: vec![Index {
831                name: "idx_email".into(),
832                columns: vec!["email".into()],
833                unique: true,
834                kind: Default::default(),
835                ann_quantization: Default::default(),
836                ..Default::default()
837            }],
838            foreign_keys: vec![],
839            unique_constraints: vec![],
840            check_constraints: vec![CheckConstraint {
841                name: "chk_id_positive".into(),
842                expr: "id > 0".into(),
843            }],
844        }])
845        .unwrap();
846
847        let json = serde_json::to_string(&schema).unwrap();
848        let decoded: Schema = serde_json::from_str(&json).unwrap();
849        assert_eq!(decoded.tables.len(), 1);
850        assert_eq!(decoded.table("users").unwrap().columns.len(), 2);
851    }
852
853    #[test]
854    fn dense_ann_quantization_roundtrips_and_old_json_defaults_binary() {
855        let dense = Index {
856            name: "idx_embedding".into(),
857            columns: vec!["embedding".into()],
858            unique: false,
859            kind: IndexKind::Ann,
860            ann_quantization: AnnQuantization::Dense,
861            predicate: Some("embedding IS NOT NULL".into()),
862            ann_m: Some(24),
863            ann_ef_construction: Some(96),
864            ann_ef_search: Some(48),
865            ..Default::default()
866        };
867        let json = serde_json::to_value(&dense).unwrap();
868        assert_eq!(json["ann_quantization"], "dense");
869        assert_eq!(json["ann_m"], 24);
870        assert_eq!(json["predicate"], "embedding IS NOT NULL");
871        assert_eq!(
872            serde_json::from_value::<Index>(json)
873                .unwrap()
874                .ann_quantization,
875            AnnQuantization::Dense
876        );
877
878        let old: Index = serde_json::from_value(serde_json::json!({
879            "name": "idx_embedding",
880            "columns": ["embedding"],
881            "unique": false,
882            "kind": "ann"
883        }))
884        .unwrap();
885        assert_eq!(old.ann_quantization, AnnQuantization::BinarySign);
886    }
887
888    #[test]
889    fn swappable_ann_algorithm_and_product_quantization_roundtrip() {
890        let diskann = Index {
891            name: "idx_diskann".into(),
892            columns: vec!["embedding".into()],
893            unique: false,
894            kind: IndexKind::Ann,
895            ann_algorithm: AnnAlgorithm::DiskAnn,
896            ann_quantization: AnnQuantization::Dense,
897            ann_diskann_r: Some(128),
898            ann_diskann_l: Some(256),
899            ann_diskann_beam_width: Some(4),
900            ann_diskann_alpha: Some(130),
901            ..Default::default()
902        };
903        let json = serde_json::to_value(&diskann).unwrap();
904        assert_eq!(json["ann_algorithm"], "diskann");
905        assert_eq!(json["ann_diskann_r"], 128);
906        assert_eq!(json["ann_diskann_alpha"], 130);
907        let decoded: Index = serde_json::from_value(json).unwrap();
908        assert_eq!(decoded.ann_algorithm, AnnAlgorithm::DiskAnn);
909        assert_eq!(decoded.ann_diskann_l, Some(256));
910
911        let ivf = Index {
912            name: "idx_ivf".into(),
913            columns: vec!["embedding".into()],
914            kind: IndexKind::Ann,
915            ann_algorithm: AnnAlgorithm::Ivf,
916            ann_quantization: AnnQuantization::Dense,
917            ann_ivf_nlist: Some(512),
918            ann_ivf_nprobe: Some(16),
919            ann_ivf_training_samples: Some(20_000),
920            ..Default::default()
921        };
922        let json = serde_json::to_value(&ivf).unwrap();
923        assert_eq!(json["ann_algorithm"], "ivf");
924        assert_eq!(json["ann_ivf_nlist"], 512);
925        assert_eq!(json["ann_ivf_training_samples"], 20_000);
926
927        let pq = Index {
928            name: "idx_pq".into(),
929            columns: vec!["embedding".into()],
930            kind: IndexKind::Ann,
931            ann_quantization: AnnQuantization::Product {
932                num_subvectors: 32,
933                bits: 8,
934            },
935            ann_pq_training_samples: Some(10_000),
936            ann_pq_seed: Some(42),
937            ann_pq_rerank_factor: Some(3),
938            ..Default::default()
939        };
940        let json = serde_json::to_value(&pq).unwrap();
941        assert_eq!(json["ann_quantization"]["product"]["num_subvectors"], 32);
942        assert_eq!(json["ann_quantization"]["product"]["bits"], 8);
943        assert_eq!(json["ann_pq_seed"], 42);
944        let decoded: Index = serde_json::from_value(json).unwrap();
945        assert_eq!(
946            decoded.ann_quantization,
947            AnnQuantization::Product {
948                num_subvectors: 32,
949                bits: 8
950            }
951        );
952
953        // Default algorithm is HNSW and is omitted on the wire (backward compat).
954        let hnsw = Index {
955            name: "idx_hnsw".into(),
956            columns: vec!["embedding".into()],
957            kind: IndexKind::Ann,
958            ann_quantization: AnnQuantization::Dense,
959            ..Default::default()
960        };
961        let json = serde_json::to_value(&hnsw).unwrap();
962        assert_eq!(json["ann_algorithm"], "hnsw");
963    }
964
965    #[test]
966    fn embedding_source_roundtrips_json() {
967        let mut emb = Column::new(2, "vec", ColumnType::Embedding);
968        emb.embedding_dim = Some(4);
969        emb.embedding_source = Some(EmbeddingSource::LocalModel {
970            model_path: "/models/kit-mini".into(),
971            model_id: "kit-mini".into(),
972        });
973        let schema = Schema::new(vec![Table {
974            id: 1,
975            name: "docs".into(),
976            columns: vec![Column::new(1, "id", ColumnType::Int64), emb],
977            primary_key: vec!["id".into()],
978            indexes: vec![],
979            foreign_keys: vec![],
980            unique_constraints: vec![],
981            check_constraints: vec![],
982        }])
983        .unwrap();
984        let json = serde_json::to_string(&schema).unwrap();
985        assert!(json.contains("local_model"));
986        assert!(json.contains("kit-mini"));
987        let decoded: Schema = serde_json::from_str(&json).unwrap();
988        let col = decoded.table("docs").unwrap().column("vec").unwrap();
989        assert_eq!(
990            col.embedding_source,
991            Some(EmbeddingSource::LocalModel {
992                model_path: "/models/kit-mini".into(),
993                model_id: "kit-mini".into(),
994            })
995        );
996    }
997
998    #[test]
999    fn embedding_source_rejected_on_non_embedding() {
1000        let mut col = Column::new(2, "name", ColumnType::Text);
1001        col.embedding_source = Some(EmbeddingSource::SuppliedByApplication);
1002        let err = Schema::new(vec![Table {
1003            id: 1,
1004            name: "t".into(),
1005            columns: vec![Column::new(1, "id", ColumnType::Int64), col],
1006            primary_key: vec!["id".into()],
1007            indexes: vec![],
1008            foreign_keys: vec![],
1009            unique_constraints: vec![],
1010            check_constraints: vec![],
1011        }])
1012        .unwrap_err();
1013        assert!(matches!(
1014            err,
1015            SchemaError::EmbeddingSourceOnNonEmbedding(_, _)
1016        ));
1017    }
1018
1019    #[test]
1020    fn generated_embedding_requires_dim() {
1021        let mut emb = Column::new(2, "vec", ColumnType::Embedding);
1022        emb.embedding_source = Some(EmbeddingSource::GeneratedColumn {
1023            provider: "local-test".into(),
1024        });
1025        let err = Schema::new(vec![Table {
1026            id: 1,
1027            name: "t".into(),
1028            columns: vec![Column::new(1, "id", ColumnType::Int64), emb],
1029            primary_key: vec!["id".into()],
1030            indexes: vec![],
1031            foreign_keys: vec![],
1032            unique_constraints: vec![],
1033            check_constraints: vec![],
1034        }])
1035        .unwrap_err();
1036        assert!(matches!(err, SchemaError::EmbeddingSourceMissingDim(_, _)));
1037    }
1038
1039    #[test]
1040    fn generated_embedding_spec_roundtrips_and_validates_sources() {
1041        let mut emb = Column::new(3, "vec", ColumnType::Embedding);
1042        emb.embedding_dim = Some(4);
1043        emb.embedding_source = Some(EmbeddingSource::GeneratedColumnSpec {
1044            spec: GeneratedEmbeddingSpec {
1045                provider_id: "provider".into(),
1046                model_id: "model".into(),
1047                model_version: "1".into(),
1048                source_columns: vec![2],
1049                input_template: "{body}".into(),
1050                dimension: 4,
1051                normalization: EmbeddingSpecNormalization::None,
1052                failure_policy: EmbeddingWriteFailurePolicy::AbortWrite,
1053            },
1054        });
1055        let schema = Schema::new(vec![Table {
1056            id: 1,
1057            name: "docs".into(),
1058            columns: vec![
1059                Column::new(1, "id", ColumnType::Int64),
1060                Column::new(2, "body", ColumnType::Text),
1061                emb,
1062            ],
1063            primary_key: vec!["id".into()],
1064            indexes: vec![],
1065            foreign_keys: vec![],
1066            unique_constraints: vec![],
1067            check_constraints: vec![],
1068        }])
1069        .unwrap();
1070        let json = serde_json::to_string(&schema).unwrap();
1071        assert!(json.contains("generated_column_spec"));
1072        assert_eq!(serde_json::from_str::<Schema>(&json).unwrap(), schema);
1073    }
1074}