Skip to main content

sea_query/
prepare.rs

1//! Helper for preparing SQL statements.
2
3use crate::*;
4pub use std::fmt::Write;
5
6pub trait SqlWriter: Write + Sized + ToString {
7    fn push_param<T: QueryBuilder>(&mut self, value: Value, query_builder: &T);
8
9    /// Upcast this into parent trait. Still needed in 1.85
10    fn as_writer(&mut self) -> &mut dyn Write;
11}
12
13impl SqlWriter for String {
14    fn push_param<T: QueryBuilder>(&mut self, value: Value, query_builder: &T) {
15        query_builder.write_value(self, &value).unwrap();
16    }
17
18    fn as_writer(&mut self) -> &mut dyn Write {
19        self as _
20    }
21}
22
23#[derive(Debug, Clone)]
24pub struct SqlWriterValues {
25    counter: usize,
26    placeholder: String,
27    numbered: bool,
28    string: String,
29    values: Vec<Value>,
30}
31
32impl SqlWriterValues {
33    pub fn new<T>(placeholder: T, numbered: bool) -> Self
34    where
35        T: Into<String>,
36    {
37        Self {
38            counter: 0,
39            placeholder: placeholder.into(),
40            numbered,
41            string: String::with_capacity(256),
42            values: Vec::new(),
43        }
44    }
45
46    pub fn into_parts(self) -> (String, Values) {
47        (self.string, Values(self.values))
48    }
49}
50
51impl Write for SqlWriterValues {
52    #[inline]
53    fn write_str(&mut self, s: &str) -> std::fmt::Result {
54        self.string.write_str(s)
55    }
56
57    #[inline]
58    fn write_char(&mut self, c: char) -> std::fmt::Result {
59        self.string.write_char(c)
60    }
61}
62
63impl std::fmt::Display for SqlWriterValues {
64    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
65        f.write_str(&self.string)
66    }
67}
68
69impl SqlWriter for SqlWriterValues {
70    fn push_param<T: QueryBuilder>(&mut self, value: Value, _: &T) {
71        self.string.push_str(&self.placeholder);
72        if self.numbered {
73            self.counter += 1;
74            write_int(&mut self.string, self.counter);
75        }
76        self.values.push(value)
77    }
78
79    fn as_writer(&mut self) -> &mut dyn Write {
80        self as _
81    }
82}
83
84#[cfg(feature = "itoa")]
85#[inline]
86pub(crate) fn write_int(w: &mut (impl Write + ?Sized), n: impl itoa::Integer) {
87    let mut buf = itoa::Buffer::new();
88    let s = buf.format(n);
89    w.write_str(s).unwrap();
90}
91
92#[cfg(not(feature = "itoa"))]
93#[inline(always)]
94pub(crate) fn write_int(w: &mut (impl Write + ?Sized), n: impl std::fmt::Display) {
95    write!(w, "{n}").unwrap();
96}
97
98pub fn inject_parameters(sql: &str, params: &[Value], query_builder: &impl QueryBuilder) -> String {
99    let mut counter = 0;
100    let mut output = String::new();
101
102    let mut tokenizer = Tokenizer::new(sql)
103        .for_query_builder(query_builder)
104        .iter()
105        .peekable();
106
107    while let Some(token) = tokenizer.next() {
108        match token {
109            Token::Punctuation(mark) => {
110                let (ph, numbered) = query_builder.placeholder();
111
112                if !numbered && mark == ph {
113                    query_builder
114                        .write_value(&mut output, &params[counter])
115                        .unwrap();
116
117                    counter += 1;
118                    continue;
119                } else if numbered && mark == ph {
120                    if let Some(Token::Unquoted(next)) = tokenizer.peek() {
121                        if let Ok(num) = next.parse::<usize>() {
122                            query_builder
123                                .write_value(&mut output, &params[num - 1])
124                                .unwrap();
125
126                            tokenizer.next();
127                            continue;
128                        }
129                    }
130                }
131                output.push_str(mark.as_ref());
132            }
133            _ => output.write_str(token.as_str()).unwrap(),
134        }
135    }
136
137    output
138}
139
140#[cfg(test)]
141#[cfg(feature = "backend-mysql")]
142mod tests_mysql {
143    use super::*;
144    use pretty_assertions::assert_eq;
145
146    #[test]
147    fn inject_parameters_1() {
148        assert_eq!(
149            inject_parameters("WHERE A = ?", &["B".into()], &MysqlQueryBuilder),
150            "WHERE A = 'B'"
151        );
152    }
153
154    #[test]
155    fn inject_parameters_2() {
156        assert_eq!(
157            inject_parameters("WHERE A = '?' AND B = ?", &["C".into()], &MysqlQueryBuilder),
158            "WHERE A = '?' AND B = 'C'"
159        );
160    }
161
162    #[test]
163    fn inject_parameters_3() {
164        assert_eq!(
165            inject_parameters(
166                "WHERE A = ? AND C = ?",
167                &["B".into(), "D".into()],
168                &MysqlQueryBuilder
169            ),
170            "WHERE A = 'B' AND C = 'D'"
171        );
172    }
173
174    #[test]
175    fn inject_parameters_4() {
176        assert_eq!(
177            inject_parameters("?", &[vec![0xABu8, 0xCD, 0xEF].into()], &MysqlQueryBuilder),
178            "x'ABCDEF'"
179        );
180    }
181}
182
183#[cfg(test)]
184#[cfg(feature = "backend-postgres")]
185mod tests_postgres {
186    use super::*;
187    use pretty_assertions::assert_eq;
188
189    #[test]
190    fn inject_parameters_5() {
191        assert_eq!(
192            inject_parameters(
193                "WHERE A = $1 AND C = $2",
194                &["B".into(), "D".into()],
195                &PostgresQueryBuilder
196            ),
197            "WHERE A = 'B' AND C = 'D'"
198        );
199    }
200
201    #[test]
202    fn inject_parameters_6() {
203        assert_eq!(
204            inject_parameters(
205                "WHERE A = $2 AND C = $1",
206                &["B".into(), "D".into()],
207                &PostgresQueryBuilder
208            ),
209            "WHERE A = 'D' AND C = 'B'"
210        );
211    }
212
213    #[test]
214    fn inject_parameters_7() {
215        assert_eq!(
216            inject_parameters("WHERE A = $1", &[Value::from("B'C")], &PostgresQueryBuilder),
217            "WHERE A = E'B\\'C'"
218        );
219    }
220}