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    /// Product-quantizer training sample cap. Engine default when omitted.
304    #[serde(default, skip_serializing_if = "Option::is_none")]
305    pub ann_pq_training_samples: Option<usize>,
306    /// Product-quantizer deterministic training seed. Engine default when omitted.
307    #[serde(default, skip_serializing_if = "Option::is_none")]
308    pub ann_pq_seed: Option<u64>,
309    /// Product-quantizer exact-rerank factor (0 disables). Engine default when omitted.
310    #[serde(default, skip_serializing_if = "Option::is_none")]
311    pub ann_pq_rerank_factor: Option<usize>,
312    /// MinHash permutation count. Engine default when omitted.
313    #[serde(default, skip_serializing_if = "Option::is_none")]
314    pub minhash_permutations: Option<usize>,
315    /// MinHash LSH band count. Engine default when omitted.
316    #[serde(default, skip_serializing_if = "Option::is_none")]
317    pub minhash_bands: Option<usize>,
318    /// Learned-range error bound. Engine default when omitted.
319    #[serde(default, skip_serializing_if = "Option::is_none")]
320    pub learned_range_epsilon: Option<usize>,
321}
322
323/// A uniqueness constraint over one or more columns.
324#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
325pub struct UniqueConstraint {
326    pub name: String,
327    pub columns: Vec<String>,
328}
329
330/// A foreign-key reference from child columns to parent columns.
331#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
332pub struct ForeignKey {
333    pub name: String,
334    pub columns: Vec<String>,
335    pub references_table: String,
336    pub references_columns: Vec<String>,
337    #[serde(default)]
338    pub on_delete: ForeignKeyAction,
339}
340
341/// Action taken when a referenced parent row is deleted.
342#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, Serialize, Deserialize)]
343#[serde(rename_all = "snake_case")]
344pub enum ForeignKeyAction {
345    #[default]
346    Restrict,
347    Cascade,
348    SetNull,
349}
350
351/// A named table-level check constraint.
352#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
353pub struct CheckConstraint {
354    pub name: String,
355    pub expr: String,
356}
357
358/// A monotonic sequence allocator.
359#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
360pub struct Sequence {
361    pub name: String,
362    pub next_value: i64,
363}
364
365/// A table definition.
366#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
367pub struct Table {
368    /// Stable table identifier. IDs must be unique within a schema.
369    pub id: u32,
370    pub name: String,
371    pub columns: Vec<Column>,
372    pub primary_key: Vec<String>,
373    #[serde(default, skip_serializing_if = "Vec::is_empty")]
374    pub indexes: Vec<Index>,
375    #[serde(default, skip_serializing_if = "Vec::is_empty")]
376    pub foreign_keys: Vec<ForeignKey>,
377    #[serde(default, skip_serializing_if = "Vec::is_empty")]
378    pub unique_constraints: Vec<UniqueConstraint>,
379    #[serde(default, skip_serializing_if = "Vec::is_empty")]
380    pub check_constraints: Vec<CheckConstraint>,
381}
382
383impl Table {
384    /// Find a column by name.
385    pub fn column(&self, name: &str) -> Option<&Column> {
386        self.columns.iter().find(|c| c.name == name)
387    }
388
389    /// Whether the named column is part of the primary key.
390    pub fn is_pk_column(&self, name: &str) -> bool {
391        self.primary_key.iter().any(|c| c == name)
392    }
393}
394
395/// Errors that can occur while constructing a [`Schema`].
396#[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)]
397pub enum SchemaError {
398    #[error("duplicate table name \"{0}\"")]
399    DuplicateTableName(String),
400    #[error("duplicate table id {0}")]
401    DuplicateTableId(u32),
402    #[error("duplicate column name \"{1}\" in table \"{0}\"")]
403    DuplicateColumnName(String, String),
404    #[error("duplicate column id {1} in table \"{0}\"")]
405    DuplicateColumnId(String, u32),
406    #[error("primary key column \"{1}\" not found in table \"{0}\"")]
407    MissingPrimaryKeyColumn(String, String),
408    #[error("index \"{1}\" references unknown column \"{2}\" in table \"{0}\"")]
409    MissingIndexColumn(String, String, String),
410    #[error("unique constraint \"{1}\" references unknown column \"{2}\" in table \"{0}\"")]
411    MissingUniqueColumn(String, String, String),
412    #[error("foreign key \"{1}\" references unknown column \"{2}\" in table \"{0}\"")]
413    MissingForeignKeyColumn(String, String, String),
414    #[error("foreign key \"{1}\" references unknown table \"{2}\"")]
415    MissingReferencedTable(String, String, String),
416    #[error("foreign key \"{1}\" references unknown column \"{2}\" on table \"{3}\"")]
417    MissingReferencedColumn(String, String, String, String),
418    #[error(
419        "column \"{1}\" on table \"{0}\" sets embedding_source but is not an embedding column"
420    )]
421    EmbeddingSourceOnNonEmbedding(String, String),
422    #[error(
423        "embedding column \"{1}\" on table \"{0}\" with LocalModel/GeneratedColumn source requires embedding_dim > 0"
424    )]
425    EmbeddingSourceMissingDim(String, String),
426    #[error("generated embedding column \"{1}\" on table \"{0}\" is invalid: {2}")]
427    InvalidGeneratedEmbeddingSpec(String, String, String),
428}
429
430/// A validated collection of tables.
431#[derive(Debug, Clone, PartialEq, Serialize)]
432pub struct Schema {
433    pub tables: Vec<Table>,
434    by_name: HashMap<String, usize>,
435    by_id: HashMap<u32, usize>,
436}
437
438impl<'de> serde::Deserialize<'de> for Schema {
439    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
440    where
441        D: serde::Deserializer<'de>,
442    {
443        #[derive(serde::Deserialize)]
444        struct SchemaHelper {
445            tables: Vec<Table>,
446        }
447        let helper = SchemaHelper::deserialize(deserializer)?;
448        Schema::new(helper.tables).map_err(serde::de::Error::custom)
449    }
450}
451
452/// A unique index also enforces uniqueness (guard-backed), matching SQL where a
453/// UNIQUE index is a UNIQUE constraint. Synthesize a constraint for each unique
454/// index unless an existing (or already-synthesized) unique constraint already
455/// covers exactly the same columns. Mirrors the TypeScript kit's `table()`.
456fn synthesize_unique_from_indexes(table: &mut Table) {
457    let mut synthesized: Vec<UniqueConstraint> = Vec::new();
458    for idx in &table.indexes {
459        if !idx.unique {
460            continue;
461        }
462        let covered = table
463            .unique_constraints
464            .iter()
465            .chain(synthesized.iter())
466            .any(|u| u.columns == idx.columns);
467        if !covered {
468            synthesized.push(UniqueConstraint {
469                name: idx.name.clone(),
470                columns: idx.columns.clone(),
471            });
472        }
473    }
474    table.unique_constraints.extend(synthesized);
475}
476
477impl Schema {
478    /// Build and validate a schema from a list of tables.
479    pub fn new(mut tables: Vec<Table>) -> Result<Self, SchemaError> {
480        for table in &mut tables {
481            synthesize_unique_from_indexes(table);
482        }
483
484        let mut by_name = HashMap::with_capacity(tables.len());
485        let mut by_id = HashMap::with_capacity(tables.len());
486
487        for (idx, table) in tables.iter().enumerate() {
488            if by_name.contains_key(&table.name) {
489                return Err(SchemaError::DuplicateTableName(table.name.clone()));
490            }
491            if by_id.contains_key(&table.id) {
492                return Err(SchemaError::DuplicateTableId(table.id));
493            }
494            by_name.insert(table.name.clone(), idx);
495            by_id.insert(table.id, idx);
496        }
497
498        for table in &tables {
499            Self::validate_table(table, &by_name)?;
500        }
501
502        Ok(Self {
503            tables,
504            by_name,
505            by_id,
506        })
507    }
508
509    fn validate_table(
510        table: &Table,
511        table_names: &HashMap<String, usize>,
512    ) -> Result<(), SchemaError> {
513        let mut column_names = HashMap::with_capacity(table.columns.len());
514        let mut column_ids = HashMap::with_capacity(table.columns.len());
515
516        for col in &table.columns {
517            if column_names.contains_key(&col.name) {
518                return Err(SchemaError::DuplicateColumnName(
519                    table.name.clone(),
520                    col.name.clone(),
521                ));
522            }
523            if column_ids.contains_key(&col.id) {
524                return Err(SchemaError::DuplicateColumnId(table.name.clone(), col.id));
525            }
526            if col.embedding_source.is_some() && col.storage_type != ColumnType::Embedding {
527                return Err(SchemaError::EmbeddingSourceOnNonEmbedding(
528                    table.name.clone(),
529                    col.name.clone(),
530                ));
531            }
532            if matches!(
533                col.embedding_source,
534                Some(
535                    EmbeddingSource::LocalModel { .. }
536                        | EmbeddingSource::ConfiguredModel { .. }
537                        | EmbeddingSource::GeneratedColumn { .. }
538                        | EmbeddingSource::GeneratedColumnSpec { .. }
539                )
540            ) && col.embedding_dim.unwrap_or(0) == 0
541            {
542                return Err(SchemaError::EmbeddingSourceMissingDim(
543                    table.name.clone(),
544                    col.name.clone(),
545                ));
546            }
547            column_names.insert(col.name.clone(), col.id);
548            column_ids.insert(col.id, col.name.clone());
549        }
550
551        for col in &table.columns {
552            if let Some(EmbeddingSource::ConfiguredModel {
553                provider_id,
554                model_id,
555                model_version,
556            }) = col.embedding_source.as_ref()
557            {
558                if provider_id.is_empty() || model_id.is_empty() || model_version.is_empty() {
559                    return Err(SchemaError::InvalidGeneratedEmbeddingSpec(
560                        table.name.clone(),
561                        col.name.clone(),
562                        "provider, model, and version are required".into(),
563                    ));
564                }
565            }
566        }
567
568        for col in &table.columns {
569            let Some(EmbeddingSource::GeneratedColumnSpec { spec }) = col.embedding_source.as_ref()
570            else {
571                continue;
572            };
573            if spec.provider_id.is_empty()
574                || spec.model_id.is_empty()
575                || spec.model_version.is_empty()
576                || spec.source_columns.is_empty()
577            {
578                return Err(SchemaError::InvalidGeneratedEmbeddingSpec(
579                    table.name.clone(),
580                    col.name.clone(),
581                    "provider, model, version, and source columns are required".into(),
582                ));
583            }
584            if col.embedding_dim != Some(spec.dimension) {
585                return Err(SchemaError::InvalidGeneratedEmbeddingSpec(
586                    table.name.clone(),
587                    col.name.clone(),
588                    "spec dimension must match embedding_dim".into(),
589                ));
590            }
591            let mut seen = HashSet::new();
592            if spec.source_columns.iter().any(|source_id| {
593                *source_id == col.id
594                    || !seen.insert(*source_id)
595                    || !column_ids.contains_key(source_id)
596            }) {
597                return Err(SchemaError::InvalidGeneratedEmbeddingSpec(
598                    table.name.clone(),
599                    col.name.clone(),
600                    "source columns must exist, be unique, and exclude the target".into(),
601                ));
602            }
603        }
604
605        for pk in &table.primary_key {
606            if !column_names.contains_key(pk) {
607                return Err(SchemaError::MissingPrimaryKeyColumn(
608                    table.name.clone(),
609                    pk.clone(),
610                ));
611            }
612        }
613
614        for idx in &table.indexes {
615            for col in &idx.columns {
616                if !column_names.contains_key(col) {
617                    return Err(SchemaError::MissingIndexColumn(
618                        table.name.clone(),
619                        idx.name.clone(),
620                        col.clone(),
621                    ));
622                }
623            }
624        }
625
626        for uq in &table.unique_constraints {
627            for col in &uq.columns {
628                if !column_names.contains_key(col) {
629                    return Err(SchemaError::MissingUniqueColumn(
630                        table.name.clone(),
631                        uq.name.clone(),
632                        col.clone(),
633                    ));
634                }
635            }
636        }
637
638        for fk in &table.foreign_keys {
639            for col in &fk.columns {
640                if !column_names.contains_key(col) {
641                    return Err(SchemaError::MissingForeignKeyColumn(
642                        table.name.clone(),
643                        fk.name.clone(),
644                        col.clone(),
645                    ));
646                }
647            }
648            if !table_names.contains_key(&fk.references_table) {
649                return Err(SchemaError::MissingReferencedTable(
650                    table.name.clone(),
651                    fk.name.clone(),
652                    fk.references_table.clone(),
653                ));
654            }
655        }
656
657        Ok(())
658    }
659
660    /// Look up a table by name.
661    pub fn table(&self, name: &str) -> Option<&Table> {
662        self.by_name.get(name).map(|&idx| &self.tables[idx])
663    }
664
665    /// Look up a table by stable id.
666    pub fn table_by_id(&self, id: u32) -> Option<&Table> {
667        self.by_id.get(&id).map(|&idx| &self.tables[idx])
668    }
669
670    /// Whether the schema contains a table with the given name.
671    pub fn has_table(&self, name: &str) -> bool {
672        self.by_name.contains_key(name)
673    }
674
675    /// Rename a table in place, keeping the `by_name` index in sync. Returns
676    /// `false` if `from` is absent or `to` is already in use (no change made).
677    /// Does *not* retarget foreign keys — callers that need that should do it
678    /// before/after on the tables they own.
679    pub fn rename_table(&mut self, from: &str, to: &str) -> bool {
680        if from == to {
681            return self.has_table(from);
682        }
683        if !self.has_table(from) || self.has_table(to) {
684            return false;
685        }
686        let idx = *self.by_name.get(from).unwrap();
687        self.tables[idx].name = to.to_string();
688        self.by_name.remove(from);
689        self.by_name.insert(to.to_string(), idx);
690        true
691    }
692}
693
694#[cfg(test)]
695mod tests {
696    use super::*;
697
698    fn make_table(name: &str, id: u32) -> Table {
699        Table {
700            id,
701            name: name.into(),
702            columns: vec![Column::new(1, "id", ColumnType::Int64)],
703            primary_key: vec!["id".into()],
704            indexes: vec![],
705            foreign_keys: vec![],
706            unique_constraints: vec![],
707            check_constraints: vec![],
708        }
709    }
710
711    #[test]
712    fn schema_rejects_duplicate_table_name() {
713        let err = Schema::new(vec![make_table("a", 1), make_table("a", 2)]).unwrap_err();
714        assert!(matches!(err, SchemaError::DuplicateTableName(n) if n == "a"));
715    }
716
717    #[test]
718    fn schema_rejects_duplicate_table_id() {
719        let err = Schema::new(vec![make_table("a", 1), make_table("b", 1)]).unwrap_err();
720        assert!(matches!(err, SchemaError::DuplicateTableId(1)));
721    }
722
723    #[test]
724    fn schema_rejects_missing_pk_column() {
725        let t = Table {
726            id: 1,
727            name: "t".into(),
728            columns: vec![Column::new(1, "x", ColumnType::Text)],
729            primary_key: vec!["id".into()],
730            indexes: vec![],
731            foreign_keys: vec![],
732            unique_constraints: vec![],
733            check_constraints: vec![],
734        };
735        let err = Schema::new(vec![t]).unwrap_err();
736        assert!(matches!(err, SchemaError::MissingPrimaryKeyColumn(_, _)));
737    }
738
739    #[test]
740    fn unique_index_synthesizes_unique_constraint() {
741        let schema = Schema::new(vec![Table {
742            id: 1,
743            name: "users".into(),
744            columns: vec![
745                Column::new(1, "id", ColumnType::Int64),
746                Column::new(2, "email", ColumnType::Text),
747                Column::new(3, "handle", ColumnType::Text),
748            ],
749            primary_key: vec!["id".into()],
750            indexes: vec![
751                Index {
752                    name: "idx_email".into(),
753                    columns: vec!["email".into()],
754                    unique: true,
755                    kind: Default::default(),
756                    ann_quantization: Default::default(),
757                    ..Default::default()
758                },
759                // A non-unique index must NOT synthesize a constraint.
760                Index {
761                    name: "idx_handle".into(),
762                    columns: vec!["handle".into()],
763                    unique: false,
764                    kind: Default::default(),
765                    ann_quantization: Default::default(),
766                    ..Default::default()
767                },
768            ],
769            foreign_keys: vec![],
770            unique_constraints: vec![],
771            check_constraints: vec![],
772        }])
773        .unwrap();
774        let table = schema.table("users").unwrap();
775        assert_eq!(table.unique_constraints.len(), 1);
776        assert_eq!(
777            table.unique_constraints[0].columns,
778            vec!["email".to_string()]
779        );
780    }
781
782    #[test]
783    fn unique_index_does_not_duplicate_existing_constraint() {
784        let schema = Schema::new(vec![Table {
785            id: 1,
786            name: "users".into(),
787            columns: vec![
788                Column::new(1, "id", ColumnType::Int64),
789                Column::new(2, "email", ColumnType::Text),
790            ],
791            primary_key: vec!["id".into()],
792            indexes: vec![Index {
793                name: "idx_email".into(),
794                columns: vec!["email".into()],
795                unique: true,
796                kind: Default::default(),
797                ann_quantization: Default::default(),
798                ..Default::default()
799            }],
800            foreign_keys: vec![],
801            unique_constraints: vec![UniqueConstraint {
802                name: "uq_email".into(),
803                columns: vec!["email".into()],
804            }],
805            check_constraints: vec![],
806        }])
807        .unwrap();
808        // The pre-existing constraint already covers `email`; no synthesis.
809        let table = schema.table("users").unwrap();
810        assert_eq!(table.unique_constraints.len(), 1);
811        assert_eq!(table.unique_constraints[0].name, "uq_email");
812    }
813
814    #[test]
815    fn schema_roundtrips_json() {
816        let schema = Schema::new(vec![Table {
817            id: 1,
818            name: "users".into(),
819            columns: vec![
820                Column::new(1, "id", ColumnType::Int64),
821                Column {
822                    nullable: true,
823                    ..Column::new(2, "email", ColumnType::Text)
824                },
825            ],
826            primary_key: vec!["id".into()],
827            indexes: vec![Index {
828                name: "idx_email".into(),
829                columns: vec!["email".into()],
830                unique: true,
831                kind: Default::default(),
832                ann_quantization: Default::default(),
833                ..Default::default()
834            }],
835            foreign_keys: vec![],
836            unique_constraints: vec![],
837            check_constraints: vec![CheckConstraint {
838                name: "chk_id_positive".into(),
839                expr: "id > 0".into(),
840            }],
841        }])
842        .unwrap();
843
844        let json = serde_json::to_string(&schema).unwrap();
845        let decoded: Schema = serde_json::from_str(&json).unwrap();
846        assert_eq!(decoded.tables.len(), 1);
847        assert_eq!(decoded.table("users").unwrap().columns.len(), 2);
848    }
849
850    #[test]
851    fn dense_ann_quantization_roundtrips_and_old_json_defaults_binary() {
852        let dense = Index {
853            name: "idx_embedding".into(),
854            columns: vec!["embedding".into()],
855            unique: false,
856            kind: IndexKind::Ann,
857            ann_quantization: AnnQuantization::Dense,
858            predicate: Some("embedding IS NOT NULL".into()),
859            ann_m: Some(24),
860            ann_ef_construction: Some(96),
861            ann_ef_search: Some(48),
862            ..Default::default()
863        };
864        let json = serde_json::to_value(&dense).unwrap();
865        assert_eq!(json["ann_quantization"], "dense");
866        assert_eq!(json["ann_m"], 24);
867        assert_eq!(json["predicate"], "embedding IS NOT NULL");
868        assert_eq!(
869            serde_json::from_value::<Index>(json)
870                .unwrap()
871                .ann_quantization,
872            AnnQuantization::Dense
873        );
874
875        let old: Index = serde_json::from_value(serde_json::json!({
876            "name": "idx_embedding",
877            "columns": ["embedding"],
878            "unique": false,
879            "kind": "ann"
880        }))
881        .unwrap();
882        assert_eq!(old.ann_quantization, AnnQuantization::BinarySign);
883    }
884
885    #[test]
886    fn swappable_ann_algorithm_and_product_quantization_roundtrip() {
887        let diskann = Index {
888            name: "idx_diskann".into(),
889            columns: vec!["embedding".into()],
890            unique: false,
891            kind: IndexKind::Ann,
892            ann_algorithm: AnnAlgorithm::DiskAnn,
893            ann_quantization: AnnQuantization::Dense,
894            ann_diskann_r: Some(128),
895            ann_diskann_l: Some(256),
896            ann_diskann_beam_width: Some(4),
897            ann_diskann_alpha: Some(130),
898            ..Default::default()
899        };
900        let json = serde_json::to_value(&diskann).unwrap();
901        assert_eq!(json["ann_algorithm"], "diskann");
902        assert_eq!(json["ann_diskann_r"], 128);
903        assert_eq!(json["ann_diskann_alpha"], 130);
904        let decoded: Index = serde_json::from_value(json).unwrap();
905        assert_eq!(decoded.ann_algorithm, AnnAlgorithm::DiskAnn);
906        assert_eq!(decoded.ann_diskann_l, Some(256));
907
908        let ivf = Index {
909            name: "idx_ivf".into(),
910            columns: vec!["embedding".into()],
911            kind: IndexKind::Ann,
912            ann_algorithm: AnnAlgorithm::Ivf,
913            ann_quantization: AnnQuantization::Dense,
914            ann_ivf_nlist: Some(512),
915            ann_ivf_nprobe: Some(16),
916            ..Default::default()
917        };
918        let json = serde_json::to_value(&ivf).unwrap();
919        assert_eq!(json["ann_algorithm"], "ivf");
920        assert_eq!(json["ann_ivf_nlist"], 512);
921
922        let pq = Index {
923            name: "idx_pq".into(),
924            columns: vec!["embedding".into()],
925            kind: IndexKind::Ann,
926            ann_quantization: AnnQuantization::Product {
927                num_subvectors: 32,
928                bits: 8,
929            },
930            ann_pq_training_samples: Some(10_000),
931            ann_pq_seed: Some(42),
932            ann_pq_rerank_factor: Some(3),
933            ..Default::default()
934        };
935        let json = serde_json::to_value(&pq).unwrap();
936        assert_eq!(json["ann_quantization"]["product"]["num_subvectors"], 32);
937        assert_eq!(json["ann_quantization"]["product"]["bits"], 8);
938        assert_eq!(json["ann_pq_seed"], 42);
939        let decoded: Index = serde_json::from_value(json).unwrap();
940        assert_eq!(
941            decoded.ann_quantization,
942            AnnQuantization::Product {
943                num_subvectors: 32,
944                bits: 8
945            }
946        );
947
948        // Default algorithm is HNSW and is omitted on the wire (backward compat).
949        let hnsw = Index {
950            name: "idx_hnsw".into(),
951            columns: vec!["embedding".into()],
952            kind: IndexKind::Ann,
953            ann_quantization: AnnQuantization::Dense,
954            ..Default::default()
955        };
956        let json = serde_json::to_value(&hnsw).unwrap();
957        assert_eq!(json["ann_algorithm"], "hnsw");
958    }
959
960    #[test]
961    fn embedding_source_roundtrips_json() {
962        let mut emb = Column::new(2, "vec", ColumnType::Embedding);
963        emb.embedding_dim = Some(4);
964        emb.embedding_source = Some(EmbeddingSource::LocalModel {
965            model_path: "/models/kit-mini".into(),
966            model_id: "kit-mini".into(),
967        });
968        let schema = Schema::new(vec![Table {
969            id: 1,
970            name: "docs".into(),
971            columns: vec![Column::new(1, "id", ColumnType::Int64), emb],
972            primary_key: vec!["id".into()],
973            indexes: vec![],
974            foreign_keys: vec![],
975            unique_constraints: vec![],
976            check_constraints: vec![],
977        }])
978        .unwrap();
979        let json = serde_json::to_string(&schema).unwrap();
980        assert!(json.contains("local_model"));
981        assert!(json.contains("kit-mini"));
982        let decoded: Schema = serde_json::from_str(&json).unwrap();
983        let col = decoded.table("docs").unwrap().column("vec").unwrap();
984        assert_eq!(
985            col.embedding_source,
986            Some(EmbeddingSource::LocalModel {
987                model_path: "/models/kit-mini".into(),
988                model_id: "kit-mini".into(),
989            })
990        );
991    }
992
993    #[test]
994    fn embedding_source_rejected_on_non_embedding() {
995        let mut col = Column::new(2, "name", ColumnType::Text);
996        col.embedding_source = Some(EmbeddingSource::SuppliedByApplication);
997        let err = Schema::new(vec![Table {
998            id: 1,
999            name: "t".into(),
1000            columns: vec![Column::new(1, "id", ColumnType::Int64), col],
1001            primary_key: vec!["id".into()],
1002            indexes: vec![],
1003            foreign_keys: vec![],
1004            unique_constraints: vec![],
1005            check_constraints: vec![],
1006        }])
1007        .unwrap_err();
1008        assert!(matches!(
1009            err,
1010            SchemaError::EmbeddingSourceOnNonEmbedding(_, _)
1011        ));
1012    }
1013
1014    #[test]
1015    fn generated_embedding_requires_dim() {
1016        let mut emb = Column::new(2, "vec", ColumnType::Embedding);
1017        emb.embedding_source = Some(EmbeddingSource::GeneratedColumn {
1018            provider: "local-test".into(),
1019        });
1020        let err = Schema::new(vec![Table {
1021            id: 1,
1022            name: "t".into(),
1023            columns: vec![Column::new(1, "id", ColumnType::Int64), emb],
1024            primary_key: vec!["id".into()],
1025            indexes: vec![],
1026            foreign_keys: vec![],
1027            unique_constraints: vec![],
1028            check_constraints: vec![],
1029        }])
1030        .unwrap_err();
1031        assert!(matches!(err, SchemaError::EmbeddingSourceMissingDim(_, _)));
1032    }
1033
1034    #[test]
1035    fn generated_embedding_spec_roundtrips_and_validates_sources() {
1036        let mut emb = Column::new(3, "vec", ColumnType::Embedding);
1037        emb.embedding_dim = Some(4);
1038        emb.embedding_source = Some(EmbeddingSource::GeneratedColumnSpec {
1039            spec: GeneratedEmbeddingSpec {
1040                provider_id: "provider".into(),
1041                model_id: "model".into(),
1042                model_version: "1".into(),
1043                source_columns: vec![2],
1044                input_template: "{body}".into(),
1045                dimension: 4,
1046                normalization: EmbeddingSpecNormalization::None,
1047                failure_policy: EmbeddingWriteFailurePolicy::AbortWrite,
1048            },
1049        });
1050        let schema = Schema::new(vec![Table {
1051            id: 1,
1052            name: "docs".into(),
1053            columns: vec![
1054                Column::new(1, "id", ColumnType::Int64),
1055                Column::new(2, "body", ColumnType::Text),
1056                emb,
1057            ],
1058            primary_key: vec!["id".into()],
1059            indexes: vec![],
1060            foreign_keys: vec![],
1061            unique_constraints: vec![],
1062            check_constraints: vec![],
1063        }])
1064        .unwrap();
1065        let json = serde_json::to_string(&schema).unwrap();
1066        assert!(json.contains("generated_column_spec"));
1067        assert_eq!(serde_json::from_str::<Schema>(&json).unwrap(), schema);
1068    }
1069}