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 {
96                expr,
97                alias,
98                leading_comments,
99                trailing_comment,
100            } => crate::sql::parser::ast::SelectItem::Expression {
101                expr: Self::remove_from_expression(expr),
102                alias,
103                leading_comments,
104                trailing_comment,
105            },
106            other => other,
107        }
108    }
109
110    /// Remove INTO from expressions (handles subqueries)
111    fn remove_from_expression(
112        expr: crate::sql::parser::ast::SqlExpression,
113    ) -> crate::sql::parser::ast::SqlExpression {
114        use crate::sql::parser::ast::SqlExpression;
115
116        match expr {
117            SqlExpression::ScalarSubquery { query } => SqlExpression::ScalarSubquery {
118                query: Box::new(Self::remove_from_statement(*query)),
119            },
120            SqlExpression::InSubquery { expr, subquery } => SqlExpression::InSubquery {
121                expr: Box::new(Self::remove_from_expression(*expr)),
122                subquery: Box::new(Self::remove_from_statement(*subquery)),
123            },
124            SqlExpression::NotInSubquery { expr, subquery } => SqlExpression::NotInSubquery {
125                expr: Box::new(Self::remove_from_expression(*expr)),
126                subquery: Box::new(Self::remove_from_statement(*subquery)),
127            },
128            SqlExpression::BinaryOp { left, op, right } => SqlExpression::BinaryOp {
129                left: Box::new(Self::remove_from_expression(*left)),
130                op,
131                right: Box::new(Self::remove_from_expression(*right)),
132            },
133            SqlExpression::FunctionCall {
134                name,
135                args,
136                distinct,
137            } => SqlExpression::FunctionCall {
138                name,
139                args: args
140                    .into_iter()
141                    .map(|arg| Self::remove_from_expression(arg))
142                    .collect(),
143                distinct,
144            },
145            SqlExpression::CaseExpression {
146                when_branches,
147                else_branch,
148            } => SqlExpression::CaseExpression {
149                when_branches: when_branches
150                    .into_iter()
151                    .map(|branch| crate::sql::parser::ast::WhenBranch {
152                        condition: Box::new(Self::remove_from_expression(*branch.condition)),
153                        result: Box::new(Self::remove_from_expression(*branch.result)),
154                    })
155                    .collect(),
156                else_branch: else_branch.map(|e| Box::new(Self::remove_from_expression(*e))),
157            },
158            SqlExpression::SimpleCaseExpression {
159                expr,
160                when_branches,
161                else_branch,
162            } => SqlExpression::SimpleCaseExpression {
163                expr: Box::new(Self::remove_from_expression(*expr)),
164                when_branches: when_branches
165                    .into_iter()
166                    .map(|branch| crate::sql::parser::ast::SimpleWhenBranch {
167                        value: Box::new(Self::remove_from_expression(*branch.value)),
168                        result: Box::new(Self::remove_from_expression(*branch.result)),
169                    })
170                    .collect(),
171                else_branch: else_branch.map(|e| Box::new(Self::remove_from_expression(*e))),
172            },
173            SqlExpression::InList { expr, values } => SqlExpression::InList {
174                expr: Box::new(Self::remove_from_expression(*expr)),
175                values: values
176                    .into_iter()
177                    .map(|e| Self::remove_from_expression(e))
178                    .collect(),
179            },
180            SqlExpression::NotInList { expr, values } => SqlExpression::NotInList {
181                expr: Box::new(Self::remove_from_expression(*expr)),
182                values: values
183                    .into_iter()
184                    .map(|e| Self::remove_from_expression(e))
185                    .collect(),
186            },
187            SqlExpression::Between { expr, lower, upper } => SqlExpression::Between {
188                expr: Box::new(Self::remove_from_expression(*expr)),
189                lower: Box::new(Self::remove_from_expression(*lower)),
190                upper: Box::new(Self::remove_from_expression(*upper)),
191            },
192            SqlExpression::Not { expr } => SqlExpression::Not {
193                expr: Box::new(Self::remove_from_expression(*expr)),
194            },
195            // Terminal expressions don't contain subqueries
196            other => other,
197        }
198    }
199}
200
201#[cfg(test)]
202mod tests {
203    use super::*;
204    use crate::sql::parser::ast::IntoTable;
205
206    #[test]
207    fn test_remove_simple_into() {
208        let stmt = SelectStatement {
209            distinct: false,
210            columns: vec!["col1".to_string()],
211            select_items: vec![],
212            from_table: Some("table1".to_string()),
213            from_subquery: None,
214            from_function: None,
215            from_alias: None,
216            joins: vec![],
217            where_clause: None,
218            order_by: None,
219            group_by: None,
220            having: None,
221            limit: None,
222            offset: None,
223            ctes: vec![],
224            into_table: Some(IntoTable {
225                name: "#temp".to_string(),
226            }),
227            set_operations: vec![],
228            leading_comments: vec![],
229            trailing_comment: None,
230        };
231
232        let result = IntoClauseRemover::remove_into_clause(stmt);
233        assert!(result.into_table.is_none());
234        assert_eq!(result.from_table, Some("table1".to_string()));
235    }
236
237    #[test]
238    fn test_remove_into_from_subquery() {
239        let subquery = SelectStatement {
240            distinct: false,
241            columns: vec![],
242            select_items: vec![],
243            from_table: Some("inner_table".to_string()),
244            from_subquery: None,
245            from_function: None,
246            from_alias: None,
247            joins: vec![],
248            where_clause: None,
249            order_by: None,
250            group_by: None,
251            having: None,
252            limit: None,
253            offset: None,
254            ctes: vec![],
255            into_table: Some(IntoTable {
256                name: "#inner_temp".to_string(),
257            }),
258            set_operations: vec![],
259            leading_comments: vec![],
260            trailing_comment: None,
261        };
262
263        let stmt = SelectStatement {
264            distinct: false,
265            columns: vec![],
266            select_items: vec![],
267            from_table: None,
268            from_subquery: Some(Box::new(subquery)),
269            from_function: None,
270            from_alias: Some("subq".to_string()),
271            joins: vec![],
272            where_clause: None,
273            order_by: None,
274            group_by: None,
275            having: None,
276            limit: None,
277            offset: None,
278            ctes: vec![],
279            into_table: Some(IntoTable {
280                name: "#outer_temp".to_string(),
281            }),
282            set_operations: vec![],
283            leading_comments: vec![],
284            trailing_comment: None,
285        };
286
287        let result = IntoClauseRemover::remove_into_clause(stmt);
288
289        // Both outer and inner INTO should be removed
290        assert!(result.into_table.is_none());
291        assert!(result.from_subquery.as_ref().unwrap().into_table.is_none());
292    }
293}
sql_cli/query_plan/into_clause_remover.rs

sql_cli/query_plan/
into_clause_remover.rs
