vibesql_executor/cache/
table_extractor.rs

1//! Extract table names from AST for cache invalidation
2
3use std::collections::HashSet;
4
5/// Extract all table names referenced in a SELECT statement
6pub fn extract_tables_from_select(stmt: &vibesql_ast::SelectStmt) -> HashSet<String> {
7    let mut tables = HashSet::new();
8
9    // Extract from FROM clause
10    if let Some(from_clause) = &stmt.from {
11        extract_from_from_clause(from_clause, &mut tables);
12    }
13
14    // Extract from subqueries in SELECT list
15    for select_item in &stmt.select_list {
16        if let vibesql_ast::SelectItem::Expression { expr, .. } = select_item {
17            extract_from_expression(expr, &mut tables);
18        }
19    }
20
21    // Extract from WHERE clause
22    if let Some(where_clause) = &stmt.where_clause {
23        extract_from_expression(where_clause, &mut tables);
24    }
25
26    // Extract from GROUP BY
27    if let Some(group_by) = &stmt.group_by {
28        for expr in group_by.all_expressions() {
29            extract_from_expression(expr, &mut tables);
30        }
31    }
32
33    // Extract from HAVING
34    if let Some(having) = &stmt.having {
35        extract_from_expression(having, &mut tables);
36    }
37
38    // Extract from ORDER BY
39    if let Some(order_by) = &stmt.order_by {
40        for order_item in order_by {
41            extract_from_expression(&order_item.expr, &mut tables);
42        }
43    }
44
45    // Extract from CTEs (WITH clause)
46    if let Some(with_clause) = &stmt.with_clause {
47        for cte in with_clause {
48            let cte_tables = extract_tables_from_select(&cte.query);
49            tables.extend(cte_tables);
50        }
51    }
52
53    // Extract from set operations (UNION, INTERSECT, EXCEPT)
54    if let Some(set_op) = &stmt.set_operation {
55        let right_tables = extract_tables_from_select(&set_op.right);
56        tables.extend(right_tables);
57    }
58
59    tables
60}
61
62/// Extract table names from a FROM clause
63fn extract_from_from_clause(from: &vibesql_ast::FromClause, tables: &mut HashSet<String>) {
64    match from {
65        vibesql_ast::FromClause::Table { name, .. } => {
66            // Handle qualified names (schema.table) - we want just the table name
67            let table_name = if let Some(pos) = name.rfind('.') { &name[pos + 1..] } else { name };
68            tables.insert(table_name.to_string());
69        }
70        vibesql_ast::FromClause::Join { left, right, condition, .. } => {
71            extract_from_from_clause(left, tables);
72            extract_from_from_clause(right, tables);
73            if let Some(cond) = condition {
74                extract_from_expression(cond, tables);
75            }
76        }
77        vibesql_ast::FromClause::Subquery { query, .. } => {
78            let subquery_tables = extract_tables_from_select(query);
79            tables.extend(subquery_tables);
80        }
81        vibesql_ast::FromClause::Values { .. } => {
82            // VALUES clauses don't reference any tables
83        }
84    }
85}
86
87/// Extract table names from an expression (for subqueries)
88fn extract_from_expression(expr: &vibesql_ast::Expression, tables: &mut HashSet<String>) {
89    match expr {
90        vibesql_ast::Expression::ScalarSubquery(stmt) => {
91            let subquery_tables = extract_tables_from_select(stmt);
92            tables.extend(subquery_tables);
93        }
94        vibesql_ast::Expression::BinaryOp { left, right, .. } => {
95            extract_from_expression(left, tables);
96            extract_from_expression(right, tables);
97        }
98        vibesql_ast::Expression::UnaryOp { expr, .. } => {
99            extract_from_expression(expr, tables);
100        }
101        vibesql_ast::Expression::Function { args, .. }
102        | vibesql_ast::Expression::AggregateFunction { args, .. } => {
103            for arg in args {
104                extract_from_expression(arg, tables);
105            }
106        }
107        vibesql_ast::Expression::Case { operand, when_clauses, else_result, .. } => {
108            if let Some(op) = operand {
109                extract_from_expression(op, tables);
110            }
111            for when_clause in when_clauses {
112                for condition in &when_clause.conditions {
113                    extract_from_expression(condition, tables);
114                }
115                extract_from_expression(&when_clause.result, tables);
116            }
117            if let Some(else_expr) = else_result {
118                extract_from_expression(else_expr, tables);
119            }
120        }
121        vibesql_ast::Expression::In { expr, subquery, .. } => {
122            extract_from_expression(expr, tables);
123            let subquery_tables = extract_tables_from_select(subquery);
124            tables.extend(subquery_tables);
125        }
126        vibesql_ast::Expression::InList { expr, values, .. } => {
127            extract_from_expression(expr, tables);
128            for val in values {
129                extract_from_expression(val, tables);
130            }
131        }
132        vibesql_ast::Expression::Exists { subquery, .. } => {
133            let subquery_tables = extract_tables_from_select(subquery);
134            tables.extend(subquery_tables);
135        }
136        vibesql_ast::Expression::Between { expr, low, high, .. } => {
137            extract_from_expression(expr, tables);
138            extract_from_expression(low, tables);
139            extract_from_expression(high, tables);
140        }
141        vibesql_ast::Expression::IsNull { expr, .. } => {
142            extract_from_expression(expr, tables);
143        }
144        vibesql_ast::Expression::IsDistinctFrom { left, right, .. } => {
145            extract_from_expression(left, tables);
146            extract_from_expression(right, tables);
147        }
148        vibesql_ast::Expression::IsTruthValue { expr, .. } => {
149            extract_from_expression(expr, tables);
150        }
151        vibesql_ast::Expression::Cast { expr, .. } => {
152            extract_from_expression(expr, tables);
153        }
154        vibesql_ast::Expression::Like { expr, pattern, .. }
155        | vibesql_ast::Expression::Glob { expr, pattern, .. } => {
156            extract_from_expression(expr, tables);
157            extract_from_expression(pattern, tables);
158        }
159        vibesql_ast::Expression::Position { substring, string, .. } => {
160            extract_from_expression(substring, tables);
161            extract_from_expression(string, tables);
162        }
163        vibesql_ast::Expression::Trim { removal_char, string, .. } => {
164            if let Some(removal) = removal_char {
165                extract_from_expression(removal, tables);
166            }
167            extract_from_expression(string, tables);
168        }
169        vibesql_ast::Expression::Extract { expr, .. } => {
170            extract_from_expression(expr, tables);
171        }
172        vibesql_ast::Expression::QuantifiedComparison { expr, subquery, .. } => {
173            extract_from_expression(expr, tables);
174            let subquery_tables = extract_tables_from_select(subquery);
175            tables.extend(subquery_tables);
176        }
177        vibesql_ast::Expression::Conjunction(children)
178        | vibesql_ast::Expression::Disjunction(children)
179        | vibesql_ast::Expression::RowValueConstructor(children) => {
180            for child in children {
181                extract_from_expression(child, tables);
182            }
183        }
184
185        vibesql_ast::Expression::Collate { expr, .. } => {
186            extract_from_expression(expr, tables);
187        }
188
189        vibesql_ast::Expression::Raise { error_message, .. } => {
190            if let Some(msg) = error_message {
191                extract_from_expression(msg, tables);
192            }
193        }
194
195        // Leaf expressions - no tables to extract
196        vibesql_ast::Expression::Literal(_)
197        | vibesql_ast::Expression::Placeholder(_)
198        | vibesql_ast::Expression::NumberedPlaceholder(_)
199        | vibesql_ast::Expression::NamedPlaceholder(_)
200        | vibesql_ast::Expression::ColumnRef(_)
201        | vibesql_ast::Expression::Wildcard
202        | vibesql_ast::Expression::CurrentDate
203        | vibesql_ast::Expression::CurrentTime { .. }
204        | vibesql_ast::Expression::CurrentTimestamp { .. }
205        | vibesql_ast::Expression::Interval { .. }
206        | vibesql_ast::Expression::Default
207        | vibesql_ast::Expression::DuplicateKeyValue { .. }
208        | vibesql_ast::Expression::WindowFunction { .. }
209        | vibesql_ast::Expression::NextValue { .. }
210        | vibesql_ast::Expression::MatchAgainst { .. }
211        | vibesql_ast::Expression::PseudoVariable { .. }
212        | vibesql_ast::Expression::SessionVariable { .. } => {}
213    }
214}
215
216/// Extract table names from any statement (for comprehensive cache invalidation)
217pub fn extract_tables_from_statement(stmt: &vibesql_ast::Statement) -> HashSet<String> {
218    match stmt {
219        vibesql_ast::Statement::Select(select) => extract_tables_from_select(select),
220        vibesql_ast::Statement::Insert(insert) => {
221            let mut tables = HashSet::new();
222            // Extract table name being inserted into
223            let table_name = if let Some(pos) = insert.table_name.rfind('.') {
224                &insert.table_name[pos + 1..]
225            } else {
226                &insert.table_name
227            };
228            tables.insert(table_name.to_string());
229
230            // Extract from source (VALUES or SELECT)
231            match &insert.source {
232                vibesql_ast::InsertSource::Values(values) => {
233                    for row in values {
234                        for expr in row {
235                            extract_from_expression(expr, &mut tables);
236                        }
237                    }
238                }
239                vibesql_ast::InsertSource::Select(select) => {
240                    let select_tables = extract_tables_from_select(select);
241                    tables.extend(select_tables);
242                }
243                vibesql_ast::InsertSource::DefaultValues => {
244                    // No expressions to extract from DEFAULT VALUES
245                }
246            }
247
248            tables
249        }
250        vibesql_ast::Statement::Update(update) => {
251            let mut tables = HashSet::new();
252            // Extract table being updated
253            let table_name = if let Some(pos) = update.table_name.rfind('.') {
254                &update.table_name[pos + 1..]
255            } else {
256                &update.table_name
257            };
258            tables.insert(table_name.to_string());
259
260            // Extract from SET assignments
261            for assignment in &update.assignments {
262                extract_from_expression(&assignment.value, &mut tables);
263            }
264
265            // Extract from WHERE clause
266            if let Some(vibesql_ast::WhereClause::Condition(expr)) = &update.where_clause {
267                extract_from_expression(expr, &mut tables);
268            }
269
270            tables
271        }
272        vibesql_ast::Statement::Delete(delete) => {
273            let mut tables = HashSet::new();
274            // Extract table being deleted from
275            let table_name = if let Some(pos) = delete.table_name.rfind('.') {
276                &delete.table_name[pos + 1..]
277            } else {
278                &delete.table_name
279            };
280            tables.insert(table_name.to_string());
281
282            // Extract from WHERE clause
283            if let Some(vibesql_ast::WhereClause::Condition(expr)) = &delete.where_clause {
284                extract_from_expression(expr, &mut tables);
285            }
286
287            tables
288        }
289        // DDL statements don't reference tables in a way that matters for SELECT caching
290        _ => HashSet::new(),
291    }
292}
293
294#[cfg(test)]
295mod tests {
296    use vibesql_parser::Parser;
297
298    use super::*;
299
300    #[test]
301    fn test_extract_simple_select() {
302        let sql = "SELECT * FROM users";
303        let stmt = Parser::parse_sql(sql).unwrap();
304
305        if let vibesql_ast::Statement::Select(select) = stmt {
306            let tables = extract_tables_from_select(&select);
307            assert_eq!(tables.len(), 1);
308            // Parser uppercases identifiers
309            assert!(tables.contains("users"));
310        } else {
311            panic!("Expected SELECT statement");
312        }
313    }
314
315    #[test]
316    fn test_extract_join() {
317        let sql = "SELECT * FROM users u JOIN orders o ON u.id = o.user_id";
318        let stmt = Parser::parse_sql(sql).unwrap();
319
320        if let vibesql_ast::Statement::Select(select) = stmt {
321            let tables = extract_tables_from_select(&select);
322            assert_eq!(tables.len(), 2);
323            // Parser uppercases identifiers
324            assert!(tables.contains("users"));
325            assert!(tables.contains("orders"));
326        } else {
327            panic!("Expected SELECT statement");
328        }
329    }
330
331    #[test]
332    fn test_extract_qualified_table_name() {
333        let sql = "SELECT * FROM public.users";
334        let stmt = Parser::parse_sql(sql).unwrap();
335
336        if let vibesql_ast::Statement::Select(select) = stmt {
337            let tables = extract_tables_from_select(&select);
338            assert_eq!(tables.len(), 1);
339            // Should extract just the table name, not the schema
340            assert!(tables.contains("users"));
341        } else {
342            panic!("Expected SELECT statement");
343        }
344    }
345
346    #[test]
347    fn test_extract_subquery_in_from() {
348        let sql = "SELECT * FROM (SELECT * FROM users) AS u";
349        let stmt = Parser::parse_sql(sql).unwrap();
350
351        if let vibesql_ast::Statement::Select(select) = stmt {
352            let tables = extract_tables_from_select(&select);
353            assert_eq!(tables.len(), 1);
354            assert!(tables.contains("users"));
355        } else {
356            panic!("Expected SELECT statement");
357        }
358    }
359
360    #[test]
361    fn test_extract_from_insert() {
362        let sql = "INSERT INTO users VALUES (1, 'Alice')";
363        let stmt = Parser::parse_sql(sql).unwrap();
364        let tables = extract_tables_from_statement(&stmt);
365
366        assert_eq!(tables.len(), 1);
367        assert!(tables.contains("users"));
368    }
369
370    #[test]
371    fn test_extract_from_update() {
372        let sql = "UPDATE users SET name = 'Bob' WHERE id = 1";
373        let stmt = Parser::parse_sql(sql).unwrap();
374        let tables = extract_tables_from_statement(&stmt);
375
376        assert_eq!(tables.len(), 1);
377        assert!(tables.contains("users"));
378    }
379
380    #[test]
381    fn test_extract_from_delete() {
382        let sql = "DELETE FROM users WHERE id = 1";
383        let stmt = Parser::parse_sql(sql).unwrap();
384        let tables = extract_tables_from_statement(&stmt);
385
386        assert_eq!(tables.len(), 1);
387        assert!(tables.contains("users"));
388    }
389}
vibesql_executor/cache/table_extractor.rs

vibesql_executor/cache/
table_extractor.rs
