Skip to main content

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.
30///
31/// Bounded on `NestedStatement = Statement<Raw>` (true for every `AstInfo` —
32/// both `Raw` and `Aug` nest a raw statement) so that `DECLARE`/`PREPARE` can
33/// recurse into their inner statement's pretty doc, which is what keeps a secret
34/// in the inner statement (e.g. a role password) lossless rather than redacted.
35pub fn to_pretty<T: AstInfo<NestedStatement = Statement<Raw>>>(
36    stmt: &Statement<T>,
37    config: PrettyConfig,
38) -> String {
39    format!("{};", Pretty { config }.to_doc(stmt).pretty(config.width))
40}
41
42/// Parses `str` into SQL statements and pretty prints them.
43pub fn pretty_strs(str: &str, config: PrettyConfig) -> Result<Vec<String>, Error> {
44    let stmts = parse_statements(str)?;
45    Ok(stmts.iter().map(|s| to_pretty(&s.ast, config)).collect())
46}
47
48/// Parses `str` into a single SQL statement and pretty prints it.
49pub fn pretty_str(str: &str, config: PrettyConfig) -> Result<String, Error> {
50    let stmts = parse_statements(str)?;
51    if stmts.len() != 1 {
52        return Err(Error::ExpectedOne);
53    }
54    Ok(to_pretty(&stmts[0].ast, config))
55}
56
57/// Parses `str` into SQL statements and pretty prints them in `Simple` mode.
58pub fn pretty_strs_simple(str: &str, width: usize) -> Result<Vec<String>, Error> {
59    pretty_strs(
60        str,
61        PrettyConfig {
62            width,
63            format_mode: FormatMode::Simple,
64        },
65    )
66}
67
68/// Parses `str` into a single SQL statement and pretty prints it in `Simple` mode.
69pub fn pretty_str_simple(str: &str, width: usize) -> Result<String, Error> {
70    pretty_str(
71        str,
72        PrettyConfig {
73            width,
74            format_mode: FormatMode::Simple,
75        },
76    )
77}
78
79#[derive(Error, Debug)]
80pub enum Error {
81    #[error(transparent)]
82    Parser(#[from] ParserStatementError),
83    #[error("expected exactly one statement")]
84    ExpectedOne,
85}
86
87/// (Public only for tests)
88pub struct Pretty {
89    pub config: PrettyConfig,
90}
91
92impl Pretty {
93    fn to_doc<'a, T: AstInfo<NestedStatement = Statement<Raw>>>(
94        &'a self,
95        v: &'a Statement<T>,
96    ) -> RcDoc<'a> {
97        match v {
98            Statement::Select(v) => self.doc_select_statement(v),
99            Statement::Insert(v) => self.doc_insert(v),
100            Statement::CreateView(v) => self.doc_create_view(v),
101            Statement::CreateMaterializedView(v) => self.doc_create_materialized_view(v),
102            Statement::Copy(v) => self.doc_copy(v),
103            Statement::Subscribe(v) => self.doc_subscribe(v),
104            Statement::CreateSource(v) => self.doc_create_source(v),
105            Statement::CreateWebhookSource(v) => self.doc_create_webhook_source(v),
106            Statement::CreateTable(v) => self.doc_create_table(v),
107            Statement::CreateTableFromSource(v) => self.doc_create_table_from_source(v),
108            Statement::CreateConnection(v) => self.doc_create_connection(v),
109            Statement::CreateSink(v) => self.doc_create_sink(v),
110            Statement::CreateSubsource(v) => self.doc_create_subsource(v),
111            Statement::CreateCluster(v) => self.doc_create_cluster(v),
112            Statement::CreateClusterReplica(v) => self.doc_create_cluster_replica(v),
113            Statement::CreateNetworkPolicy(v) => self.doc_create_network_policy(v),
114            Statement::CreateIndex(v) => self.doc_create_index(v),
115            Statement::CreateRole(v) => self.doc_create_role(v),
116            Statement::AlterRole(v) => self.doc_alter_role(v),
117            // Recurse into the inner statement so a secret it carries (e.g. a
118            // role password) is printed losslessly by the inner doc rather than
119            // redacted by the fallback AstDisplay path.
120            Statement::Declare(v) => self.doc_declare(v),
121            Statement::Prepare(v) => self.doc_prepare(v),
122            _ => self.doc_display(v, "statement"),
123        }
124    }
125}