proof_of_sql/base/database/
test_schema_accessor.rs

1use super::{ColumnType, SchemaAccessor, TableRef};
2use crate::base::map::IndexMap;
3use alloc::vec::Vec;
4use sqlparser::ast::Ident;
5/// A simple in-memory `SchemaAccessor` for testing intermediate AST -> Provable AST conversion.
6#[derive(Clone)]
7pub struct TestSchemaAccessor {
8    schemas: IndexMap<TableRef, IndexMap<Ident, ColumnType>>,
9}
10
11impl TestSchemaAccessor {
12    /// Create a new `TestSchemaAccessor` with the given schema.
13    #[must_use]
14    pub fn new(schemas: IndexMap<TableRef, IndexMap<Ident, ColumnType>>) -> Self {
15        Self { schemas }
16    }
17}
18
19impl SchemaAccessor for TestSchemaAccessor {
20    fn lookup_column(&self, table_ref: &TableRef, column_id: &Ident) -> Option<ColumnType> {
21        self.schemas.get(table_ref)?.get(column_id).copied()
22    }
23
24    fn lookup_schema(&self, table_ref: &TableRef) -> Vec<(Ident, ColumnType)> {
25        self.schemas
26            .get(table_ref)
27            .unwrap_or(&IndexMap::default())
28            .iter()
29            .map(|(id, col)| (id.clone(), *col))
30            .collect()
31    }
32}
33
34#[cfg(test)]
35mod tests {
36    use super::*;
37    use crate::base::map::indexmap;
38
39    fn sample_test_schema_accessor() -> TestSchemaAccessor {
40        let table1 = TableRef::new("schema", "table1");
41        let table2 = TableRef::new("schema", "table2");
42        TestSchemaAccessor::new(indexmap! {
43            table1 => indexmap! {
44                "col1".into() => ColumnType::BigInt,
45                "col2".into() => ColumnType::VarChar,
46            },
47            table2 => indexmap! {
48                "col1".into() => ColumnType::BigInt,
49            },
50        })
51    }
52
53    #[test]
54    fn test_lookup_column() {
55        let accessor = sample_test_schema_accessor();
56        let table1 = TableRef::new("schema", "table1");
57        let table2 = TableRef::new("schema", "table2");
58        let not_a_table = TableRef::new("schema", "not_a_table");
59        assert_eq!(
60            accessor.lookup_column(&table1, &"col1".into()),
61            Some(ColumnType::BigInt)
62        );
63        assert_eq!(
64            accessor.lookup_column(&table1, &"col2".into()),
65            Some(ColumnType::VarChar)
66        );
67        assert_eq!(accessor.lookup_column(&table1, &"not_a_col".into()), None);
68        assert_eq!(
69            accessor.lookup_column(&table2, &"col1".into()),
70            Some(ColumnType::BigInt)
71        );
72        assert_eq!(accessor.lookup_column(&table2, &"col2".into()), None);
73        assert_eq!(accessor.lookup_column(&not_a_table, &"col1".into()), None);
74        assert_eq!(accessor.lookup_column(&not_a_table, &"col2".into()), None);
75        assert_eq!(
76            accessor.lookup_column(&not_a_table, &"not_a_col".into()),
77            None
78        );
79    }
80
81    #[test]
82    fn test_lookup_schema() {
83        let accessor = sample_test_schema_accessor();
84        let table1 = TableRef::new("schema", "table1");
85        let table2 = TableRef::new("schema", "table2");
86        let not_a_table = TableRef::new("schema", "not_a_table");
87        assert_eq!(
88            accessor.lookup_schema(&table1),
89            vec![
90                ("col1".into(), ColumnType::BigInt),
91                ("col2".into(), ColumnType::VarChar),
92            ]
93        );
94        assert_eq!(
95            accessor.lookup_schema(&table2),
96            vec![("col1".into(), ColumnType::BigInt),]
97        );
98        assert_eq!(accessor.lookup_schema(&not_a_table), vec![]);
99    }
100}