risingwave_sqlparser/
quote_ident.rs

1// Licensed under the Apache License, Version 2.0 (the "License");
2// you may not use this file except in compliance with the License.
3// You may obtain a copy of the License at
4//
5//     http://www.apache.org/licenses/LICENSE-2.0
6//
7// Unless required by applicable law or agreed to in writing, software
8// distributed under the License is distributed on an "AS IS" BASIS,
9// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
10// See the License for the specific language governing permissions and
11// limitations under the License.
12
13use std::fmt::{self, Display};
14
15use crate::keywords;
16
17/// A wrapper that returns the given string suitably quoted to be used as an identifier in an SQL
18/// statement string in its `Display` implementation.
19/// Quotes are added only if necessary (i.e., if the string contains non-identifier characters,
20/// would be case-folded, or is a SQL keyword). Embedded quotes are properly doubled.
21///
22/// Refer to <https://github.com/postgres/postgres/blob/90189eefc1e11822794e3386d9bafafd3ba3a6e8/src/backend/utils/adt/ruleutils.c#L11506>
23pub struct QuoteIdent<'a>(pub &'a str);
24
25impl QuoteIdent<'_> {
26    pub(crate) fn needs_quotes(ident: &str) -> bool {
27        !is_simple_identifier(ident) || is_keyword(ident)
28    }
29}
30
31impl Display for QuoteIdent<'_> {
32    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
33        if !Self::needs_quotes(self.0) {
34            return f.write_str(self.0);
35        }
36
37        f.write_str("\"")?;
38        for c in self.0.chars() {
39            if c == '"' {
40                f.write_str("\"\"")?;
41            } else {
42                write!(f, "{c}")?;
43            }
44        }
45        f.write_str("\"")
46    }
47}
48
49fn is_simple_identifier(ident: &str) -> bool {
50    let mut chars = ident.chars();
51    let Some(first) = chars.next() else {
52        return false;
53    };
54
55    if !matches!(first, 'a'..='z' | '_') {
56        return false;
57    }
58
59    for ch in chars {
60        if !matches!(ch, 'a'..='z' | '0'..='9' | '_') {
61            return false;
62        }
63    }
64
65    true
66}
67
68fn is_keyword(ident: &str) -> bool {
69    let ident_upper = ident.to_ascii_uppercase();
70    // NOTE: PostgreSQL only quotes keywords whose category != UNRESERVED
71    // (e.g. RESERVED / COL_NAME / TYPE_FUNC). We stay conservative and
72    // quote any token in ALL_KEYWORDS to avoid ambiguity.
73    keywords::ALL_KEYWORDS
74        .binary_search(&ident_upper.as_str())
75        .is_ok()
76}
77
78#[cfg(test)]
79mod tests {
80    use super::QuoteIdent;
81
82    #[track_caller]
83    fn check(input: &str, expect: &str) {
84        assert_eq!(QuoteIdent(input).to_string(), expect);
85    }
86
87    #[test]
88    fn quote_ident_basic() {
89        check("foo", "foo");
90        check("foo_bar", "foo_bar");
91        check("foo1", "foo1");
92        check("1foo", "\"1foo\"");
93        check("foo-bar", "\"foo-bar\"");
94        check("FooBar", "\"FooBar\"");
95    }
96
97    #[test]
98    fn quote_ident_keyword() {
99        check("select", "\"select\"");
100        check("with", "\"with\"");
101    }
102
103    #[test]
104    fn quote_ident_embedded_quote() {
105        check("foo\"bar", "\"foo\"\"bar\"");
106    }
107}
risingwave_sqlparser/quote_ident.rs

risingwave_sqlparser/
quote_ident.rs
