sql_cli/query_plan/
into_clause_remover.rs

1use crate::sql::parser::ast::SelectStatement;
2
3/// INTO Clause Remover - Removes INTO clause from AST for execution
4///
5/// This module implements AST rewriting to remove the INTO clause from
6/// SELECT statements. The INTO clause is used to store query results in
7/// temporary tables, but the query executor doesn't understand this syntax.
8///
9/// The removal is done at the AST level (not via regex) to ensure correctness
10/// and maintainability. The caller is responsible for capturing the INTO table
11/// information before removal and storing the results after execution.
12///
13/// Example transformation:
14/// ```sql
15/// -- Input:
16/// SELECT col1, col2 INTO #temp FROM table WHERE x > 5
17///
18/// -- Output (for execution):
19/// SELECT col1, col2 FROM table WHERE x > 5
20/// ```
21pub struct IntoClauseRemover;
22
23impl IntoClauseRemover {
24    /// Remove INTO clause from statement and all nested subqueries
25    ///
26    /// This creates a new statement with `into_table` set to None.
27    /// The original statement is not modified.
28    ///
29    /// # Arguments
30    /// * `statement` - The SELECT statement to process
31    ///
32    /// # Returns
33    /// A new statement with INTO clause removed from all levels
34    pub fn remove_into_clause(statement: SelectStatement) -> SelectStatement {
35        Self::remove_from_statement(statement)
36    }
37
38    /// Recursively remove INTO clause from a statement and its subqueries
39    fn remove_from_statement(mut statement: SelectStatement) -> SelectStatement {
40        // Remove the INTO clause from this statement
41        statement.into_table = None;
42
43        // Remove from subquery in FROM clause
44        if let Some(subquery) = statement.from_subquery.take() {
45            statement.from_subquery = Some(Box::new(Self::remove_from_statement(*subquery)));
46        }
47
48        // Remove from JOIN subqueries
49        statement.joins = statement
50            .joins
51            .into_iter()
52            .map(|mut join| {
53                if let crate::sql::parser::ast::TableSource::DerivedTable { query, alias } =
54                    join.table
55                {
56                    join.table = crate::sql::parser::ast::TableSource::DerivedTable {
57                        query: Box::new(Self::remove_from_statement(*query)),
58                        alias,
59                    };
60                }
61                join
62            })
63            .collect();
64
65        // Remove from scalar subqueries and other expression subqueries
66        statement.select_items = statement
67            .select_items
68            .into_iter()
69            .map(|item| Self::remove_from_select_item(item))
70            .collect();
71
72        // Remove from WHERE clause subqueries
73        if let Some(mut where_clause) = statement.where_clause.take() {
74            for condition in &mut where_clause.conditions {
75                condition.expr = Self::remove_from_expression(condition.expr.clone());
76            }
77            statement.where_clause = Some(where_clause);
78        }
79
80        // Remove from set operation queries (UNION, INTERSECT, EXCEPT)
81        statement.set_operations = statement
82            .set_operations
83            .into_iter()
84            .map(|(op, query)| (op, Box::new(Self::remove_from_statement(*query))))
85            .collect();
86
87        statement
88    }
89
90    /// Remove INTO from SELECT items (handles subqueries in expressions)
91    fn remove_from_select_item(
92        item: crate::sql::parser::ast::SelectItem,
93    ) -> crate::sql::parser::ast::SelectItem {
94        match item {
95            crate::sql::parser::ast::SelectItem::Expression { expr, alias } => {
96                crate::sql::parser::ast::SelectItem::Expression {
97                    expr: Self::remove_from_expression(expr),
98                    alias,
99                }
100            }
101            other => other,
102        }
103    }
104
105    /// Remove INTO from expressions (handles subqueries)
106    fn remove_from_expression(
107        expr: crate::sql::parser::ast::SqlExpression,
108    ) -> crate::sql::parser::ast::SqlExpression {
109        use crate::sql::parser::ast::SqlExpression;
110
111        match expr {
112            SqlExpression::ScalarSubquery { query } => SqlExpression::ScalarSubquery {
113                query: Box::new(Self::remove_from_statement(*query)),
114            },
115            SqlExpression::InSubquery { expr, subquery } => SqlExpression::InSubquery {
116                expr: Box::new(Self::remove_from_expression(*expr)),
117                subquery: Box::new(Self::remove_from_statement(*subquery)),
118            },
119            SqlExpression::NotInSubquery { expr, subquery } => SqlExpression::NotInSubquery {
120                expr: Box::new(Self::remove_from_expression(*expr)),
121                subquery: Box::new(Self::remove_from_statement(*subquery)),
122            },
123            SqlExpression::BinaryOp { left, op, right } => SqlExpression::BinaryOp {
124                left: Box::new(Self::remove_from_expression(*left)),
125                op,
126                right: Box::new(Self::remove_from_expression(*right)),
127            },
128            SqlExpression::FunctionCall {
129                name,
130                args,
131                distinct,
132            } => SqlExpression::FunctionCall {
133                name,
134                args: args
135                    .into_iter()
136                    .map(|arg| Self::remove_from_expression(arg))
137                    .collect(),
138                distinct,
139            },
140            SqlExpression::CaseExpression {
141                when_branches,
142                else_branch,
143            } => SqlExpression::CaseExpression {
144                when_branches: when_branches
145                    .into_iter()
146                    .map(|branch| crate::sql::parser::ast::WhenBranch {
147                        condition: Box::new(Self::remove_from_expression(*branch.condition)),
148                        result: Box::new(Self::remove_from_expression(*branch.result)),
149                    })
150                    .collect(),
151                else_branch: else_branch.map(|e| Box::new(Self::remove_from_expression(*e))),
152            },
153            SqlExpression::SimpleCaseExpression {
154                expr,
155                when_branches,
156                else_branch,
157            } => SqlExpression::SimpleCaseExpression {
158                expr: Box::new(Self::remove_from_expression(*expr)),
159                when_branches: when_branches
160                    .into_iter()
161                    .map(|branch| crate::sql::parser::ast::SimpleWhenBranch {
162                        value: Box::new(Self::remove_from_expression(*branch.value)),
163                        result: Box::new(Self::remove_from_expression(*branch.result)),
164                    })
165                    .collect(),
166                else_branch: else_branch.map(|e| Box::new(Self::remove_from_expression(*e))),
167            },
168            SqlExpression::InList { expr, values } => SqlExpression::InList {
169                expr: Box::new(Self::remove_from_expression(*expr)),
170                values: values
171                    .into_iter()
172                    .map(|e| Self::remove_from_expression(e))
173                    .collect(),
174            },
175            SqlExpression::NotInList { expr, values } => SqlExpression::NotInList {
176                expr: Box::new(Self::remove_from_expression(*expr)),
177                values: values
178                    .into_iter()
179                    .map(|e| Self::remove_from_expression(e))
180                    .collect(),
181            },
182            SqlExpression::Between { expr, lower, upper } => SqlExpression::Between {
183                expr: Box::new(Self::remove_from_expression(*expr)),
184                lower: Box::new(Self::remove_from_expression(*lower)),
185                upper: Box::new(Self::remove_from_expression(*upper)),
186            },
187            SqlExpression::Not { expr } => SqlExpression::Not {
188                expr: Box::new(Self::remove_from_expression(*expr)),
189            },
190            // Terminal expressions don't contain subqueries
191            other => other,
192        }
193    }
194}
195
196#[cfg(test)]
197mod tests {
198    use super::*;
199    use crate::sql::parser::ast::IntoTable;
200
201    #[test]
202    fn test_remove_simple_into() {
203        let stmt = SelectStatement {
204            distinct: false,
205            columns: vec!["col1".to_string()],
206            select_items: vec![],
207            from_table: Some("table1".to_string()),
208            from_subquery: None,
209            from_function: None,
210            from_alias: None,
211            joins: vec![],
212            where_clause: None,
213            order_by: None,
214            group_by: None,
215            having: None,
216            limit: None,
217            offset: None,
218            ctes: vec![],
219            into_table: Some(IntoTable {
220                name: "#temp".to_string(),
221            }),
222            set_operations: vec![],
223        };
224
225        let result = IntoClauseRemover::remove_into_clause(stmt);
226        assert!(result.into_table.is_none());
227        assert_eq!(result.from_table, Some("table1".to_string()));
228    }
229
230    #[test]
231    fn test_remove_into_from_subquery() {
232        let subquery = SelectStatement {
233            distinct: false,
234            columns: vec![],
235            select_items: vec![],
236            from_table: Some("inner_table".to_string()),
237            from_subquery: None,
238            from_function: None,
239            from_alias: None,
240            joins: vec![],
241            where_clause: None,
242            order_by: None,
243            group_by: None,
244            having: None,
245            limit: None,
246            offset: None,
247            ctes: vec![],
248            into_table: Some(IntoTable {
249                name: "#inner_temp".to_string(),
250            }),
251            set_operations: vec![],
252        };
253
254        let stmt = SelectStatement {
255            distinct: false,
256            columns: vec![],
257            select_items: vec![],
258            from_table: None,
259            from_subquery: Some(Box::new(subquery)),
260            from_function: None,
261            from_alias: Some("subq".to_string()),
262            joins: vec![],
263            where_clause: None,
264            order_by: None,
265            group_by: None,
266            having: None,
267            limit: None,
268            offset: None,
269            ctes: vec![],
270            into_table: Some(IntoTable {
271                name: "#outer_temp".to_string(),
272            }),
273            set_operations: vec![],
274        };
275
276        let result = IntoClauseRemover::remove_into_clause(stmt);
277
278        // Both outer and inner INTO should be removed
279        assert!(result.into_table.is_none());
280        assert!(result.from_subquery.as_ref().unwrap().into_table.is_none());
281    }
282}