sqlfuzz/
sql_writer.rs

1// Copyright 2022 Andy Grove
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7// http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use crate::fuzz_sql::{SQLExpr, SQLJoin, SQLSelect, SQLSubqueryAlias};
16use crate::SQLRelation;
17use datafusion::common::Result;
18
19/// Generate a SQL string from a SQLRelation struct
20pub fn plan_to_sql(plan: &SQLRelation, indent: usize) -> Result<String> {
21    let indent_str = "  ".repeat(indent);
22    match plan {
23        SQLRelation::Select(SQLSelect {
24            projection,
25            filter,
26            input,
27        }) => {
28            let expr: Vec<String> = projection
29                .iter()
30                .map(|e| expr_to_sql(e, indent))
31                .collect::<Result<Vec<_>>>()?;
32            let input = plan_to_sql(input, indent + 1)?;
33            let where_clause = if let Some(predicate) = filter {
34                let predicate = expr_to_sql(predicate, indent)?;
35                format!("\n{}WHERE {}", indent_str, predicate)
36            } else {
37                "".to_string()
38            };
39            Ok(format!(
40                "SELECT {}\n{}FROM ({}){}",
41                expr.join(", "),
42                indent_str,
43                input,
44                where_clause
45            ))
46        }
47        SQLRelation::TableScan(scan) => Ok(scan.table_name.clone()),
48        SQLRelation::Join(SQLJoin {
49            left,
50            right,
51            on,
52            join_type,
53            ..
54        }) => {
55            let l = plan_to_sql(left, indent + 1)?;
56            let r = plan_to_sql(right, indent + 1)?;
57            let join_condition = on
58                .iter()
59                .map(|(l, r)| format!("{} = {}", l.flat_name(), r.flat_name()))
60                .collect::<Vec<_>>()
61                .join(" AND ");
62            Ok(format!(
63                "\n{}({})\n{}{} JOIN\n{}({})\n{}ON {}",
64                indent_str,
65                l,
66                indent_str,
67                join_type.to_string().to_uppercase(),
68                indent_str,
69                r,
70                indent_str,
71                join_condition
72            ))
73        }
74        SQLRelation::SubqueryAlias(SQLSubqueryAlias { input, alias, .. }) => {
75            let sql = plan_to_sql(input, indent + 1)?;
76            Ok(format!("({}) {}", sql, alias))
77        }
78    }
79}
80
81/// Generate a SQL string from an expression
82fn expr_to_sql(expr: &SQLExpr, indent: usize) -> Result<String> {
83    Ok(match expr {
84        SQLExpr::Alias { expr, alias } => format!("{} AS {}", expr_to_sql(expr, indent)?, alias),
85        SQLExpr::Column(col) => col.flat_name(),
86        SQLExpr::BinaryExpr { left, op, right } => {
87            let l = expr_to_sql(left, indent)?;
88            let r = expr_to_sql(right, indent)?;
89            format!("{} {} {}", l, op, r)
90        }
91        SQLExpr::Exists { subquery, negated } => {
92            let sql = plan_to_sql(&SQLRelation::Select(subquery.as_ref().clone()), indent)?;
93            if *negated {
94                format!("NOT EXISTS ({})", sql)
95            } else {
96                format!("EXISTS ({})", sql)
97            }
98        }
99    })
100}
sqlfuzz/sql_writer.rs

sqlfuzz/
sql_writer.rs
