cdk_sql_common/
stmt.rs

1//! Stataments mod
2use std::collections::HashMap;
3use std::sync::{Arc, LazyLock, RwLock};
4
5use cdk_common::database::Error;
6
7use crate::database::DatabaseExecutor;
8use crate::value::Value;
9
10/// The Column type
11pub type Column = Value;
12
13/// Expected response type for a given SQL statement
14#[derive(Debug, Clone, Copy, Default)]
15pub enum ExpectedSqlResponse {
16    /// A single row
17    SingleRow,
18    /// All the rows that matches a query
19    #[default]
20    ManyRows,
21    /// How many rows were affected by the query
22    AffectedRows,
23    /// Return the first column of the first row
24    Pluck,
25    /// Batch
26    Batch,
27}
28
29/// Part value
30#[derive(Debug, Clone)]
31pub enum PlaceholderValue {
32    /// Value
33    Value(Value),
34    /// Set
35    Set(Vec<Value>),
36}
37
38impl From<Value> for PlaceholderValue {
39    fn from(value: Value) -> Self {
40        PlaceholderValue::Value(value)
41    }
42}
43
44impl From<Vec<Value>> for PlaceholderValue {
45    fn from(value: Vec<Value>) -> Self {
46        PlaceholderValue::Set(value)
47    }
48}
49
50/// SQL Part
51#[derive(Debug, Clone)]
52pub enum SqlPart {
53    /// Raw SQL statement
54    Raw(Arc<str>),
55    /// Placeholder
56    Placeholder(Arc<str>, Option<PlaceholderValue>),
57}
58
59/// SQL parser error
60#[derive(Debug, PartialEq, thiserror::Error)]
61pub enum SqlParseError {
62    /// Invalid SQL
63    #[error("Unterminated String literal")]
64    UnterminatedStringLiteral,
65    /// Invalid placeholder name
66    #[error("Invalid placeholder name")]
67    InvalidPlaceholder,
68}
69
70/// Rudimentary SQL parser.
71///
72/// This function does not validate the SQL statement, it only extracts the placeholder to be
73/// database agnostic.
74pub fn split_sql_parts(input: &str) -> Result<Vec<SqlPart>, SqlParseError> {
75    let mut parts = Vec::new();
76    let mut current = String::new();
77    let mut chars = input.chars().peekable();
78
79    while let Some(&c) = chars.peek() {
80        match c {
81            '\'' | '"' => {
82                // Start of string literal
83                let quote = c;
84                current.push(
85                    chars
86                        .next()
87                        .ok_or(SqlParseError::UnterminatedStringLiteral)?,
88                );
89
90                let mut closed = false;
91                while let Some(&next) = chars.peek() {
92                    current.push(
93                        chars
94                            .next()
95                            .ok_or(SqlParseError::UnterminatedStringLiteral)?,
96                    );
97
98                    if next == quote {
99                        if chars.peek() == Some(&quote) {
100                            // Escaped quote (e.g. '' inside strings)
101                            current.push(
102                                chars
103                                    .next()
104                                    .ok_or(SqlParseError::UnterminatedStringLiteral)?,
105                            );
106                        } else {
107                            closed = true;
108                            break;
109                        }
110                    }
111                }
112
113                if !closed {
114                    return Err(SqlParseError::UnterminatedStringLiteral);
115                }
116            }
117
118            '-' => {
119                current.push(
120                    chars
121                        .next()
122                        .ok_or(SqlParseError::UnterminatedStringLiteral)?,
123                );
124
125                if chars.peek() == Some(&'-') {
126                    while let Some(&next) = chars.peek() {
127                        current.push(
128                            chars
129                                .next()
130                                .ok_or(SqlParseError::UnterminatedStringLiteral)?,
131                        );
132                        if next == '\n' {
133                            break;
134                        }
135                    }
136                }
137            }
138
139            ':' => {
140                chars.next(); // consume ':'
141
142                if chars.peek() == Some(&':') {
143                    current.push(':');
144                    current.push(
145                        chars
146                            .next()
147                            .ok_or(SqlParseError::UnterminatedStringLiteral)?,
148                    );
149                    continue;
150                }
151
152                // Flush current raw SQL
153                if !current.is_empty() {
154                    parts.push(SqlPart::Raw(current.clone().into()));
155                    current.clear();
156                }
157
158                let mut name = String::new();
159
160                while let Some(&next) = chars.peek() {
161                    if next.is_alphanumeric() || next == '_' {
162                        name.push(
163                            chars
164                                .next()
165                                .ok_or(SqlParseError::UnterminatedStringLiteral)?,
166                        );
167                    } else {
168                        break;
169                    }
170                }
171
172                if name.is_empty() {
173                    return Err(SqlParseError::InvalidPlaceholder);
174                }
175
176                parts.push(SqlPart::Placeholder(name.into(), None));
177            }
178
179            _ => {
180                current.push(
181                    chars
182                        .next()
183                        .ok_or(SqlParseError::UnterminatedStringLiteral)?,
184                );
185            }
186        }
187    }
188
189    if !current.is_empty() {
190        parts.push(SqlPart::Raw(current.into()));
191    }
192
193    Ok(parts)
194}
195
196type Cache = HashMap<String, (Vec<SqlPart>, Option<Arc<str>>)>;
197
198/// Sql message
199#[derive(Debug, Default)]
200pub struct Statement {
201    cache: Arc<RwLock<Cache>>,
202    cached_sql: Option<Arc<str>>,
203    sql: Option<String>,
204    /// The SQL statement
205    pub parts: Vec<SqlPart>,
206    /// The expected response type
207    pub expected_response: ExpectedSqlResponse,
208}
209
210impl Statement {
211    /// Creates a new statement
212    fn new(sql: &str, cache: Arc<RwLock<Cache>>) -> Result<Self, SqlParseError> {
213        let parsed = cache
214            .read()
215            .map(|cache| cache.get(sql).cloned())
216            .ok()
217            .flatten();
218
219        if let Some((parts, cached_sql)) = parsed {
220            Ok(Self {
221                parts,
222                cached_sql,
223                sql: None,
224                cache,
225                ..Default::default()
226            })
227        } else {
228            let parts = split_sql_parts(sql)?;
229
230            if let Ok(mut cache) = cache.write() {
231                cache.insert(sql.to_owned(), (parts.clone(), None));
232            } else {
233                tracing::warn!("Failed to acquire write lock for SQL statement cache");
234            }
235
236            Ok(Self {
237                parts,
238                sql: Some(sql.to_owned()),
239                cache,
240                ..Default::default()
241            })
242        }
243    }
244
245    /// Convert Statement into a SQL statement and the list of placeholders
246    ///
247    /// By default it converts the statement into placeholder using $1..$n placeholders which seems
248    /// to be more widely supported, although it can be reimplemented with other formats since part
249    /// is public
250    pub fn to_sql(self) -> Result<(String, Vec<Value>), Error> {
251        let has_set_placeholder = self.parts.iter().any(|part| {
252            matches!(
253                part,
254                SqlPart::Placeholder(_, Some(PlaceholderValue::Set(_)))
255            )
256        });
257
258        if let (false, Some(cached_sql)) = (has_set_placeholder, self.cached_sql) {
259            let sql = cached_sql.to_string();
260            let values = self
261                .parts
262                .into_iter()
263                .map(|x| match x {
264                    SqlPart::Placeholder(name, value) => {
265                        match value.ok_or(Error::MissingPlaceholder(name.to_string()))? {
266                            PlaceholderValue::Value(value) => Ok(vec![value]),
267                            PlaceholderValue::Set(values) => Ok(values),
268                        }
269                    }
270                    SqlPart::Raw(_) => Ok(vec![]),
271                })
272                .collect::<Result<Vec<_>, Error>>()?
273                .into_iter()
274                .flatten()
275                .collect::<Vec<_>>();
276            return Ok((sql, values));
277        }
278
279        let mut placeholder_values = Vec::new();
280        let mut can_be_cached = true;
281        let sql = self
282            .parts
283            .into_iter()
284            .map(|x| match x {
285                SqlPart::Placeholder(name, value) => {
286                    match value.ok_or(Error::MissingPlaceholder(name.to_string()))? {
287                        PlaceholderValue::Value(value) => {
288                            placeholder_values.push(value);
289                            Ok::<_, Error>(format!("${}", placeholder_values.len()))
290                        }
291                        PlaceholderValue::Set(mut values) => {
292                            can_be_cached = false;
293                            let start_size = placeholder_values.len();
294                            placeholder_values.append(&mut values);
295                            let placeholders = (start_size + 1..=placeholder_values.len())
296                                .map(|i| format!("${i}"))
297                                .collect::<Vec<_>>()
298                                .join(", ");
299                            Ok(placeholders)
300                        }
301                    }
302                }
303                SqlPart::Raw(raw) => Ok(raw.trim().to_string()),
304            })
305            .collect::<Result<Vec<String>, _>>()?
306            .join(" ");
307
308        if can_be_cached {
309            if let Some(original_sql) = self.sql {
310                let _ = self.cache.write().map(|mut cache| {
311                    if let Some((_, cached_sql)) = cache.get_mut(&original_sql) {
312                        *cached_sql = Some(sql.clone().into());
313                    }
314                });
315            }
316        }
317
318        Ok((sql, placeholder_values))
319    }
320
321    /// Binds a given placeholder to a value.
322    #[inline]
323    pub fn bind<C, V>(mut self, name: C, value: V) -> Self
324    where
325        C: ToString,
326        V: Into<Value>,
327    {
328        let name = name.to_string();
329        let value = value.into();
330        let value: PlaceholderValue = value.into();
331
332        for part in self.parts.iter_mut() {
333            if let SqlPart::Placeholder(part_name, part_value) = part {
334                if **part_name == *name.as_str() {
335                    *part_value = Some(value.clone());
336                }
337            }
338        }
339
340        self
341    }
342
343    /// Binds a single variable with a vector.
344    ///
345    /// This will rewrite the function from `:foo` (where value is vec![1, 2, 3]) to `:foo0, :foo1,
346    /// :foo2` and binds each value from the value vector accordingly.
347    ///
348    /// Returns an error if the vector is empty, as empty `IN` clauses produce invalid SQL.
349    #[inline]
350    pub fn bind_vec<C, V>(mut self, name: C, value: Vec<V>) -> Result<Self, Error>
351    where
352        C: ToString,
353        V: Into<Value>,
354    {
355        let name = name.to_string();
356
357        if value.is_empty() {
358            return Err(Error::EmptyInClause(name));
359        }
360
361        let value: PlaceholderValue = value
362            .into_iter()
363            .map(|x| x.into())
364            .collect::<Vec<Value>>()
365            .into();
366
367        for part in self.parts.iter_mut() {
368            if let SqlPart::Placeholder(part_name, part_value) = part {
369                if **part_name == *name.as_str() {
370                    *part_value = Some(value.clone());
371                }
372            }
373        }
374
375        Ok(self)
376    }
377
378    /// Executes a query and returns the affected rows
379    pub async fn pluck<C>(self, conn: &C) -> Result<Option<Value>, Error>
380    where
381        C: DatabaseExecutor,
382    {
383        conn.pluck(self).await
384    }
385
386    /// Executes a query and returns the affected rows
387    pub async fn batch<C>(self, conn: &C) -> Result<(), Error>
388    where
389        C: DatabaseExecutor,
390    {
391        conn.batch(self).await
392    }
393
394    /// Executes a query and returns the affected rows
395    pub async fn execute<C>(self, conn: &C) -> Result<usize, Error>
396    where
397        C: DatabaseExecutor,
398    {
399        conn.execute(self).await
400    }
401
402    /// Runs the query and returns the first row or None
403    pub async fn fetch_one<C>(self, conn: &C) -> Result<Option<Vec<Column>>, Error>
404    where
405        C: DatabaseExecutor,
406    {
407        conn.fetch_one(self).await
408    }
409
410    /// Runs the query and returns the first row or None
411    pub async fn fetch_all<C>(self, conn: &C) -> Result<Vec<Vec<Column>>, Error>
412    where
413        C: DatabaseExecutor,
414    {
415        conn.fetch_all(self).await
416    }
417}
418
419/// Creates a new query statement
420#[inline(always)]
421pub fn query(sql: &str) -> Result<Statement, Error> {
422    static CACHE: LazyLock<Arc<RwLock<Cache>>> =
423        LazyLock::new(|| Arc::new(RwLock::new(HashMap::new())));
424    Statement::new(sql, CACHE.clone()).map_err(|e| Error::Database(Box::new(e)))
425}
426
427#[cfg(test)]
428mod tests {
429    use super::*;
430
431    #[test]
432    fn bind_vec_errors_on_empty_vec() {
433        let stmt = query("SELECT * FROM foo WHERE id IN (:ids)").unwrap();
434        let result = stmt.bind_vec("ids", Vec::<Vec<u8>>::new());
435        assert!(result.is_err());
436        assert!(matches!(result.unwrap_err(), Error::EmptyInClause(name) if name == "ids"));
437    }
438
439    #[test]
440    fn parser_preserves_postgres_cast_operator() {
441        let stmt = query("SELECT (ord - 1)::int AS matched WHERE id = :id")
442            .unwrap()
443            .bind("id", "quote-id");
444
445        let (sql, values) = stmt.to_sql().unwrap();
446
447        assert_eq!(sql, "SELECT (ord - 1)::int AS matched WHERE id = $1");
448        assert_eq!(values.len(), 1);
449    }
450
451    #[test]
452    fn bind_vec_ignores_cached_sql_for_same_query_string() {
453        let raw_sql = "SELECT * FROM cached_sql_vec_bug WHERE id IN (:ids)";
454
455        let (cached_sql, cached_values) =
456            query(raw_sql).unwrap().bind("ids", 1_i64).to_sql().unwrap();
457        assert!(cached_sql.contains("$1"));
458        assert_eq!(cached_values.len(), 1);
459
460        let (sql, values) = query(raw_sql)
461            .unwrap()
462            .bind_vec("ids", vec![1_i64, 2, 3])
463            .unwrap()
464            .to_sql()
465            .unwrap();
466
467        assert!(sql.contains("$1, $2, $3"));
468        assert_eq!(values.len(), 3);
469    }
470}
cdk_sql_common/stmt.rs

cdk_sql_common/
stmt.rs
