Skip to main content

drizzle_core/sql/
chunk.rs

1use crate::SQLConstraintKind;
2use crate::prelude::*;
3use crate::{Dialect, Param, Placeholder, SQLParam, sql::tokens::Token};
4
5// ==================== Dialect enums ====================
6
7/// Dialect-specific column metadata. Const-compatible enum grouping
8/// fields that only apply to one dialect.
9#[derive(Clone, Copy, Debug, PartialEq, Eq)]
10pub enum ColumnDialect {
11    SQLite {
12        autoincrement: bool,
13        default: Option<&'static str>,
14        generated_expression: Option<&'static str>,
15        generated_stored: bool,
16        collate: Option<&'static str>,
17    },
18    PostgreSQL {
19        postgres_type: &'static str,
20        dimensions: Option<i32>,
21        is_serial: bool,
22        is_bigserial: bool,
23        is_generated_identity: bool,
24        is_identity_always: bool,
25        default: Option<&'static str>,
26        generated_expression: Option<&'static str>,
27        generated_stored: bool,
28        collate: Option<&'static str>,
29        comment: Option<&'static str>,
30    },
31    MySQL {
32        auto_increment: bool,
33        default: Option<&'static str>,
34        generated_expression: Option<&'static str>,
35        generated_stored: bool,
36        charset: Option<&'static str>,
37        collate: Option<&'static str>,
38        on_update: Option<&'static str>,
39    },
40}
41
42/// Dialect-specific table metadata.
43#[derive(Clone, Copy, Debug, PartialEq, Eq)]
44pub enum TableDialect {
45    PostgreSQL {
46        is_unlogged: bool,
47        is_temporary: bool,
48        inherits: Option<&'static str>,
49        tablespace: Option<&'static str>,
50        is_rls_enabled: bool,
51        comment: Option<&'static str>,
52    },
53    SQLite {
54        without_rowid: bool,
55        strict: bool,
56    },
57    MySQL {
58        is_temporary: bool,
59        engine: Option<&'static str>,
60        charset: Option<&'static str>,
61        collate: Option<&'static str>,
62        comment: Option<&'static str>,
63    },
64}
65
66impl Default for TableDialect {
67    fn default() -> Self {
68        Self::PostgreSQL {
69            is_unlogged: false,
70            is_temporary: false,
71            inherits: None,
72            tablespace: None,
73            is_rls_enabled: false,
74            comment: None,
75        }
76    }
77}
78
79// ==================== Ref structs ====================
80
81/// Foreign key reference as a const Copy struct.
82#[derive(Clone, Copy, Debug, PartialEq, Eq)]
83pub struct ForeignKeyRef {
84    pub name: &'static str,
85    pub name_explicit: bool,
86    pub target_table: &'static str,
87    pub target_schema: &'static str,
88    pub source_columns: &'static [&'static str],
89    pub target_columns: &'static [&'static str],
90    pub on_delete: Option<&'static str>,
91    pub on_update: Option<&'static str>,
92    pub deferrable: bool,
93    pub initially_deferred: bool,
94}
95
96/// Primary key reference as a const Copy struct.
97#[derive(Clone, Copy, Debug, PartialEq, Eq)]
98pub struct PrimaryKeyRef {
99    pub columns: &'static [&'static str],
100}
101
102/// Constraint reference as a const Copy struct.
103#[derive(Clone, Copy, Debug, PartialEq, Eq)]
104pub struct ConstraintRef {
105    pub name: Option<&'static str>,
106    pub name_explicit: bool,
107    pub kind: SQLConstraintKind,
108    pub columns: &'static [&'static str],
109    pub check_expression: Option<&'static str>,
110    pub deferrable: bool,
111    pub initially_deferred: bool,
112}
113
114// ==================== Enhanced TableRef and ColumnRef ====================
115
116/// Table reference with full schema metadata.
117///
118/// Carries both the SQL rendering fields (`name`, `column_names`) and
119/// complete schema metadata (columns, keys, constraints). SQL rendering
120/// code only uses `name`/`column_names` and ignores extra fields.
121#[derive(Clone, Copy, Debug, PartialEq, Eq)]
122pub struct TableRef {
123    // SQL rendering fields
124    pub name: &'static str,
125    pub column_names: &'static [&'static str],
126
127    // Schema metadata
128    pub schema: Option<&'static str>,
129    pub qualified_name: &'static str,
130    pub columns: &'static [ColumnRef],
131    pub primary_key: Option<PrimaryKeyRef>,
132    pub foreign_keys: &'static [ForeignKeyRef],
133    pub constraints: &'static [ConstraintRef],
134    pub dependency_names: &'static [&'static str],
135
136    // Dialect-specific
137    pub dialect: TableDialect,
138}
139
140impl TableRef {
141    /// Creates a lightweight `TableRef` for SQL rendering only.
142    ///
143    /// Only `name` and `column_names` are populated; metadata fields use
144    /// empty defaults. Use a full struct literal for metadata-carrying refs.
145    #[must_use]
146    pub const fn sql(name: &'static str, column_names: &'static [&'static str]) -> Self {
147        Self {
148            name,
149            column_names,
150            schema: None,
151            qualified_name: "",
152            columns: &[],
153            primary_key: None,
154            foreign_keys: &[],
155            constraints: &[],
156            dependency_names: &[],
157            dialect: TableDialect::PostgreSQL {
158                is_unlogged: false,
159                is_temporary: false,
160                inherits: None,
161                tablespace: None,
162                is_rls_enabled: false,
163                comment: None,
164            },
165        }
166    }
167}
168
169/// Table fields needed by SQL rendering.
170#[derive(Clone, Copy, Debug, PartialEq, Eq)]
171pub struct TableSqlRef {
172    pub schema: Option<&'static str>,
173    pub name: &'static str,
174    pub column_names: &'static [&'static str],
175}
176
177impl TableSqlRef {
178    #[inline]
179    #[must_use]
180    pub const fn from_table_ref(table: TableRef) -> Self {
181        Self {
182            schema: table.schema,
183            name: table.name,
184            column_names: table.column_names,
185        }
186    }
187
188    #[inline]
189    #[must_use]
190    pub const fn from_table_ref_ref(table: &TableRef) -> Self {
191        Self {
192            schema: table.schema,
193            name: table.name,
194            column_names: table.column_names,
195        }
196    }
197}
198
199impl From<&TableRef> for TableSqlRef {
200    #[inline]
201    fn from(value: &TableRef) -> Self {
202        Self::from_table_ref_ref(value)
203    }
204}
205
206impl From<TableRef> for TableSqlRef {
207    #[inline]
208    fn from(value: TableRef) -> Self {
209        Self::from_table_ref(value)
210    }
211}
212
213/// Packed column metadata flags for [`ColumnRef`].
214///
215/// Encodes the nullability, primary-key, unique, and has-default bits in a
216/// single byte so that [`ColumnRef`] stays below the "too many bools" threshold
217/// while keeping each bit independently addressable.
218#[derive(Clone, Copy, Debug, PartialEq, Eq, Default)]
219pub struct ColumnFlags(u8);
220
221impl ColumnFlags {
222    /// Column is declared `NOT NULL`.
223    pub const NOT_NULL: Self = Self(1 << 0);
224    /// Column participates in the table's primary key.
225    pub const PRIMARY_KEY: Self = Self(1 << 1);
226    /// Column has a `UNIQUE` constraint.
227    pub const UNIQUE: Self = Self(1 << 2);
228    /// Column has a `DEFAULT` clause.
229    pub const HAS_DEFAULT: Self = Self(1 << 3);
230
231    /// Returns a flag set with no bits set.
232    #[must_use]
233    pub const fn empty() -> Self {
234        Self(0)
235    }
236
237    /// Reconstructs a flag set from its raw byte representation.
238    #[must_use]
239    pub const fn from_bits(bits: u8) -> Self {
240        Self(bits)
241    }
242
243    /// Returns the raw byte representation.
244    #[must_use]
245    pub const fn bits(self) -> u8 {
246        self.0
247    }
248
249    /// Returns `true` when every bit in `other` is set in `self`.
250    #[must_use]
251    pub const fn contains(self, other: Self) -> bool {
252        (self.0 & other.0) == other.0
253    }
254
255    /// Returns the union of two flag sets.
256    #[must_use]
257    pub const fn union(self, other: Self) -> Self {
258        Self(self.0 | other.0)
259    }
260}
261
262impl core::ops::BitOr for ColumnFlags {
263    type Output = Self;
264    fn bitor(self, rhs: Self) -> Self {
265        self.union(rhs)
266    }
267}
268
269impl core::ops::BitOrAssign for ColumnFlags {
270    fn bitor_assign(&mut self, rhs: Self) {
271        *self = self.union(rhs);
272    }
273}
274
275/// Column reference with full schema metadata.
276///
277/// Carries both the SQL rendering fields (`table`, `name`) and
278/// complete column metadata. SQL rendering code only uses the name fields
279/// and ignores extra metadata.
280#[derive(Clone, Copy, Debug, PartialEq, Eq)]
281pub struct ColumnRef {
282    // SQL rendering fields
283    pub table: &'static str,
284    pub name: &'static str,
285
286    // Schema metadata
287    pub sql_type: &'static str,
288    pub flags: ColumnFlags,
289
290    // Dialect-specific
291    pub dialect: ColumnDialect,
292}
293
294impl ColumnRef {
295    /// Creates a lightweight `ColumnRef` for SQL rendering only.
296    ///
297    /// Only `table` and `name` are populated; metadata fields
298    /// use empty defaults. Use a full struct literal for metadata-carrying refs.
299    #[must_use]
300    pub const fn sql(table: &'static str, name: &'static str) -> Self {
301        Self {
302            table,
303            name,
304            sql_type: "",
305            flags: ColumnFlags::empty(),
306            dialect: ColumnDialect::SQLite {
307                autoincrement: false,
308                default: None,
309                generated_expression: None,
310                generated_stored: false,
311                collate: None,
312            },
313        }
314    }
315
316    /// Returns `true` if this column is declared `NOT NULL`.
317    #[must_use]
318    pub const fn not_null(&self) -> bool {
319        self.flags.contains(ColumnFlags::NOT_NULL)
320    }
321
322    /// Returns `true` if this column participates in the primary key.
323    #[must_use]
324    pub const fn primary_key(&self) -> bool {
325        self.flags.contains(ColumnFlags::PRIMARY_KEY)
326    }
327
328    /// Returns `true` if this column has a `UNIQUE` constraint.
329    #[must_use]
330    pub const fn unique(&self) -> bool {
331        self.flags.contains(ColumnFlags::UNIQUE)
332    }
333
334    /// Returns `true` if this column has a `DEFAULT` clause.
335    #[must_use]
336    pub const fn has_default(&self) -> bool {
337        self.flags.contains(ColumnFlags::HAS_DEFAULT)
338    }
339}
340
341/// Column fields needed by SQL rendering.
342#[derive(Clone, Copy, Debug, PartialEq, Eq)]
343pub struct ColumnSqlRef {
344    pub table: &'static str,
345    pub name: &'static str,
346}
347
348impl ColumnSqlRef {
349    #[inline]
350    #[must_use]
351    pub const fn from_column_ref(column: ColumnRef) -> Self {
352        Self {
353            table: column.table,
354            name: column.name,
355        }
356    }
357
358    #[inline]
359    #[must_use]
360    pub const fn from_column_ref_ref(column: &ColumnRef) -> Self {
361        Self {
362            table: column.table,
363            name: column.name,
364        }
365    }
366}
367
368impl From<&ColumnRef> for ColumnSqlRef {
369    #[inline]
370    fn from(value: &ColumnRef) -> Self {
371        Self::from_column_ref_ref(value)
372    }
373}
374
375impl From<ColumnRef> for ColumnSqlRef {
376    #[inline]
377    fn from(value: ColumnRef) -> Self {
378        Self::from_column_ref(value)
379    }
380}
381
382// ==================== Identifier quoting ====================
383
384/// Writes a SQL identifier enclosed in double quotes.
385///
386/// Embedded double quotes are doubled. This function retains the original
387/// dialect-neutral interface used by downstream code. Core SQL rendering
388/// dispatches through the active value dialect internally.
389#[inline]
390pub fn write_quoted_ident(buf: &mut impl core::fmt::Write, name: &str) {
391    write_dialect_quoted_ident(Dialect::SQLite, buf, name);
392}
393
394/// Writes a SQL identifier using the delimiter required by `dialect`.
395///
396/// `PostgreSQL` and `SQLite` use `"..."`; `MySQL` uses backticks. Embedded
397/// delimiter characters are doubled so a runtime identifier cannot terminate
398/// the quoted identifier.
399///
400/// Identifiers without an embedded delimiter use a three-write fast path.
401#[inline]
402pub(crate) fn write_dialect_quoted_ident(
403    dialect: Dialect,
404    buf: &mut impl core::fmt::Write,
405    name: &str,
406) {
407    let delimiter = match dialect {
408        Dialect::MySQL => '`',
409        Dialect::SQLite | Dialect::PostgreSQL => '"',
410    };
411
412    let _ = buf.write_char(delimiter);
413    if name.contains(delimiter) {
414        for ch in name.chars() {
415            if ch == delimiter {
416                let _ = buf.write_char(delimiter);
417                let _ = buf.write_char(delimiter);
418            } else {
419                let _ = buf.write_char(ch);
420            }
421        }
422    } else {
423        let _ = buf.write_str(name);
424    }
425    let _ = buf.write_char(delimiter);
426}
427
428// ==================== SQLChunk ====================
429
430/// A SQL chunk represents a part of an SQL statement.
431///
432/// Each variant has a clear semantic purpose:
433/// - `Token` - SQL keywords and operators (SELECT, FROM, =, etc.)
434/// - `Ident` - Quoted identifiers ("`table_name`", "`column_name`")
435/// - `Raw` - Unquoted raw SQL text (function names, expressions)
436/// - `Param` - Parameter placeholders with values
437/// - `Table` - Table reference via `TableSqlRef`
438/// - `Column` - Column reference via `ColumnSqlRef`
439#[derive(Clone)]
440pub enum SQLChunk<'a, V: SQLParam> {
441    /// SQL keywords and operators: SELECT, FROM, WHERE, =, AND, etc.
442    /// Renders as: keyword with automatic spacing rules
443    Token(Token),
444
445    /// Quoted identifier for user-provided names
446    /// Renders as: "name" (with quotes)
447    /// Use for: table names, column names, alias names
448    Ident(Cow<'a, str>),
449
450    /// Raw SQL text (unquoted) for expressions, function names
451    /// Renders as: text (no quotes, as-is)
452    /// Use for: function names like COUNT, expressions, numeric literals
453    Raw(Cow<'a, str>),
454
455    /// Unsigned integer SQL literal rendered directly without heap allocation.
456    ///
457    /// Primarily used for clauses like LIMIT/OFFSET where numeric literals are
458    /// embedded directly in SQL text rather than parameterized.
459    Number(usize),
460
461    /// Parameter with value and placeholder
462    /// Renders as: ? or $1 or :name depending on placeholder style
463    Param(Param<'a, V>),
464
465    /// Table reference with static name and column names.
466    /// Renders as: "`table_name`"
467    /// Column names used for SELECT * expansion.
468    Table(TableSqlRef),
469
470    /// Column reference with static table and column names.
471    /// Renders as: "`table_name"."column_name`"
472    Column(ColumnSqlRef),
473}
474
475impl<'a, V: SQLParam> SQLChunk<'a, V> {
476    // ==================== const constructors ====================
477
478    /// Creates a token chunk - const
479    #[inline]
480    #[must_use]
481    pub const fn token(t: Token) -> Self {
482        Self::Token(t)
483    }
484
485    /// Creates a quoted identifier from a static string - const
486    #[inline]
487    #[must_use]
488    pub const fn ident_static(name: &'static str) -> Self {
489        Self::Ident(Cow::Borrowed(name))
490    }
491
492    /// Creates raw SQL text from a static string - const
493    #[inline]
494    #[must_use]
495    pub const fn raw_static(text: &'static str) -> Self {
496        Self::Raw(Cow::Borrowed(text))
497    }
498
499    /// Creates a table chunk - const
500    #[inline]
501    #[must_use]
502    pub const fn table(table: TableRef) -> Self {
503        Self::Table(TableSqlRef::from_table_ref(table))
504    }
505
506    /// Creates a column chunk - const
507    #[inline]
508    #[must_use]
509    pub const fn column(column: ColumnRef) -> Self {
510        Self::Column(ColumnSqlRef::from_column_ref(column))
511    }
512
513    /// Creates a parameter chunk with borrowed value - const
514    #[inline]
515    pub const fn param_borrowed(value: &'a V, placeholder: Placeholder) -> Self {
516        Self::Param(Param {
517            value: Some(Cow::Borrowed(value)),
518            placeholder,
519        })
520    }
521
522    // ==================== non-const constructors ====================
523
524    /// Creates a quoted identifier from a runtime string
525    #[inline]
526    pub fn ident(name: impl Into<Cow<'a, str>>) -> Self {
527        Self::Ident(name.into())
528    }
529
530    /// Creates raw SQL text from a runtime string
531    #[inline]
532    pub fn raw(text: impl Into<Cow<'a, str>>) -> Self {
533        Self::Raw(text.into())
534    }
535
536    /// Creates an unsigned integer SQL literal chunk.
537    #[inline]
538    #[must_use]
539    pub const fn number(value: usize) -> Self {
540        Self::Number(value)
541    }
542
543    /// Creates a parameter chunk with owned value
544    #[inline]
545    pub fn param(value: impl Into<Cow<'a, V>>, placeholder: Placeholder) -> Self {
546        Self::Param(Param {
547            value: Some(value.into()),
548            placeholder,
549        })
550    }
551
552    // ==================== write implementation ====================
553
554    /// Write chunk content to buffer
555    #[inline]
556    pub(crate) fn write(&self, buf: &mut impl core::fmt::Write) {
557        match self {
558            SQLChunk::Token(token) => {
559                let _ = buf.write_str(token.as_str());
560            }
561            SQLChunk::Ident(name) => {
562                write_dialect_quoted_ident(V::DIALECT, buf, name);
563            }
564            SQLChunk::Raw(text) => {
565                let _ = buf.write_str(text);
566            }
567            SQLChunk::Number(value) => {
568                let _ = write!(buf, "{value}");
569            }
570            SQLChunk::Param(Param { placeholder, .. }) => {
571                let _ = write!(buf, "{placeholder}");
572            }
573            SQLChunk::Table(t) => {
574                if let Some(schema) = t.schema {
575                    write_dialect_quoted_ident(V::DIALECT, buf, schema);
576                    let _ = buf.write_char('.');
577                }
578                write_dialect_quoted_ident(V::DIALECT, buf, t.name);
579            }
580            SQLChunk::Column(c) => {
581                write_dialect_quoted_ident(V::DIALECT, buf, c.table);
582                let _ = buf.write_char('.');
583                write_dialect_quoted_ident(V::DIALECT, buf, c.name);
584            }
585        }
586    }
587
588    /// Check if this chunk is "word-like" (needs space separation from other word-like chunks)
589    #[inline]
590    pub(crate) const fn is_word_like(&self) -> bool {
591        match self {
592            SQLChunk::Token(t) => !t.is_punctuation() && !t.is_operator(),
593            SQLChunk::Ident(_)
594            | SQLChunk::Raw(_)
595            | SQLChunk::Number(_)
596            | SQLChunk::Param(_)
597            | SQLChunk::Table(_)
598            | SQLChunk::Column(_) => true,
599        }
600    }
601}
602
603impl<V: SQLParam + core::fmt::Debug> core::fmt::Debug for SQLChunk<'_, V> {
604    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
605        match self {
606            SQLChunk::Token(token) => f.debug_tuple("Token").field(token).finish(),
607            SQLChunk::Ident(name) => f.debug_tuple("Ident").field(name).finish(),
608            SQLChunk::Raw(text) => f.debug_tuple("Raw").field(text).finish(),
609            SQLChunk::Number(value) => f.debug_tuple("Number").field(value).finish(),
610            SQLChunk::Param(param) => f.debug_tuple("Param").field(param).finish(),
611            SQLChunk::Table(t) => f
612                .debug_tuple("Table")
613                .field(&t.schema)
614                .field(&t.name)
615                .finish(),
616            SQLChunk::Column(c) => f
617                .debug_tuple("Column")
618                .field(&format!("{}.{}", c.table, c.name))
619                .finish(),
620        }
621    }
622}
623
624// ==================== From implementations ====================
625
626impl<V: SQLParam> From<Token> for SQLChunk<'_, V> {
627    #[inline]
628    fn from(value: Token) -> Self {
629        Self::Token(value)
630    }
631}
632
633impl<V: SQLParam> From<TableRef> for SQLChunk<'_, V> {
634    #[inline]
635    fn from(value: TableRef) -> Self {
636        Self::Table(value.into())
637    }
638}
639
640impl<V: SQLParam> From<ColumnRef> for SQLChunk<'_, V> {
641    #[inline]
642    fn from(value: ColumnRef) -> Self {
643        Self::Column(value.into())
644    }
645}
646
647impl<'a, V: SQLParam> From<Param<'a, V>> for SQLChunk<'a, V> {
648    #[inline]
649    fn from(value: Param<'a, V>) -> Self {
650        Self::Param(value)
651    }
652}
653
654#[cfg(test)]
655mod tests {
656    use super::*;
657    use crate::dialect::{Dialect, MySQLDialect, SQLiteDialect};
658    use core::mem::size_of;
659
660    #[allow(dead_code)]
661    #[derive(Clone, Debug)]
662    struct TestParam([usize; 4]);
663
664    impl SQLParam for TestParam {
665        const DIALECT: Dialect = Dialect::SQLite;
666        type DialectMarker = SQLiteDialect;
667    }
668
669    #[derive(Clone, Debug)]
670    struct MySQLTestParam;
671
672    impl SQLParam for MySQLTestParam {
673        const DIALECT: Dialect = Dialect::MySQL;
674        type DialectMarker = MySQLDialect;
675    }
676
677    #[test]
678    fn sql_chunk_stays_slim() {
679        // Param is the dominant variant for this 32-byte test parameter.
680        assert!(size_of::<SQLChunk<'static, TestParam>>() <= 64);
681    }
682
683    #[test]
684    fn quoted_ident_uses_the_dialect_delimiter() {
685        let mut sqlite = String::new();
686        write_dialect_quoted_ident(Dialect::SQLite, &mut sqlite, "account\"owner");
687        assert_eq!(sqlite, "\"account\"\"owner\"");
688
689        let mut postgres = String::new();
690        write_dialect_quoted_ident(Dialect::PostgreSQL, &mut postgres, "account\"owner");
691        assert_eq!(postgres, "\"account\"\"owner\"");
692
693        let mut mysql = String::new();
694        write_dialect_quoted_ident(Dialect::MySQL, &mut mysql, "account`owner");
695        assert_eq!(mysql, "`account``owner`");
696    }
697
698    #[test]
699    fn quoted_ident_keeps_injection_text_inside_the_identifier() {
700        let mut mysql = String::new();
701        write_dialect_quoted_ident(Dialect::MySQL, &mut mysql, "users`; DROP TABLE audit; --");
702        assert_eq!(mysql, "`users``; DROP TABLE audit; --`");
703    }
704
705    #[test]
706    fn table_chunk_preserves_structured_mysql_database_qualification() {
707        let table = TableRef {
708            schema: Some("tenant`db"),
709            ..TableRef::sql("user`accounts", &["id"])
710        };
711        let chunk = SQLChunk::<MySQLTestParam>::table(table);
712        let mut sql = String::new();
713
714        chunk.write(&mut sql);
715
716        assert_eq!(sql, "`tenant``db`.`user``accounts`");
717    }
718}