Skip to main content

postrust_sql/
identifier.rs

1//! Safe SQL identifier handling.
2//!
3//! Provides functions for safely escaping and quoting SQL identifiers
4//! and literals to prevent SQL injection.
5
6/// Escape a SQL identifier (table name, column name, etc.).
7///
8/// This function wraps the identifier in double quotes and escapes
9/// any embedded double quotes by doubling them.
10///
11/// # Examples
12///
13/// ```
14/// use postrust_sql::escape_ident;
15///
16/// assert_eq!(escape_ident("users"), "\"users\"");
17/// assert_eq!(escape_ident("user\"name"), "\"user\"\"name\"");
18/// assert_eq!(escape_ident("My Table"), "\"My Table\"");
19/// ```
20pub fn escape_ident(name: &str) -> String {
21    format!("\"{}\"", name.replace('"', "\"\""))
22}
23
24/// Quote a SQL literal string.
25///
26/// This function wraps the string in single quotes and escapes
27/// any embedded single quotes by doubling them. This is only for
28/// cases where parameterized queries can't be used (e.g., SET commands).
29///
30/// # Warning
31///
32/// Prefer using parameterized queries with `SqlParam` instead of this function
33/// whenever possible. This function should only be used for SQL constructs
34/// that don't support parameters (like some SET commands).
35///
36/// # Examples
37///
38/// ```
39/// use postrust_sql::quote_literal;
40///
41/// assert_eq!(quote_literal("hello"), "'hello'");
42/// assert_eq!(quote_literal("it's"), "'it''s'");
43/// assert_eq!(quote_literal(r"a\b"), r"E'a\\b'");
44/// ```
45///
46/// A string with a backslash in it is written as an escape string, `E'...'`,
47/// with the backslash doubled. A plain `'...'` literal reads a backslash as
48/// itself only while `standard_conforming_strings` is on: a database, role or
49/// server that turns it off reads `\'` as a quote inside the literal, and the
50/// doubled quote after it ends the literal early. An escape string means the
51/// same under either setting.
52pub fn quote_literal(s: &str) -> String {
53    let doubled = s.replace('\'', "''");
54    match s.contains('\\') {
55        true => format!("E'{}'", doubled.replace('\\', "\\\\")),
56        false => format!("'{}'", doubled),
57    }
58}
59
60/// Qualified identifier (schema.name).
61#[derive(Clone, Debug, PartialEq, Eq, Hash)]
62pub struct QualifiedIdentifier {
63    /// The schema. Empty when the identifier is unqualified, in which case
64    /// resolution is left to the connection's `search_path`.
65    pub schema: String,
66    /// The object's own name, unqualified and unquoted.
67    pub name: String,
68}
69
70impl QualifiedIdentifier {
71    /// A `schema.name` identifier.
72    pub fn new(schema: impl Into<String>, name: impl Into<String>) -> Self {
73        Self {
74            schema: schema.into(),
75            name: name.into(),
76        }
77    }
78
79    /// A bare `name`, to be resolved against the connection's `search_path`.
80    ///
81    /// Prefer [`new`](Self::new) wherever the schema is known: which object a
82    /// bare name reaches depends on session state rather than on the query.
83    pub fn unqualified(name: impl Into<String>) -> Self {
84        Self {
85            schema: String::new(),
86            name: name.into(),
87        }
88    }
89}
90
91/// Convert a qualified identifier to a safe SQL string.
92///
93/// # Examples
94///
95/// ```
96/// use postrust_sql::{from_qi, identifier::QualifiedIdentifier};
97///
98/// let qi = QualifiedIdentifier::new("public", "users");
99/// assert_eq!(from_qi(&qi), "\"public\".\"users\"");
100///
101/// let qi = QualifiedIdentifier::unqualified("users");
102/// assert_eq!(from_qi(&qi), "\"users\"");
103/// ```
104pub fn from_qi(qi: &QualifiedIdentifier) -> String {
105    if qi.schema.is_empty() {
106        escape_ident(&qi.name)
107    } else {
108        format!("{}.{}", escape_ident(&qi.schema), escape_ident(&qi.name))
109    }
110}
111
112/// Check if a string is a valid unquoted identifier.
113///
114/// PostgreSQL unquoted identifiers must start with a letter or underscore,
115/// and can contain letters, digits, underscores, and dollar signs.
116pub fn is_valid_identifier(s: &str) -> bool {
117    if s.is_empty() {
118        return false;
119    }
120
121    let mut chars = s.chars();
122    let first = chars.next().unwrap();
123
124    if !first.is_ascii_alphabetic() && first != '_' {
125        return false;
126    }
127
128    chars.all(|c| c.is_ascii_alphanumeric() || c == '_' || c == '$')
129}
130
131/// Check if a string is a SQL keyword that should be quoted.
132pub fn is_keyword(s: &str) -> bool {
133    const KEYWORDS: &[&str] = &[
134        "all",
135        "and",
136        "any",
137        "array",
138        "as",
139        "asc",
140        "between",
141        "by",
142        "case",
143        "cast",
144        "check",
145        "column",
146        "constraint",
147        "create",
148        "cross",
149        "current",
150        "default",
151        "delete",
152        "desc",
153        "distinct",
154        "drop",
155        "else",
156        "end",
157        "exists",
158        "false",
159        "for",
160        "foreign",
161        "from",
162        "full",
163        "group",
164        "having",
165        "in",
166        "index",
167        "inner",
168        "insert",
169        "into",
170        "is",
171        "join",
172        "key",
173        "left",
174        "like",
175        "limit",
176        "not",
177        "null",
178        "offset",
179        "on",
180        "or",
181        "order",
182        "outer",
183        "primary",
184        "references",
185        "right",
186        "select",
187        "set",
188        "table",
189        "then",
190        "to",
191        "true",
192        "union",
193        "unique",
194        "update",
195        "using",
196        "values",
197        "when",
198        "where",
199        "with",
200    ];
201
202    KEYWORDS.contains(&s.to_lowercase().as_str())
203}
204
205#[cfg(test)]
206mod tests {
207    use super::*;
208
209    #[test]
210    fn test_escape_ident() {
211        assert_eq!(escape_ident("users"), "\"users\"");
212        assert_eq!(escape_ident("user_table"), "\"user_table\"");
213        assert_eq!(escape_ident("user\"name"), "\"user\"\"name\"");
214        assert_eq!(escape_ident("My Table"), "\"My Table\"");
215        assert_eq!(escape_ident(""), "\"\"");
216    }
217
218    #[test]
219    fn test_quote_literal() {
220        assert_eq!(quote_literal("hello"), "'hello'");
221        assert_eq!(quote_literal("it's"), "'it''s'");
222        assert_eq!(quote_literal(""), "''");
223        // Whatever `standard_conforming_strings` is.
224        assert_eq!(quote_literal(r"\"), r"E'\\'");
225        assert_eq!(quote_literal(r"\' OR true --"), r"E'\\'' OR true --'");
226    }
227
228    #[test]
229    fn test_from_qi() {
230        let qi = QualifiedIdentifier::new("public", "users");
231        assert_eq!(from_qi(&qi), "\"public\".\"users\"");
232
233        let qi = QualifiedIdentifier::unqualified("users");
234        assert_eq!(from_qi(&qi), "\"users\"");
235
236        let qi = QualifiedIdentifier::new("my schema", "my\"table");
237        assert_eq!(from_qi(&qi), "\"my schema\".\"my\"\"table\"");
238    }
239
240    #[test]
241    fn test_is_valid_identifier() {
242        assert!(is_valid_identifier("users"));
243        assert!(is_valid_identifier("_private"));
244        assert!(is_valid_identifier("user123"));
245        assert!(is_valid_identifier("user$table"));
246
247        assert!(!is_valid_identifier(""));
248        assert!(!is_valid_identifier("123users"));
249        assert!(!is_valid_identifier("my-table"));
250        assert!(!is_valid_identifier("my table"));
251    }
252
253    #[test]
254    fn test_is_keyword() {
255        assert!(is_keyword("select"));
256        assert!(is_keyword("SELECT"));
257        assert!(is_keyword("from"));
258        assert!(is_keyword("WHERE"));
259
260        assert!(!is_keyword("users"));
261        assert!(!is_keyword("my_column"));
262    }
263}