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

proof_of_sql/base/database/
test_schema_accessor.rs
