toasty_driver_integration_suite/
helpers.rs

1/// Helper function to look up TableId by table name (handles database-specific prefixes)
2pub fn table_id(db: &toasty::Db, table_name: &str) -> toasty_core::schema::db::TableId {
3    let schema = db.schema();
4
5    // First try exact match
6    if let Some(position) = schema.db.tables.iter().position(|t| t.name == table_name) {
7        return toasty_core::schema::db::TableId(position);
8    }
9
10    // If not found, try to find a table that ends with the given name (for database prefixes)
11    if let Some(position) = schema
12        .db
13        .tables
14        .iter()
15        .position(|t| t.name.ends_with(table_name))
16    {
17        return toasty_core::schema::db::TableId(position);
18    }
19
20    // If still not found, show available tables for debugging
21    let available_tables: Vec<_> = schema.db.tables.iter().map(|t| &t.name).collect();
22    panic!(
23        "Table '{}' not found. Available tables: {:?}",
24        table_name, available_tables
25    );
26}
27
28/// Helper function to get a single ColumnId for specified table and column
29pub fn column(
30    db: &toasty::Db,
31    table_name: &str,
32    column_name: &str,
33) -> toasty_core::schema::db::ColumnId {
34    columns(db, table_name, &[column_name])[0]
35}
36
37/// Helper function to generate a `Vec<ColumnId>` for specified table and columns
38pub fn columns(
39    db: &toasty::Db,
40    table_name: &str,
41    column_names: &[&str],
42) -> Vec<toasty_core::schema::db::ColumnId> {
43    let schema = db.schema();
44
45    // Find the table using the same logic as table_id (handles prefixes)
46    let table = schema
47        .db
48        .tables
49        .iter()
50        .find(|t| t.name == table_name || t.name.ends_with(table_name))
51        .unwrap_or_else(|| panic!("Table '{}' not found", table_name));
52
53    let table_id = table_id(db, table_name);
54
55    column_names
56        .iter()
57        .map(|col_name| {
58            let index = table
59                .columns
60                .iter()
61                .position(|c| c.name == *col_name)
62                .unwrap_or_else(|| {
63                    panic!("Column '{}' not found in table '{}'", col_name, table_name)
64                });
65
66            toasty_core::schema::db::ColumnId {
67                table: table_id,
68                index,
69            }
70        })
71        .collect()
72}
73
74/// Helper function to look up the storage type of a column from the schema
75/// (handles database-specific table prefixes)
76pub fn column_storage_ty(
77    db: &toasty::Db,
78    table_name: &str,
79    column_name: &str,
80) -> toasty_core::schema::db::Type {
81    let schema = db.schema();
82    let table = schema
83        .db
84        .tables
85        .iter()
86        .find(|t| t.name == table_name || t.name.ends_with(table_name))
87        .unwrap_or_else(|| panic!("table '{table_name}' not in schema"));
88    table
89        .columns
90        .iter()
91        .find(|c| c.name == column_name)
92        .unwrap_or_else(|| panic!("column '{column_name}' not in table '{table_name}'"))
93        .storage_ty
94        .clone()
95}
96
97use std::collections::BTreeMap;
98
99use toasty_core::{
100    driver::{Operation, operation::TypedValue},
101    stmt::{Assignment, Expr, ExprSet, Statement, Value},
102};
103
104use crate::Test;
105
106/// Resolve a value expression: a literal `Value`, or a bound param looked up in
107/// `params` (SQL drivers replace scalars with `Expr::Arg` placeholders).
108fn expr_value(expr: &Expr, params: &[Value]) -> Value {
109    match expr {
110        Expr::Value(value) => value.clone(),
111        Expr::Arg(arg) => params[arg.position].clone(),
112        other => panic!("expected a value expression, got {other:#?}"),
113    }
114}
115
116fn params_of(params: Vec<TypedValue>) -> Vec<Value> {
117    params.into_iter().map(|tv| tv.value).collect()
118}
119
120/// Pop the next logged op (a create) and return `column index -> inserted
121/// value`, normalizing SQL (`QuerySql`) vs key-value (`Insert`) ops and inlining
122/// bound params. Useful for asserting exactly what an `INSERT` writes.
123pub fn pop_insert(test: &mut Test) -> BTreeMap<usize, Value> {
124    let (op, _) = test.log().pop();
125    let (stmt, params) = match op {
126        Operation::QuerySql(q) => (q.stmt, params_of(q.params)),
127        Operation::Insert(i) => (i.stmt, params_of(i.params)),
128        other => panic!("expected an insert op, got {other:#?}"),
129    };
130    let Statement::Insert(insert) = stmt else {
131        panic!("expected an Insert statement");
132    };
133    let toasty_core::stmt::InsertTarget::Table(target) = &insert.target else {
134        panic!("expected a table insert target");
135    };
136    let ExprSet::Values(values) = &insert.source.body else {
137        panic!("expected a VALUES source");
138    };
139    // The single row is an `Expr::Record` of value/param exprs (SQL) or an
140    // already-evaluated `Value::Record` (key-value drivers).
141    let row: Vec<Value> = match &values.rows[0] {
142        Expr::Record(record) => record
143            .fields
144            .iter()
145            .map(|expr| expr_value(expr, &params))
146            .collect(),
147        Expr::Value(Value::Record(record)) => record.fields.clone(),
148        other => panic!("expected a record row, got {other:#?}"),
149    };
150    target
151        .columns
152        .iter()
153        .zip(row)
154        .map(|(col, value)| (col.index, value))
155        .collect()
156}
157
158/// Pop the next logged op (an update) and return `column index -> assigned
159/// value`, normalizing SQL (`QuerySql`) vs key-value (`UpdateByKey`) ops and
160/// inlining bound params.
161pub fn pop_update(test: &mut Test) -> BTreeMap<usize, Value> {
162    let (op, _) = test.log().pop();
163    let (assignments, params) = match op {
164        Operation::QuerySql(q) => match q.stmt {
165            Statement::Update(update) => (update.assignments, params_of(q.params)),
166            other => panic!("expected an Update statement, got {other:#?}"),
167        },
168        Operation::UpdateByKey(update) => (update.assignments, vec![]),
169        other => panic!("expected an update op, got {other:#?}"),
170    };
171    assignments
172        .iter()
173        .map(|(projection, assignment)| {
174            let Assignment::Set(expr) = assignment else {
175                panic!("expected a Set assignment, got {assignment:#?}");
176            };
177            (projection.as_slice()[0], expr_value(expr, &params))
178        })
179        .collect()
180}
181
182/// Pop the next logged op (a read) and return its filter predicate, normalizing
183/// SQL (`QuerySql` select) vs key-value (`Scan`) ops.
184pub fn pop_filter(test: &mut Test) -> Expr {
185    let (op, _) = test.log().pop();
186    match op {
187        Operation::QuerySql(q) => {
188            let Statement::Query(query) = q.stmt else {
189                panic!("expected a Query statement");
190            };
191            let ExprSet::Select(select) = query.body else {
192                panic!("expected a Select body");
193            };
194            select.filter.expr.expect("filter predicate present")
195        }
196        Operation::Scan(scan) => scan.filter.expect("filter predicate present"),
197        other => panic!("expected a query/scan op, got {other:#?}"),
198    }
199}
toasty_driver_integration_suite/helpers.rs

toasty_driver_integration_suite/
helpers.rs
