mz_sql_pretty/
lib.rs

1// Copyright Materialize, Inc. and contributors. All rights reserved.
2//
3// Use of this software is governed by the Business Source License
4// included in the LICENSE file.
5//
6// As of the Change Date specified in that file, in accordance with
7// the Business Source License, use of this software will be governed
8// by the Apache License, Version 2.0.
9
10mod doc;
11mod util;
12
13use mz_sql_parser::ast::display::FormatMode;
14use mz_sql_parser::ast::*;
15use mz_sql_parser::parser::{ParserStatementError, parse_statements};
16use pretty::RcDoc;
17use thiserror::Error;
18
19pub const DEFAULT_WIDTH: usize = 100;
20
21const TAB: isize = 4;
22
23#[derive(Clone, Copy)]
24pub struct PrettyConfig {
25    pub width: usize,
26    pub format_mode: FormatMode,
27}
28
29/// Pretty prints a statement at a width.
30pub fn to_pretty<T: AstInfo>(stmt: &Statement<T>, config: PrettyConfig) -> String {
31    format!("{};", Pretty { config }.to_doc(stmt).pretty(config.width))
32}
33
34/// Parses `str` into SQL statements and pretty prints them.
35pub fn pretty_strs(str: &str, config: PrettyConfig) -> Result<Vec<String>, Error> {
36    let stmts = parse_statements(str)?;
37    Ok(stmts.iter().map(|s| to_pretty(&s.ast, config)).collect())
38}
39
40/// Parses `str` into a single SQL statement and pretty prints it.
41pub fn pretty_str(str: &str, config: PrettyConfig) -> Result<String, Error> {
42    let stmts = parse_statements(str)?;
43    if stmts.len() != 1 {
44        return Err(Error::ExpectedOne);
45    }
46    Ok(to_pretty(&stmts[0].ast, config))
47}
48
49/// Parses `str` into SQL statements and pretty prints them in `Simple` mode.
50pub fn pretty_strs_simple(str: &str, width: usize) -> Result<Vec<String>, Error> {
51    pretty_strs(
52        str,
53        PrettyConfig {
54            width,
55            format_mode: FormatMode::Simple,
56        },
57    )
58}
59
60/// Parses `str` into a single SQL statement and pretty prints it in `Simple` mode.
61pub fn pretty_str_simple(str: &str, width: usize) -> Result<String, Error> {
62    pretty_str(
63        str,
64        PrettyConfig {
65            width,
66            format_mode: FormatMode::Simple,
67        },
68    )
69}
70
71#[derive(Error, Debug)]
72pub enum Error {
73    #[error(transparent)]
74    Parser(#[from] ParserStatementError),
75    #[error("expected exactly one statement")]
76    ExpectedOne,
77}
78
79/// (Public only for tests)
80pub struct Pretty {
81    pub config: PrettyConfig,
82}
83
84impl Pretty {
85    fn to_doc<'a, T: AstInfo>(&'a self, v: &'a Statement<T>) -> RcDoc<'a> {
86        match v {
87            Statement::Select(v) => self.doc_select_statement(v),
88            Statement::Insert(v) => self.doc_insert(v),
89            Statement::CreateView(v) => self.doc_create_view(v),
90            Statement::CreateMaterializedView(v) => self.doc_create_materialized_view(v),
91            Statement::Copy(v) => self.doc_copy(v),
92            Statement::Subscribe(v) => self.doc_subscribe(v),
93            Statement::CreateSource(v) => self.doc_create_source(v),
94            _ => self.doc_display(v, "statement"),
95        }
96    }
97}