elefant_tools/
quoting.rs

1use std::collections::HashMap;
2
3/// Provides utilities for quoting identifiers in PostgreSQL as needed.
4#[derive(Debug)]
5pub struct IdentifierQuoter {
6    /// Keywords that might need to be escaped, and whether they are allowed to be used as column names or type/function names.
7    keywords: HashMap<String, AllowedKeywordUsage>,
8}
9
10/// How a keyword is allowed to be used.
11#[derive(Debug, Copy, Clone)]
12pub struct AllowedKeywordUsage {
13    pub column_name: bool,
14    pub type_or_function_name: bool,
15}
16
17/// How an identifier is attempted to be used.
18#[derive(Debug, Copy, Clone, Eq, PartialEq)]
19pub enum AttemptedKeywordUsage {
20    ColumnName,
21    TypeOrFunctionName,
22    Other,
23}
24
25impl IdentifierQuoter {
26    /// Creates a new IdentifierQuoter with the specified keywords and their allowed usages.
27    pub fn new(keywords: HashMap<String, AllowedKeywordUsage>) -> Self {
28        Self { keywords }
29    }
30
31    /// Creates a new IdentifierQuoter with no keywords.
32    ///
33    /// This is mainly useful for testing as it doesn't require connecting to Postgres.
34    pub fn empty() -> Self {
35        Self {
36            keywords: HashMap::new(),
37        }
38    }
39
40    /// Quotes an identifier as needed.
41    ///
42    /// Ported from <https://github.com/postgres/postgres/blob/97957fdbaa429c7c582d4753b108cb1e23e1b28a/src/backend/utils/adt/ruleutils.c#L11975>
43    pub fn quote(&self, identifier: impl AsRef<str>, usage: AttemptedKeywordUsage) -> String {
44        let identifier = identifier.as_ref();
45
46        if identifier.is_empty() {
47            return "\"\"".to_string();
48        }
49
50        let mut chars = identifier.chars();
51
52        let safe = if let Some(allowed) = self.keywords.get(identifier) {
53            match usage {
54                AttemptedKeywordUsage::ColumnName => allowed.column_name,
55                AttemptedKeywordUsage::TypeOrFunctionName => allowed.type_or_function_name,
56                AttemptedKeywordUsage::Other => false,
57            }
58        } else {
59            matches!(chars.next(), Some('a'..='z' | '_'))
60                && chars.all(|c| matches!(c, 'a'..='z' | '0'..='9' | '_'))
61        };
62
63        if safe {
64            identifier.to_string()
65        } else {
66            let escaped = identifier.replace('"', r#""""#);
67
68            format!("\"{}\"", escaped)
69        }
70    }
71
72    /// Quotes multiple identifiers as needed.
73    pub fn quote_iter<'a, 's, S: AsRef<str>, I: IntoIterator<Item = S>>(
74        &'a self,
75        identifiers: I,
76        usage: AttemptedKeywordUsage,
77    ) -> impl Iterator<Item = String> + 'a
78    where
79        <I as IntoIterator>::IntoIter: 'a,
80    {
81        identifiers.into_iter().map(move |i| self.quote(i, usage))
82    }
83}
84
85/// A trait for types that can be quoted.
86pub(crate) trait Quotable {
87    /// Quotes the value as needed.
88    fn quote(&self, quoter: &IdentifierQuoter, usage: AttemptedKeywordUsage) -> String;
89}
90
91impl<S> Quotable for S
92where
93    S: AsRef<str>,
94{
95    fn quote(&self, quoter: &IdentifierQuoter, usage: AttemptedKeywordUsage) -> String {
96        quoter.quote(self, usage)
97    }
98}
99
100/// A trait for types that can be quoted as an iterator.
101pub(crate) trait QuotableIter: Sized {
102    fn quote(self, quoter: &IdentifierQuoter, usage: AttemptedKeywordUsage)
103        -> IteratorQuoter<Self>;
104}
105
106impl<I> QuotableIter for I
107where
108    I: Iterator,
109    I::Item: AsRef<str>,
110{
111    fn quote(
112        self,
113        quoter: &IdentifierQuoter,
114        usage: AttemptedKeywordUsage,
115    ) -> IteratorQuoter<Self> {
116        IteratorQuoter {
117            quoter,
118            usage,
119            iter: self,
120        }
121    }
122}
123
124/// The iterator implementation used then quoting an iterator of values
125pub(crate) struct IteratorQuoter<'q, I> {
126    quoter: &'q IdentifierQuoter,
127    usage: AttemptedKeywordUsage,
128    iter: I,
129}
130
131impl<I> Iterator for IteratorQuoter<'_, I>
132where
133    I: Iterator,
134    I::Item: AsRef<str>,
135{
136    type Item = String;
137
138    fn next(&mut self) -> Option<Self::Item> {
139        self.iter.next().map(|i| self.quoter.quote(i, self.usage))
140    }
141}
142
143/// Quotes a a string value for usage in Postgres.
144pub(crate) fn quote_value_string(s: &str) -> String {
145    format!("'{}'", s.replace('\'', "''"))
146}
147
148#[cfg(test)]
149mod tests {
150    use crate::quoting::{AllowedKeywordUsage, AttemptedKeywordUsage};
151    use std::collections::HashMap;
152
153    #[test]
154    fn quoting() {
155        let quoter = super::IdentifierQuoter::new(HashMap::from([(
156            "table".to_string(),
157            AllowedKeywordUsage {
158                type_or_function_name: false,
159                column_name: false,
160            },
161        )]));
162
163        macro_rules! test_quote {
164            ($identifier:literal, $expected:literal) => {
165                let quoted = quoter.quote($identifier, AttemptedKeywordUsage::Other);
166                assert_eq!(quoted, $expected);
167            };
168        }
169
170        test_quote!("table", "\"table\"");
171        test_quote!("table1", "table1");
172        test_quote!("table_1", "table_1");
173        test_quote!("table-1", "\"table-1\"");
174        test_quote!("table 1", "\"table 1\"");
175        test_quote!("1table", "\"1table\"");
176        test_quote!("my_table", "my_table");
177        test_quote!("MyTable", "\"MyTable\"");
178        test_quote!("my\"table", "\"my\"\"table\"");
179        test_quote!("", "\"\"");
180    }
181}
elefant_tools/quoting.rs

elefant_tools/
quoting.rs
