polyglot_sql/
schema.rs

1//! Schema management for SQL queries
2//!
3//! This module provides functionality for:
4//! - Representing database schemas (tables, columns, types)
5//! - Looking up column types for query optimization
6//! - Normalizing identifiers per dialect
7//!
8//! Based on the Python implementation in `sqlglot/schema.py`.
9
10use crate::dialects::DialectType;
11use crate::expressions::DataType;
12use crate::trie::{Trie, TrieResult};
13use std::collections::{HashMap, HashSet};
14use thiserror::Error;
15
16/// Errors that can occur during schema operations
17#[derive(Debug, Error, Clone)]
18pub enum SchemaError {
19    #[error("Table not found: {0}")]
20    TableNotFound(String),
21
22    #[error("Ambiguous table: {table} matches multiple tables: {matches}")]
23    AmbiguousTable { table: String, matches: String },
24
25    #[error("Column not found: {column} in table {table}")]
26    ColumnNotFound { table: String, column: String },
27
28    #[error("Schema nesting depth mismatch: expected {expected}, got {actual}")]
29    DepthMismatch { expected: usize, actual: usize },
30
31    #[error("Invalid schema structure: {0}")]
32    InvalidStructure(String),
33}
34
35/// Result type for schema operations
36pub type SchemaResult<T> = Result<T, SchemaError>;
37
38/// Supported table argument names
39pub const TABLE_PARTS: &[&str] = &["this", "db", "catalog"];
40
41/// Abstract trait for database schemas
42pub trait Schema {
43    /// Get the dialect associated with this schema (if any)
44    fn dialect(&self) -> Option<DialectType>;
45
46    /// Add or update a table in the schema
47    fn add_table(
48        &mut self,
49        table: &str,
50        columns: &[(String, DataType)],
51        dialect: Option<DialectType>,
52    ) -> SchemaResult<()>;
53
54    /// Get column names for a table
55    fn column_names(&self, table: &str) -> SchemaResult<Vec<String>>;
56
57    /// Get the type of a column in a table
58    fn get_column_type(&self, table: &str, column: &str) -> SchemaResult<DataType>;
59
60    /// Check if a column exists in a table
61    fn has_column(&self, table: &str, column: &str) -> bool;
62
63    /// Get supported table argument levels
64    fn supported_table_args(&self) -> &[&str];
65
66    /// Check if the schema is empty
67    fn is_empty(&self) -> bool;
68
69    /// Get the nesting depth of the schema
70    fn depth(&self) -> usize;
71
72    /// Find which table(s) contain the given column name.
73    /// Returns table names that have the column. Used for correlated subquery resolution.
74    fn find_tables_for_column(&self, column: &str) -> Vec<String>;
75}
76
77/// A column with its type and visibility
78#[derive(Debug, Clone)]
79pub struct ColumnInfo {
80    pub data_type: DataType,
81    pub visible: bool,
82}
83
84impl ColumnInfo {
85    pub fn new(data_type: DataType) -> Self {
86        Self {
87            data_type,
88            visible: true,
89        }
90    }
91
92    pub fn with_visibility(data_type: DataType, visible: bool) -> Self {
93        Self { data_type, visible }
94    }
95}
96
97/// A mapping-based schema implementation
98///
99/// Supports nested schemas with different levels:
100/// - Level 1: `{table: {col: type}}`
101/// - Level 2: `{db: {table: {col: type}}}`
102/// - Level 3: `{catalog: {db: {table: {col: type}}}}`
103#[derive(Debug, Clone)]
104pub struct MappingSchema {
105    /// The actual schema data
106    mapping: HashMap<String, SchemaNode>,
107    /// Trie for efficient table lookup
108    mapping_trie: Trie<()>,
109    /// The dialect for this schema
110    dialect: Option<DialectType>,
111    /// Whether to normalize identifiers
112    normalize: bool,
113    /// Visible columns per table
114    visible: HashMap<String, HashSet<String>>,
115    /// Declared column order per normalized table path.
116    column_order: HashMap<String, Vec<String>>,
117    /// Cached depth
118    cached_depth: usize,
119}
120
121/// A node in the schema tree
122#[derive(Debug, Clone)]
123pub enum SchemaNode {
124    /// Intermediate node (database or catalog)
125    Namespace(HashMap<String, SchemaNode>),
126    /// Leaf node (table with columns)
127    Table(HashMap<String, ColumnInfo>),
128}
129
130impl Default for MappingSchema {
131    fn default() -> Self {
132        Self::new()
133    }
134}
135
136impl MappingSchema {
137    /// Create a new empty schema
138    pub fn new() -> Self {
139        Self {
140            mapping: HashMap::new(),
141            mapping_trie: Trie::new(),
142            dialect: None,
143            normalize: true,
144            visible: HashMap::new(),
145            column_order: HashMap::new(),
146            cached_depth: 0,
147        }
148    }
149
150    /// Create a schema with a specific dialect
151    pub fn with_dialect(dialect: DialectType) -> Self {
152        Self {
153            dialect: Some(dialect),
154            ..Self::new()
155        }
156    }
157
158    /// Create a schema with normalization disabled
159    pub fn without_normalization(mut self) -> Self {
160        self.normalize = false;
161        self
162    }
163
164    /// Set visibility for columns in a table
165    pub fn set_visible_columns(&mut self, table: &str, columns: &[&str]) {
166        let key = self.normalize_name(table, true);
167        let cols: HashSet<String> = columns
168            .iter()
169            .map(|c| self.normalize_name(c, false))
170            .collect();
171        self.visible.insert(key, cols);
172    }
173
174    /// Normalize an identifier name based on dialect
175    fn normalize_name(&self, name: &str, is_table: bool) -> String {
176        if !self.normalize {
177            return name.to_string();
178        }
179
180        // Default normalization: lowercase
181        // Different dialects may have different rules
182        match self.dialect {
183            Some(DialectType::BigQuery) if is_table => {
184                // BigQuery preserves case for tables
185                name.to_string()
186            }
187            Some(DialectType::Snowflake) => {
188                // Snowflake uppercases by default
189                name.to_uppercase()
190            }
191            _ => {
192                // Most dialects lowercase
193                name.to_lowercase()
194            }
195        }
196    }
197
198    /// Parse a qualified table name into parts
199    fn parse_table_parts(&self, table: &str) -> Vec<String> {
200        table
201            .split('.')
202            .map(|s| self.normalize_name(s.trim(), true))
203            .collect()
204    }
205
206    /// Get the column mapping for a table
207    fn find_table(&self, table: &str) -> SchemaResult<&HashMap<String, ColumnInfo>> {
208        let parts = self.parse_table_parts(table);
209
210        // Use trie to find table
211        let reversed_parts: Vec<_> = parts.iter().rev().map(|s| s.as_str()).collect();
212        let key: String = reversed_parts.join(".");
213
214        let (result, _) = self.mapping_trie.in_trie(&key);
215
216        match result {
217            TrieResult::Failed => Err(SchemaError::TableNotFound(table.to_string())),
218            TrieResult::Prefix => {
219                // Ambiguous - multiple tables match
220                Err(SchemaError::AmbiguousTable {
221                    table: table.to_string(),
222                    matches: "multiple matches".to_string(),
223                })
224            }
225            TrieResult::Exists => {
226                // Navigate to the table
227                self.navigate_to_table(&parts)
228            }
229        }
230    }
231
232    /// Navigate the schema tree to find a table's columns
233    fn navigate_to_table(&self, parts: &[String]) -> SchemaResult<&HashMap<String, ColumnInfo>> {
234        let mut current = &self.mapping;
235
236        for (i, part) in parts.iter().enumerate() {
237            match current.get(part) {
238                Some(SchemaNode::Namespace(inner)) => {
239                    current = inner;
240                }
241                Some(SchemaNode::Table(cols)) => {
242                    if i == parts.len() - 1 {
243                        return Ok(cols);
244                    } else {
245                        return Err(SchemaError::InvalidStructure(format!(
246                            "Found table at {} but expected more levels",
247                            parts[..=i].join(".")
248                        )));
249                    }
250                }
251                None => {
252                    return Err(SchemaError::TableNotFound(parts.join(".")));
253                }
254            }
255        }
256
257        // We've exhausted parts but didn't find a table
258        Err(SchemaError::TableNotFound(parts.join(".")))
259    }
260
261    /// Add a table to the schema
262    fn add_table_internal(
263        &mut self,
264        parts: &[String],
265        columns: HashMap<String, ColumnInfo>,
266    ) -> SchemaResult<()> {
267        if parts.is_empty() {
268            return Err(SchemaError::InvalidStructure(
269                "Table name cannot be empty".to_string(),
270            ));
271        }
272
273        // Build trie key (reversed parts)
274        let trie_key: String = parts.iter().rev().cloned().collect::<Vec<_>>().join(".");
275        self.mapping_trie.insert(&trie_key, ());
276
277        // Navigate/create path to table
278        let mut current = &mut self.mapping;
279
280        for (i, part) in parts.iter().enumerate() {
281            let is_last = i == parts.len() - 1;
282
283            if is_last {
284                // Insert table
285                current.insert(part.clone(), SchemaNode::Table(columns));
286                return Ok(());
287            } else {
288                // Navigate or create namespace
289                let entry = current
290                    .entry(part.clone())
291                    .or_insert_with(|| SchemaNode::Namespace(HashMap::new()));
292
293                match entry {
294                    SchemaNode::Namespace(inner) => {
295                        current = inner;
296                    }
297                    SchemaNode::Table(_) => {
298                        return Err(SchemaError::InvalidStructure(format!(
299                            "Expected namespace at {} but found table",
300                            parts[..=i].join(".")
301                        )));
302                    }
303                }
304            }
305        }
306
307        Ok(())
308    }
309
310    /// Update cached depth
311    fn update_depth(&mut self) {
312        self.cached_depth = self.calculate_depth(&self.mapping);
313    }
314
315    fn calculate_depth(&self, mapping: &HashMap<String, SchemaNode>) -> usize {
316        if mapping.is_empty() {
317            return 0;
318        }
319
320        let mut max_depth = 1;
321        for node in mapping.values() {
322            match node {
323                SchemaNode::Namespace(inner) => {
324                    let d = 1 + self.calculate_depth(inner);
325                    if d > max_depth {
326                        max_depth = d;
327                    }
328                }
329                SchemaNode::Table(_) => {
330                    // Tables don't add to depth beyond their level
331                }
332            }
333        }
334        max_depth
335    }
336}
337
338impl Schema for MappingSchema {
339    fn dialect(&self) -> Option<DialectType> {
340        self.dialect
341    }
342
343    fn add_table(
344        &mut self,
345        table: &str,
346        columns: &[(String, DataType)],
347        _dialect: Option<DialectType>,
348    ) -> SchemaResult<()> {
349        let parts = self.parse_table_parts(table);
350
351        let cols: HashMap<String, ColumnInfo> = columns
352            .iter()
353            .map(|(name, dtype)| {
354                let normalized_name = self.normalize_name(name, false);
355                (normalized_name, ColumnInfo::new(dtype.clone()))
356            })
357            .collect();
358        let column_order: Vec<String> = columns
359            .iter()
360            .map(|(name, _)| self.normalize_name(name, false))
361            .collect();
362
363        self.add_table_internal(&parts, cols)?;
364        self.column_order.insert(parts.join("."), column_order);
365        self.update_depth();
366        Ok(())
367    }
368
369    fn column_names(&self, table: &str) -> SchemaResult<Vec<String>> {
370        let cols = self.find_table(table)?;
371        let table_key = self.normalize_name(table, true);
372        let ordered_columns = self
373            .column_order
374            .get(&self.parse_table_parts(table).join("."));
375
376        // Check visibility
377        if let Some(visible_cols) = self.visible.get(&table_key) {
378            Ok(ordered_columns
379                .map(|columns| {
380                    columns
381                        .iter()
382                        .filter(|column| {
383                            cols.contains_key(*column) && visible_cols.contains(*column)
384                        })
385                        .cloned()
386                        .collect()
387                })
388                .unwrap_or_else(|| {
389                    cols.keys()
390                        .filter(|column| visible_cols.contains(*column))
391                        .cloned()
392                        .collect()
393                }))
394        } else if let Some(columns) = ordered_columns {
395            Ok(columns
396                .iter()
397                .filter(|column| cols.contains_key(*column))
398                .cloned()
399                .collect())
400        } else {
401            Ok(cols.keys().cloned().collect())
402        }
403    }
404
405    fn get_column_type(&self, table: &str, column: &str) -> SchemaResult<DataType> {
406        let cols = self.find_table(table)?;
407        let normalized_col = self.normalize_name(column, false);
408
409        cols.get(&normalized_col)
410            .map(|info| info.data_type.clone())
411            .ok_or_else(|| SchemaError::ColumnNotFound {
412                table: table.to_string(),
413                column: column.to_string(),
414            })
415    }
416
417    fn has_column(&self, table: &str, column: &str) -> bool {
418        self.get_column_type(table, column).is_ok()
419    }
420
421    fn supported_table_args(&self) -> &[&str] {
422        let depth = self.depth();
423        if depth == 0 {
424            &[]
425        } else if depth <= 3 {
426            &TABLE_PARTS[..depth]
427        } else {
428            TABLE_PARTS
429        }
430    }
431
432    fn is_empty(&self) -> bool {
433        self.mapping.is_empty()
434    }
435
436    fn depth(&self) -> usize {
437        self.cached_depth
438    }
439
440    fn find_tables_for_column(&self, column: &str) -> Vec<String> {
441        let normalized = normalize_name(column, self.dialect, false, self.normalize);
442        let mut result = Vec::new();
443        for table_name in self.mapping.keys() {
444            if self.has_column(table_name, &normalized) {
445                result.push(table_name.clone());
446            }
447        }
448        result
449    }
450}
451
452/// Normalize a table or column name according to dialect rules
453pub fn normalize_name(
454    name: &str,
455    dialect: Option<DialectType>,
456    is_table: bool,
457    normalize: bool,
458) -> String {
459    if !normalize {
460        return name.to_string();
461    }
462
463    match dialect {
464        Some(DialectType::BigQuery) if is_table => name.to_string(),
465        Some(DialectType::Snowflake) => name.to_uppercase(),
466        _ => name.to_lowercase(),
467    }
468}
469
470/// Ensure we have a schema instance
471pub fn ensure_schema(schema: Option<MappingSchema>) -> MappingSchema {
472    schema.unwrap_or_default()
473}
474
475/// Helper to build a schema from a simple map
476///
477/// # Example
478///
479/// ```
480/// use polyglot_sql::schema::{MappingSchema, Schema, from_simple_map};
481/// use polyglot_sql::expressions::DataType;
482///
483/// let schema = from_simple_map(&[
484///     ("users", &[("id", DataType::Int { length: None, integer_spelling: false }), ("name", DataType::VarChar { length: Some(255), parenthesized_length: false })]),
485///     ("orders", &[("id", DataType::Int { length: None, integer_spelling: false }), ("user_id", DataType::Int { length: None, integer_spelling: false })]),
486/// ]);
487///
488/// assert_eq!(schema.column_names("users").unwrap().len(), 2);
489/// ```
490pub fn from_simple_map(tables: &[(&str, &[(&str, DataType)])]) -> MappingSchema {
491    let mut schema = MappingSchema::new();
492
493    for (table_name, columns) in tables {
494        let cols: Vec<(String, DataType)> = columns
495            .iter()
496            .map(|(name, dtype)| (name.to_string(), dtype.clone()))
497            .collect();
498
499        schema.add_table(table_name, &cols, None).ok();
500    }
501
502    schema
503}
504
505/// Flatten a nested schema to get all table paths
506pub fn flatten_schema_paths(schema: &MappingSchema) -> Vec<Vec<String>> {
507    let mut paths = Vec::new();
508    flatten_schema_paths_recursive(&schema.mapping, Vec::new(), &mut paths);
509    paths
510}
511
512fn flatten_schema_paths_recursive(
513    mapping: &HashMap<String, SchemaNode>,
514    prefix: Vec<String>,
515    paths: &mut Vec<Vec<String>>,
516) {
517    for (key, node) in mapping {
518        let mut path = prefix.clone();
519        path.push(key.clone());
520
521        match node {
522            SchemaNode::Namespace(inner) => {
523                flatten_schema_paths_recursive(inner, path, paths);
524            }
525            SchemaNode::Table(_) => {
526                paths.push(path);
527            }
528        }
529    }
530}
531
532/// Set a value in a nested dictionary-like structure
533pub fn nested_set<V: Clone>(
534    map: &mut HashMap<String, HashMap<String, V>>,
535    keys: &[String],
536    value: V,
537) {
538    if keys.is_empty() {
539        return;
540    }
541
542    if keys.len() == 1 {
543        // Can't set at single level - need at least 2 keys
544        return;
545    }
546
547    let outer_key = &keys[0];
548    let inner_key = &keys[1];
549
550    map.entry(outer_key.clone())
551        .or_insert_with(HashMap::new)
552        .insert(inner_key.clone(), value);
553}
554
555/// Get a value from a nested dictionary-like structure
556pub fn nested_get<'a, V>(
557    map: &'a HashMap<String, HashMap<String, V>>,
558    keys: &[String],
559) -> Option<&'a V> {
560    if keys.len() != 2 {
561        return None;
562    }
563
564    map.get(&keys[0])?.get(&keys[1])
565}
566
567#[cfg(test)]
568mod tests {
569    use super::*;
570
571    #[test]
572    fn test_empty_schema() {
573        let schema = MappingSchema::new();
574        assert!(schema.is_empty());
575        assert_eq!(schema.depth(), 0);
576    }
577
578    #[test]
579    fn test_add_table() {
580        let mut schema = MappingSchema::new();
581        let columns = vec![
582            (
583                "id".to_string(),
584                DataType::Int {
585                    length: None,
586                    integer_spelling: false,
587                },
588            ),
589            (
590                "name".to_string(),
591                DataType::VarChar {
592                    length: Some(255),
593                    parenthesized_length: false,
594                },
595            ),
596        ];
597
598        schema.add_table("users", &columns, None).unwrap();
599
600        assert!(!schema.is_empty());
601        assert_eq!(schema.depth(), 1);
602        assert!(schema.has_column("users", "id"));
603        assert!(schema.has_column("users", "name"));
604        assert!(!schema.has_column("users", "email"));
605    }
606
607    #[test]
608    fn test_qualified_table_names() {
609        let mut schema = MappingSchema::new();
610        let columns = vec![(
611            "id".to_string(),
612            DataType::Int {
613                length: None,
614                integer_spelling: false,
615            },
616        )];
617
618        schema.add_table("mydb.users", &columns, None).unwrap();
619
620        assert!(schema.has_column("mydb.users", "id"));
621        assert_eq!(schema.depth(), 2);
622    }
623
624    #[test]
625    fn test_catalog_db_table() {
626        let mut schema = MappingSchema::new();
627        let columns = vec![(
628            "id".to_string(),
629            DataType::Int {
630                length: None,
631                integer_spelling: false,
632            },
633        )];
634
635        schema
636            .add_table("catalog.mydb.users", &columns, None)
637            .unwrap();
638
639        assert!(schema.has_column("catalog.mydb.users", "id"));
640        assert_eq!(schema.depth(), 3);
641    }
642
643    #[test]
644    fn test_get_column_type() {
645        let mut schema = MappingSchema::new();
646        let columns = vec![
647            (
648                "id".to_string(),
649                DataType::Int {
650                    length: None,
651                    integer_spelling: false,
652                },
653            ),
654            (
655                "name".to_string(),
656                DataType::VarChar {
657                    length: Some(255),
658                    parenthesized_length: false,
659                },
660            ),
661        ];
662
663        schema.add_table("users", &columns, None).unwrap();
664
665        let id_type = schema.get_column_type("users", "id").unwrap();
666        assert!(matches!(id_type, DataType::Int { .. }));
667
668        let name_type = schema.get_column_type("users", "name").unwrap();
669        assert!(matches!(
670            name_type,
671            DataType::VarChar {
672                length: Some(255),
673                parenthesized_length: false
674            }
675        ));
676    }
677
678    #[test]
679    fn test_column_names() {
680        let mut schema = MappingSchema::new();
681        let columns = vec![
682            (
683                "id".to_string(),
684                DataType::Int {
685                    length: None,
686                    integer_spelling: false,
687                },
688            ),
689            (
690                "name".to_string(),
691                DataType::VarChar {
692                    length: None,
693                    parenthesized_length: false,
694                },
695            ),
696        ];
697
698        schema.add_table("users", &columns, None).unwrap();
699
700        let names = schema.column_names("users").unwrap();
701        assert_eq!(names, vec!["id", "name"]);
702    }
703
704    #[test]
705    fn test_table_not_found() {
706        let schema = MappingSchema::new();
707        let result = schema.column_names("nonexistent");
708        assert!(matches!(result, Err(SchemaError::TableNotFound(_))));
709    }
710
711    #[test]
712    fn test_column_not_found() {
713        let mut schema = MappingSchema::new();
714        let columns = vec![(
715            "id".to_string(),
716            DataType::Int {
717                length: None,
718                integer_spelling: false,
719            },
720        )];
721        schema.add_table("users", &columns, None).unwrap();
722
723        let result = schema.get_column_type("users", "nonexistent");
724        assert!(matches!(result, Err(SchemaError::ColumnNotFound { .. })));
725    }
726
727    #[test]
728    fn test_normalize_name_default() {
729        let name = normalize_name("MyTable", None, true, true);
730        assert_eq!(name, "mytable");
731    }
732
733    #[test]
734    fn test_normalize_name_snowflake() {
735        let name = normalize_name("MyTable", Some(DialectType::Snowflake), true, true);
736        assert_eq!(name, "MYTABLE");
737    }
738
739    #[test]
740    fn test_normalize_disabled() {
741        let name = normalize_name("MyTable", None, true, false);
742        assert_eq!(name, "MyTable");
743    }
744
745    #[test]
746    fn test_from_simple_map() {
747        let schema = from_simple_map(&[
748            (
749                "users",
750                &[
751                    (
752                        "id",
753                        DataType::Int {
754                            length: None,
755                            integer_spelling: false,
756                        },
757                    ),
758                    (
759                        "name",
760                        DataType::VarChar {
761                            length: None,
762                            parenthesized_length: false,
763                        },
764                    ),
765                ],
766            ),
767            (
768                "orders",
769                &[
770                    (
771                        "id",
772                        DataType::Int {
773                            length: None,
774                            integer_spelling: false,
775                        },
776                    ),
777                    (
778                        "user_id",
779                        DataType::Int {
780                            length: None,
781                            integer_spelling: false,
782                        },
783                    ),
784                ],
785            ),
786        ]);
787
788        assert!(schema.has_column("users", "id"));
789        assert!(schema.has_column("users", "name"));
790        assert!(schema.has_column("orders", "id"));
791        assert!(schema.has_column("orders", "user_id"));
792    }
793
794    #[test]
795    fn test_flatten_schema_paths() {
796        let mut schema = MappingSchema::new();
797        schema
798            .add_table(
799                "db1.table1",
800                &[(
801                    "id".to_string(),
802                    DataType::Int {
803                        length: None,
804                        integer_spelling: false,
805                    },
806                )],
807                None,
808            )
809            .unwrap();
810        schema
811            .add_table(
812                "db1.table2",
813                &[(
814                    "id".to_string(),
815                    DataType::Int {
816                        length: None,
817                        integer_spelling: false,
818                    },
819                )],
820                None,
821            )
822            .unwrap();
823        schema
824            .add_table(
825                "db2.table1",
826                &[(
827                    "id".to_string(),
828                    DataType::Int {
829                        length: None,
830                        integer_spelling: false,
831                    },
832                )],
833                None,
834            )
835            .unwrap();
836
837        let paths = flatten_schema_paths(&schema);
838        assert_eq!(paths.len(), 3);
839    }
840
841    #[test]
842    fn test_visible_columns() {
843        let mut schema = MappingSchema::new();
844        let columns = vec![
845            (
846                "id".to_string(),
847                DataType::Int {
848                    length: None,
849                    integer_spelling: false,
850                },
851            ),
852            (
853                "name".to_string(),
854                DataType::VarChar {
855                    length: None,
856                    parenthesized_length: false,
857                },
858            ),
859            (
860                "password".to_string(),
861                DataType::VarChar {
862                    length: None,
863                    parenthesized_length: false,
864                },
865            ),
866        ];
867        schema.add_table("users", &columns, None).unwrap();
868        schema.set_visible_columns("users", &["id", "name"]);
869
870        let names = schema.column_names("users").unwrap();
871        assert_eq!(names.len(), 2);
872        assert!(names.contains(&"id".to_string()));
873        assert!(names.contains(&"name".to_string()));
874        assert!(!names.contains(&"password".to_string()));
875    }
876}
polyglot_sql/schema.rs

polyglot_sql/
schema.rs
