drizzle_core/
sql.rs

1mod chunk;
2mod cte;
3mod owned;
4mod tokens;
5
6use crate::ToSQL;
7use crate::prelude::*;
8use crate::{
9    Param, ParamBind, Placeholder,
10    traits::{SQLColumnInfo, SQLParam, SQLTableInfo},
11};
12pub use chunk::*;
13use core::fmt::Display;
14pub use owned::*;
15use smallvec::SmallVec;
16pub use tokens::*;
17
18#[cfg(feature = "profiling")]
19use crate::profile_sql;
20
21/// SQL fragment builder with flat chunk storage.
22///
23/// Uses `SmallVec<[SQLChunk; 8]>` for inline storage of typical SQL fragments
24/// without heap allocation.
25#[derive(Debug, Clone)]
26pub struct SQL<'a, V: SQLParam> {
27    pub chunks: SmallVec<[SQLChunk<'a, V>; 8]>,
28}
29
30impl<'a, V: SQLParam> SQL<'a, V> {
31    const POSITIONAL_PLACEHOLDER: Placeholder = Placeholder::positional();
32
33    // ==================== constructors ====================
34
35    /// Creates an empty SQL fragment
36    #[inline]
37    pub const fn empty() -> Self {
38        Self {
39            chunks: SmallVec::new_const(),
40        }
41    }
42
43    /// Creates SQL with raw text at const time.
44    #[inline]
45    pub const fn raw_const(text: &'static str) -> Self {
46        Self {
47            chunks: SmallVec::from_const([
48                SQLChunk::Raw(Cow::Borrowed(text)),
49                SQLChunk::Empty,
50                SQLChunk::Empty,
51                SQLChunk::Empty,
52                SQLChunk::Empty,
53                SQLChunk::Empty,
54                SQLChunk::Empty,
55                SQLChunk::Empty,
56            ]),
57        }
58    }
59
60    // ==================== constructors ====================
61
62    /// Creates SQL with a single token
63    #[inline]
64    pub fn token(t: Token) -> Self {
65        Self {
66            chunks: smallvec::smallvec![SQLChunk::Token(t)],
67        }
68    }
69
70    /// Creates SQL with a quoted identifier
71    #[inline]
72    pub fn ident(name: impl Into<Cow<'a, str>>) -> Self {
73        Self {
74            chunks: smallvec::smallvec![SQLChunk::Ident(name.into())],
75        }
76    }
77
78    /// Creates SQL with raw text (unquoted)
79    #[inline]
80    pub fn raw(text: impl Into<Cow<'a, str>>) -> Self {
81        Self {
82            chunks: smallvec::smallvec![SQLChunk::Raw(text.into())],
83        }
84    }
85
86    /// Creates SQL with a single parameter value
87    #[inline]
88    pub fn param(value: impl Into<Cow<'a, V>>) -> Self {
89        Self {
90            chunks: smallvec::smallvec![SQLChunk::Param(Param {
91                value: Some(value.into()),
92                placeholder: Self::POSITIONAL_PLACEHOLDER,
93            })],
94        }
95    }
96
97    /// Creates SQL with a named placeholder (no value, for prepared statements)
98    #[inline]
99    pub fn placeholder(name: &'static str) -> Self {
100        Self {
101            chunks: smallvec::smallvec![SQLChunk::Param(Param {
102                value: None,
103                placeholder: Placeholder::colon(name),
104            })],
105        }
106    }
107
108    /// Creates SQL referencing a table
109    #[inline]
110    pub fn table(table: &'static dyn SQLTableInfo) -> Self {
111        Self {
112            chunks: smallvec::smallvec![SQLChunk::Table(table)],
113        }
114    }
115
116    /// Creates SQL referencing a column
117    #[inline]
118    pub fn column(column: &'static dyn SQLColumnInfo) -> Self {
119        Self {
120            chunks: smallvec::smallvec![SQLChunk::Column(column)],
121        }
122    }
123
124    /// Creates SQL for a function call: NAME(args)
125    /// Subqueries are automatically wrapped in parentheses: NAME((SELECT ...))
126    #[inline]
127    pub fn func(name: impl Into<Cow<'a, str>>, args: SQL<'a, V>) -> Self {
128        let args = if args.is_subquery() {
129            args.parens()
130        } else {
131            args
132        };
133        SQL::raw(name)
134            .push(Token::LPAREN)
135            .append(args)
136            .push(Token::RPAREN)
137    }
138
139    // ==================== builder methods ====================
140
141    /// Append another SQL fragment (flat extend)
142    #[inline]
143    pub fn append(mut self, other: impl Into<SQL<'a, V>>) -> Self {
144        #[cfg(feature = "profiling")]
145        profile_sql!("append");
146        self.chunks.extend(other.into().chunks);
147        self
148    }
149
150    /// Push a single chunk
151    #[inline]
152    pub fn push(mut self, chunk: impl Into<SQLChunk<'a, V>>) -> Self {
153        self.chunks.push(chunk.into());
154        self
155    }
156
157    /// Pre-allocates capacity for additional chunks
158    #[inline]
159    pub fn with_capacity(mut self, additional: usize) -> Self {
160        self.chunks.reserve(additional);
161        self
162    }
163
164    // ==================== combinators ====================
165
166    /// Joins multiple SQL fragments with a separator
167    pub fn join<T>(sqls: T, separator: Token) -> SQL<'a, V>
168    where
169        T: IntoIterator,
170        T::Item: ToSQL<'a, V>,
171    {
172        #[cfg(feature = "profiling")]
173        profile_sql!("join");
174
175        let mut iter = sqls.into_iter();
176        let Some(first) = iter.next() else {
177            return SQL::empty();
178        };
179
180        let mut result = first.to_sql();
181        for item in iter {
182            result = result.push(separator).append(item.to_sql());
183        }
184        result
185    }
186
187    /// Wrap in parentheses: (self)
188    #[inline]
189    pub fn parens(self) -> Self {
190        SQL::token(Token::LPAREN).append(self).push(Token::RPAREN)
191    }
192
193    /// Check if this SQL fragment is a subquery (starts with SELECT)
194    #[inline]
195    pub fn is_subquery(&self) -> bool {
196        matches!(self.chunks.first(), Some(SQLChunk::Token(Token::SELECT)))
197    }
198
199    /// Creates an aliased version: self AS "name"
200    pub fn alias(mut self, name: impl Into<Cow<'a, str>>) -> SQL<'a, V> {
201        if self.chunks.len() == 1 {
202            let inner = self.chunks.pop().unwrap();
203            SQL {
204                chunks: smallvec::smallvec![SQLChunk::Alias {
205                    inner: Box::new(inner),
206                    alias: name.into(),
207                }],
208            }
209        } else {
210            self.push(Token::AS).push(SQLChunk::Ident(name.into()))
211        }
212    }
213
214    /// Creates a comma-separated list of parameters
215    pub fn param_list<I>(values: I) -> Self
216    where
217        I: IntoIterator,
218        I::Item: Into<Cow<'a, V>>,
219    {
220        Self::join(values.into_iter().map(Self::param), Token::COMMA)
221    }
222
223    /// Creates a comma-separated list of column assignments: "col" = ?
224    pub fn assignments<I, T>(pairs: I) -> Self
225    where
226        I: IntoIterator<Item = (&'static str, T)>,
227        T: Into<Cow<'a, V>>,
228    {
229        Self::join(
230            pairs
231                .into_iter()
232                .map(|(col, val)| SQL::ident(col).push(Token::EQ).append(SQL::param(val))),
233            Token::COMMA,
234        )
235    }
236
237    // ==================== output methods ====================
238
239    pub fn into_owned(&self) -> OwnedSQL<V> {
240        OwnedSQL::from(self.clone())
241    }
242
243    /// Returns the SQL string
244    pub fn sql(&self) -> String {
245        #[cfg(feature = "profiling")]
246        profile_sql!("sql");
247        let capacity = self.estimate_capacity();
248        let mut buf = String::with_capacity(capacity);
249        self.write_to(&mut buf);
250        buf
251    }
252
253    /// Write SQL to a buffer
254    pub fn write_to(&self, buf: &mut impl core::fmt::Write) {
255        for (i, chunk) in self.chunks.iter().enumerate() {
256            self.write_chunk_to(buf, chunk, i);
257
258            if self.needs_space(i) {
259                let _ = buf.write_char(' ');
260            }
261        }
262    }
263
264    /// Write a single chunk with pattern detection
265    pub fn write_chunk_to(
266        &self,
267        buf: &mut impl core::fmt::Write,
268        chunk: &SQLChunk<'a, V>,
269        index: usize,
270    ) {
271        match chunk {
272            SQLChunk::Token(Token::SELECT) => {
273                chunk.write(buf);
274                self.write_select_columns(buf, index);
275            }
276            _ => chunk.write(buf),
277        }
278    }
279
280    /// Write appropriate columns for SELECT statement
281    fn write_select_columns(&self, buf: &mut impl core::fmt::Write, select_index: usize) {
282        let chunks = self.chunks.get(select_index + 1..select_index + 3);
283        match chunks {
284            Some([SQLChunk::Token(Token::FROM), SQLChunk::Table(table)]) => {
285                let _ = buf.write_char(' ');
286                self.write_qualified_columns(buf, *table);
287            }
288            Some([SQLChunk::Token(Token::FROM), _]) => {
289                let _ = buf.write_char(' ');
290                let _ = buf.write_str(Token::STAR.as_str());
291            }
292            _ => {}
293        }
294    }
295
296    /// Write fully qualified columns
297    pub fn write_qualified_columns(
298        &self,
299        buf: &mut impl core::fmt::Write,
300        table: &dyn SQLTableInfo,
301    ) {
302        let columns = table.columns();
303        if columns.is_empty() {
304            let _ = buf.write_char('*');
305            return;
306        }
307
308        for (i, col) in columns.iter().enumerate() {
309            if i > 0 {
310                let _ = buf.write_str(", ");
311            }
312            let _ = buf.write_char('"');
313            let _ = buf.write_str(table.name());
314            let _ = buf.write_str("\".\"");
315            let _ = buf.write_str(col.name());
316            let _ = buf.write_char('"');
317        }
318    }
319
320    fn estimate_capacity(&self) -> usize {
321        const PLACEHOLDER_SIZE: usize = 2;
322        const IDENT_OVERHEAD: usize = 2;
323        const COLUMN_OVERHEAD: usize = 5;
324        const ALIAS_OVERHEAD: usize = 6;
325
326        self.chunks
327            .iter()
328            .map(|chunk| match chunk {
329                SQLChunk::Empty => 0,
330                SQLChunk::Token(t) => t.as_str().len(),
331                SQLChunk::Ident(s) => s.len() + IDENT_OVERHEAD,
332                SQLChunk::Raw(s) => s.len(),
333                SQLChunk::Param { .. } => PLACEHOLDER_SIZE,
334                SQLChunk::Table(t) => t.name().len() + IDENT_OVERHEAD,
335                SQLChunk::Column(c) => c.table().name().len() + c.name().len() + COLUMN_OVERHEAD,
336                SQLChunk::Alias { inner, alias } => {
337                    alias.len()
338                        + ALIAS_OVERHEAD
339                        + match inner.as_ref() {
340                            SQLChunk::Ident(s) => s.len() + IDENT_OVERHEAD,
341                            SQLChunk::Raw(s) => s.len(),
342                            _ => 10,
343                        }
344                }
345            })
346            .sum::<usize>()
347            + self.chunks.len()
348    }
349
350    /// Simplified spacing logic
351    fn needs_space(&self, index: usize) -> bool {
352        let Some(next) = self.chunks.get(index + 1) else {
353            return false;
354        };
355
356        let current = &self.chunks[index];
357        chunk_needs_space(current, next)
358    }
359
360    /// Returns references to parameter values
361    pub fn params(&self) -> Vec<&V> {
362        let mut params_vec = Vec::with_capacity(self.chunks.len().min(8));
363        for chunk in &self.chunks {
364            if let SQLChunk::Param(Param {
365                value: Some(value), ..
366            }) = chunk
367            {
368                params_vec.push(value.as_ref());
369            }
370        }
371        params_vec
372    }
373
374    /// Bind named parameters
375    pub fn bind<T: SQLParam + Into<V>>(
376        self,
377        params: impl IntoIterator<Item: Into<ParamBind<'a, T>>>,
378    ) -> SQL<'a, V> {
379        #[cfg(feature = "profiling")]
380        profile_sql!("bind");
381
382        let param_map: HashMap<&str, V> = params
383            .into_iter()
384            .map(Into::into)
385            .map(|p| (p.name, p.value.into()))
386            .collect();
387
388        let bound_chunks: SmallVec<[SQLChunk<'a, V>; 8]> = self
389            .chunks
390            .into_iter()
391            .map(|chunk| match chunk {
392                SQLChunk::Param(mut param) => {
393                    if let Some(name) = param.placeholder.name
394                        && let Some(value) = param_map.get(name)
395                    {
396                        param.value = Some(Cow::Owned(value.clone()));
397                    }
398                    SQLChunk::Param(param)
399                }
400                other => other,
401            })
402            .collect();
403
404        SQL {
405            chunks: bound_chunks,
406        }
407    }
408}
409
410/// Simplified spacing logic
411fn chunk_needs_space<V: SQLParam>(current: &SQLChunk<'_, V>, next: &SQLChunk<'_, V>) -> bool {
412    // No space if current raw text ends with space
413    if let SQLChunk::Raw(text) = current
414        && text.ends_with(' ')
415    {
416        return false;
417    }
418
419    // No space if next raw text starts with space
420    if let SQLChunk::Raw(text) = next
421        && text.starts_with(' ')
422    {
423        return false;
424    }
425
426    match (current, next) {
427        // No space before closing/separator punctuation
428        (_, SQLChunk::Token(Token::RPAREN | Token::COMMA | Token::SEMI | Token::DOT)) => false,
429        // No space after opening punctuation
430        (SQLChunk::Token(Token::LPAREN | Token::DOT), _) => false,
431        // Space after comma
432        (SQLChunk::Token(Token::COMMA), _) => true,
433        // Space after closing paren if next is word-like (e.g., ") FROM")
434        (SQLChunk::Token(Token::RPAREN), next) => next.is_word_like(),
435        // Space before opening paren if preceded by word-like (e.g., "AS (")
436        (current, SQLChunk::Token(Token::LPAREN)) => current.is_word_like(),
437        // Space around comparison/arithmetic operators
438        (SQLChunk::Token(t), _) if t.is_operator() => true,
439        (_, SQLChunk::Token(t)) if t.is_operator() => true,
440        // Space between all word-like chunks
441        _ => current.is_word_like() && next.is_word_like(),
442    }
443}
444
445// ==================== trait implementations ====================
446
447impl<'a, V: SQLParam> Default for SQL<'a, V> {
448    fn default() -> Self {
449        Self::empty()
450    }
451}
452
453impl<'a, V: SQLParam + 'a> From<&'a str> for SQL<'a, V> {
454    fn from(s: &'a str) -> Self {
455        SQL::raw(s)
456    }
457}
458
459impl<'a, V: SQLParam> From<Token> for SQL<'a, V> {
460    fn from(value: Token) -> Self {
461        SQL::token(value)
462    }
463}
464
465impl<'a, V: SQLParam + 'a> AsRef<SQL<'a, V>> for SQL<'a, V> {
466    fn as_ref(&self) -> &SQL<'a, V> {
467        self
468    }
469}
470
471impl<'a, V: SQLParam + core::fmt::Display> Display for SQL<'a, V> {
472    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
473        let params = self.params();
474        write!(f, r#"sql: "{}", params: {:?}"#, self.sql(), params)
475    }
476}
477
478impl<'a, V: SQLParam + 'a> ToSQL<'a, V> for SQL<'a, V> {
479    fn to_sql(&self) -> SQL<'a, V> {
480        self.clone()
481    }
482}
483
484impl<'a, V: SQLParam, T> FromIterator<T> for SQL<'a, V>
485where
486    SQLChunk<'a, V>: From<T>,
487{
488    fn from_iter<I: IntoIterator<Item = T>>(iter: I) -> Self {
489        let chunks = SmallVec::from_iter(iter.into_iter().map(SQLChunk::from));
490        Self { chunks }
491    }
492}
493
494impl<'a, V: SQLParam> IntoIterator for SQL<'a, V> {
495    type Item = SQLChunk<'a, V>;
496    type IntoIter = smallvec::IntoIter<[SQLChunk<'a, V>; 8]>;
497
498    fn into_iter(self) -> Self::IntoIter {
499        self.chunks.into_iter()
500    }
501}