Skip to main content

squawk_syntax/
quote.rs

1use crate::SyntaxNode;
2use crate::generated::keywords::{
3    AS_LABEL_KEYWORDS, COL_NAME_KEYWORDS, RESERVED_KEYWORDS, TYPE_FUNC_NAME_KEYWORDS,
4};
5
6pub fn quote_string_literal(text: &str) -> String {
7    format!("'{}'", text.replace('\'', "''"))
8}
9
10pub fn quote_ident_always(text: &str) -> String {
11    format!(r#""{}""#, text.replace('"', r#""""#))
12}
13
14pub fn quote_column_alias(text: &str) -> String {
15    if needs_quoting(text) {
16        quote_ident_always(text)
17    } else {
18        text.to_string()
19    }
20}
21
22pub fn quote_bare_column_alias(text: &str) -> String {
23    if needs_quoting(text) || is_as_label_word(text) {
24        quote_ident_always(text)
25    } else {
26        text.to_string()
27    }
28}
29
30pub fn quote_ident(text: &str) -> String {
31    if needs_quoting(text) || is_reserved_word(text) || is_type_func_name_word(text) {
32        quote_ident_always(text)
33    } else {
34        text.to_string()
35    }
36}
37
38pub fn quote_quoted_ident(text: &str) -> String {
39    if is_col_name_word(text) {
40        quote_ident_always(text)
41    } else {
42        quote_ident(text)
43    }
44}
45
46pub fn unquote_ident(node: &SyntaxNode) -> Option<String> {
47    let text = node.text().to_string();
48
49    if !text.starts_with('"') || !text.ends_with('"') {
50        return None;
51    }
52
53    let text = &text[1..text.len() - 1];
54
55    if is_reserved_word(text) || is_type_func_name_word(text) {
56        return None;
57    }
58
59    if text.is_empty() {
60        return None;
61    }
62
63    let mut chars = text.chars();
64
65    // see: https://www.postgresql.org/docs/18/sql-syntax-lexical.html#SQL-SYNTAX-IDENTIFIERS
66    match chars.next() {
67        Some(c) if c.is_lowercase() || c == '_' => {}
68        _ => return None,
69    }
70
71    for c in chars {
72        if c.is_lowercase() || c.is_ascii_digit() || c == '_' || c == '$' {
73            continue;
74        }
75        return None;
76    }
77
78    Some(text.to_string())
79}
80
81pub fn needs_quoting(text: &str) -> bool {
82    if text.is_empty() {
83        return true;
84    }
85
86    // Column labels in AS clauses allow all keywords, so we don't need to check
87    // for reserved words. See PostgreSQL grammar:
88    // ColLabel: IDENT | unreserved_keyword | col_name_keyword | type_func_name_keyword | reserved_keyword
89
90    let mut chars = text.chars();
91
92    match chars.next() {
93        Some(c) if c.is_lowercase() || c == '_' => {}
94        _ => return true,
95    }
96
97    for c in chars {
98        if c.is_lowercase() || c.is_ascii_digit() || c == '_' || c == '$' {
99            continue;
100        }
101        return true;
102    }
103
104    false
105}
106
107pub fn is_reserved_word(text: &str) -> bool {
108    RESERVED_KEYWORDS
109        .binary_search(&text.to_ascii_lowercase().as_str())
110        .is_ok()
111}
112
113pub(crate) fn is_col_name_word(text: &str) -> bool {
114    COL_NAME_KEYWORDS
115        .binary_search(&text.to_ascii_lowercase().as_str())
116        .is_ok()
117}
118
119pub(crate) fn is_type_func_name_word(text: &str) -> bool {
120    TYPE_FUNC_NAME_KEYWORDS
121        .binary_search(&text.to_ascii_lowercase().as_str())
122        .is_ok()
123}
124
125fn is_as_label_word(text: &str) -> bool {
126    AS_LABEL_KEYWORDS
127        .binary_search(&text.to_ascii_lowercase().as_str())
128        .is_ok()
129}
130
131pub fn strip_quotes(text: &str) -> Option<&str> {
132    text.strip_prefix('\'')?.strip_suffix('\'')
133}
134
135pub fn strip_prefixed_quotes(text: &str, prefix: [char; 2]) -> Option<&str> {
136    strip_quotes(text.strip_prefix(prefix)?)
137}
138
139pub fn strip_unicode_esc_prefix(text: &str) -> Option<&str> {
140    strip_quotes(text.strip_prefix(['u', 'U'])?.strip_prefix('&')?)
141}
142
143pub fn dollar_quote_tag(text: &str) -> Option<&str> {
144    text.strip_prefix('$')?.split_once('$').map(|(tag, _)| tag)
145}
146
147pub fn strip_dollar_quotes(text: &str) -> Option<&str> {
148    let tag = dollar_quote_tag(text)?;
149    let body = &text[tag.len() + 2..];
150    let closing = format!("${tag}$");
151    body.strip_suffix(&closing)
152}
153
154#[cfg(test)]
155mod tests {
156    use insta::assert_snapshot;
157
158    use super::*;
159
160    #[test]
161    fn quote_string_literal_escapes_apostrophes() {
162        assert_snapshot!(quote_string_literal("it's"), @"'it''s'");
163    }
164
165    #[test]
166    fn quote_column_alias_handles_embedded_quotes() {
167        assert_snapshot!(quote_column_alias(r#"foo"bar"#), @r#""foo""bar""#);
168    }
169
170    #[test]
171    fn quote_column_alias_doesnt_quote_reserved_words() {
172        // Keywords are allowed as column labels in AS clauses
173        assert_snapshot!(quote_column_alias("case"), @"case");
174        assert_snapshot!(quote_column_alias("array"), @"array");
175    }
176
177    #[test]
178    fn quote_column_alias_doesnt_quote_simple_identifiers() {
179        assert_snapshot!(quote_column_alias("col_name"), @"col_name");
180    }
181
182    #[test]
183    fn quote_column_alias_handles_special_column_name() {
184        assert_snapshot!(quote_column_alias("?column?"), @r#""?column?""#);
185    }
186
187    #[test]
188    fn quote_bare_column_alias_quotes_keywords_that_need_an_as() {
189        assert_snapshot!(quote_bare_column_alias("filter"), @r#""filter""#);
190        assert_snapshot!(quote_bare_column_alias("day"), @r#""day""#);
191        // also reserved
192        assert_snapshot!(quote_bare_column_alias("array"), @r#""array""#);
193    }
194
195    #[test]
196    fn quote_bare_column_alias_doesnt_quote_bare_label_keywords() {
197        assert_snapshot!(quote_bare_column_alias("between"), @"between");
198        assert_snapshot!(quote_bare_column_alias("all"), @"all");
199        assert_snapshot!(quote_bare_column_alias("left"), @"left");
200        assert_snapshot!(quote_bare_column_alias("col_name"), @"col_name");
201    }
202
203    #[test]
204    fn quote_ident_doesnt_quote_simple_identifiers() {
205        assert_snapshot!(quote_ident("col_name"), @"col_name");
206        assert_snapshot!(quote_ident("users"), @"users");
207        assert_snapshot!(quote_ident("t2$"), @"t2$");
208    }
209
210    #[test]
211    fn quote_ident_doesnt_quote_column_or_table_keywords() {
212        // unreserved
213        assert_snapshot!(quote_ident("data"), @"data");
214        assert_snapshot!(quote_ident("value"), @"value");
215        // col name
216        assert_snapshot!(quote_ident("int"), @"int");
217    }
218
219    #[test]
220    fn quote_ident_quotes_reserved_words() {
221        assert_snapshot!(quote_ident("select"), @r#""select""#);
222        assert_snapshot!(quote_ident("array"), @r#""array""#);
223    }
224
225    #[test]
226    fn quote_ident_quotes_type_func_name_words() {
227        assert_snapshot!(quote_ident("left"), @r#""left""#);
228        assert_snapshot!(quote_ident("join"), @r#""join""#);
229    }
230
231    #[test]
232    fn quote_quoted_ident_keeps_col_name_keywords_quoted() {
233        assert_snapshot!(quote_quoted_ident("int"), @r#""int""#);
234        assert_snapshot!(quote_quoted_ident("coalesce"), @r#""coalesce""#);
235        assert_snapshot!(quote_quoted_ident("normalize"), @r#""normalize""#);
236    }
237
238    #[test]
239    fn quote_quoted_ident_unquotes_everything_else() {
240        assert_snapshot!(quote_quoted_ident("col_name"), @"col_name");
241        assert_snapshot!(quote_quoted_ident("data"), @"data");
242        assert_snapshot!(quote_quoted_ident("select"), @r#""select""#);
243        assert_snapshot!(quote_quoted_ident("Mixed"), @r#""Mixed""#);
244    }
245
246    #[test]
247    fn quote_ident_quotes_names_that_dont_fold_to_themselves() {
248        assert_snapshot!(quote_ident("Mixed"), @r#""Mixed""#);
249        assert_snapshot!(quote_ident("has space"), @r#""has space""#);
250        assert_snapshot!(quote_ident(""), @r#""""#);
251        assert_snapshot!(quote_ident(r#"foo"bar"#), @r#""foo""bar""#);
252    }
253}