sqlite_collections/
identifier.rs

1use core::fmt;
2use std::{
3    borrow::Cow,
4    ops::{Add, AddAssign, Deref},
5};
6
7mod error;
8pub use error::Error;
9
10/// An identifier wrapper for SQLite identifiers.  This checks validity and
11/// wraps the identifier in quotes with escapes for protection.
12#[derive(Debug, Clone, Eq, PartialEq, PartialOrd, Ord, Hash)]
13pub struct Identifier<'a> {
14    inner: Cow<'a, str>,
15    quoted: String,
16}
17
18impl<'a> TryFrom<Cow<'a, str>> for Identifier<'a> {
19    type Error = Error;
20
21    fn try_from(value: Cow<'a, str>) -> Result<Self, Self::Error> {
22        if value.find('\0').is_some() {
23            return Err(Error::NullCharacter);
24        }
25        let len = 2 + value.chars().filter(|&c| c == '"').count() + value.len();
26        let mut quoted = String::new();
27        quoted.reserve_exact(len);
28        {
29            quoted.push('"');
30            let mut value: &str = &value;
31            loop {
32                match value.find('"') {
33                    Some(index) => {
34                        quoted.push_str(&value[..=index]);
35                        quoted.push('"');
36
37                        value = &value[index + 1..];
38                    }
39                    None => {
40                        quoted.push_str(value);
41                        break;
42                    }
43                }
44            }
45            quoted.push('"');
46        }
47        Ok(Identifier {
48            inner: value,
49            quoted,
50        })
51    }
52}
53
54impl<'a> TryFrom<&'a str> for Identifier<'a> {
55    type Error = Error;
56
57    fn try_from(value: &'a str) -> Result<Self, Self::Error> {
58        Identifier::try_from(Cow::Borrowed(value))
59    }
60}
61
62impl TryFrom<String> for Identifier<'static> {
63    type Error = Error;
64
65    fn try_from(value: String) -> Result<Self, Self::Error> {
66        Identifier::try_from(Cow::Owned(value))
67    }
68}
69
70impl<'a> From<Identifier<'a>> for Cow<'a, str> {
71    fn from(value: Identifier<'a>) -> Self {
72        value.inner
73    }
74}
75
76impl<'a> Deref for Identifier<'a> {
77    type Target = str;
78
79    fn deref(&self) -> &Self::Target {
80        &self.inner
81    }
82}
83
84impl<'a, T> AsRef<T> for Identifier<'a>
85where
86    T: ?Sized,
87    <Identifier<'a> as Deref>::Target: AsRef<T>,
88{
89    fn as_ref(&self) -> &T {
90        self.deref().as_ref()
91    }
92}
93impl<'a, 'b> Add<&Identifier<'b>> for Identifier<'a> {
94    type Output = Identifier<'a>;
95
96    fn add(mut self, rhs: &Identifier) -> Self::Output {
97        self += rhs;
98        self
99    }
100}
101
102impl<'a, 'b> AddAssign<&Identifier<'b>> for Identifier<'a> {
103    fn add_assign(&mut self, rhs: &Identifier) {
104        let inner = self.inner.to_mut();
105        inner.reserve_exact(rhs.inner.len());
106        *inner += &rhs.inner;
107
108        // Don't need rhs's quotation marks.
109        self.quoted.reserve_exact(rhs.quoted.len() - 2);
110        self.quoted.pop();
111        self.quoted += &rhs.quoted[1..];
112    }
113}
114
115impl<'a> fmt::Display for Identifier<'a> {
116    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
117        self.quoted.fmt(f)
118    }
119}
120
121#[cfg(test)]
122mod test {
123    use super::Identifier;
124
125    #[test]
126    fn tests() {
127        assert_eq!(
128            Identifier::try_from("main").unwrap().to_string(),
129            String::from("\"main\""),
130        );
131        assert_eq!(
132            Identifier::try_from(String::from("ma\"in"))
133                .unwrap()
134                .to_string(),
135            String::from("\"ma\"\"in\""),
136        );
137        assert_eq!(
138            Identifier::try_from(String::from("\"main"))
139                .unwrap()
140                .to_string(),
141            String::from("\"\"\"main\""),
142        );
143        assert_eq!(
144            Identifier::try_from(String::from("main\""))
145                .unwrap()
146                .to_string(),
147            String::from("\"main\"\"\""),
148        );
149        assert!(Identifier::try_from(String::from("ma\0in")).is_err());
150    }
151
152    #[test]
153    fn test_add() {
154        assert_eq!(
155            (Identifier::try_from(String::from("main")).unwrap()
156                + &String::from("_after").try_into().unwrap())
157                .to_string(),
158            "\"main_after\"",
159        );
160    }
161}