drizzle-core 0.3.1

use crate::SQLConstraintKind;
use crate::prelude::*;
use crate::{Dialect, Param, Placeholder, SQLParam, sql::tokens::Token};

// ==================== Dialect enums ====================

/// One value of an enum column type.
///
/// `label` is what a text column (or a `PostgreSQL` enum) stores, and
/// `discriminant` is what an integer column stores.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct EnumVariantRef {
    /// The stored text label, usually the variant name.
    pub label: &'static str,
    /// The integer value of the variant.
    pub discriminant: i64,
}

/// The values an enum column type can hold, in declaration order.
///
/// Implemented by `#[derive(SQLiteEnum)]` and `#[derive(PostgresEnum)]`.
/// The table macros copy these into the column's [`ColumnDialect`]
/// metadata (`enum_variants`), so tools such as `drizzle-seed` can generate
/// only valid values.
pub trait SQLEnumVariants {
    /// Every variant, in declaration order.
    const VARIANTS: &'static [EnumVariantRef];
}

/// Column metadata that applies to one dialect only.
///
/// Generated by the table macros and read by migrations. The fields mirror
/// the column attributes (`default`, `generated`, `collate`, ...).
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum ColumnDialect {
    SQLite {
        autoincrement: bool,
        default: Option<&'static str>,
        generated_expression: Option<&'static str>,
        generated_stored: bool,
        collate: Option<&'static str>,
        /// The values of an enum column, `None` for other columns.
        enum_variants: Option<&'static [EnumVariantRef]>,
    },
    PostgreSQL {
        postgres_type: &'static str,
        dimensions: Option<i32>,
        is_serial: bool,
        is_bigserial: bool,
        is_generated_identity: bool,
        is_identity_always: bool,
        default: Option<&'static str>,
        generated_expression: Option<&'static str>,
        generated_stored: bool,
        collate: Option<&'static str>,
        comment: Option<&'static str>,
        /// The values of an enum column, `None` for other columns.
        enum_variants: Option<&'static [EnumVariantRef]>,
    },
    MySQL {
        auto_increment: bool,
        default: Option<&'static str>,
        generated_expression: Option<&'static str>,
        generated_stored: bool,
        charset: Option<&'static str>,
        collate: Option<&'static str>,
        on_update: Option<&'static str>,
    },
}

/// Table metadata that applies to one dialect only.
///
/// Generated by the table macros and read by migrations.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum TableDialect {
    PostgreSQL {
        is_unlogged: bool,
        is_temporary: bool,
        inherits: Option<&'static str>,
        tablespace: Option<&'static str>,
        is_rls_enabled: bool,
        comment: Option<&'static str>,
    },
    SQLite {
        without_rowid: bool,
        strict: bool,
    },
    MySQL {
        is_temporary: bool,
        engine: Option<&'static str>,
        charset: Option<&'static str>,
        collate: Option<&'static str>,
        comment: Option<&'static str>,
    },
}

impl Default for TableDialect {
    fn default() -> Self {
        Self::PostgreSQL {
            is_unlogged: false,
            is_temporary: false,
            inherits: None,
            tablespace: None,
            is_rls_enabled: false,
            comment: None,
        }
    }
}

// ==================== Ref structs ====================

/// A foreign key constraint, as static metadata generated by the table
/// macros.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct ForeignKeyRef {
    pub name: &'static str,
    pub name_explicit: bool,
    pub target_table: &'static str,
    pub target_schema: &'static str,
    pub source_columns: &'static [&'static str],
    pub target_columns: &'static [&'static str],
    pub on_delete: Option<&'static str>,
    pub on_update: Option<&'static str>,
    pub deferrable: bool,
    pub initially_deferred: bool,
}

/// A table's primary key columns, as static metadata.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct PrimaryKeyRef {
    pub columns: &'static [&'static str],
}

/// A table constraint (`UNIQUE`, `CHECK`, ...), as static metadata.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct ConstraintRef {
    pub name: Option<&'static str>,
    pub name_explicit: bool,
    pub kind: SQLConstraintKind,
    pub columns: &'static [&'static str],
    pub check_expression: Option<&'static str>,
    pub deferrable: bool,
    pub initially_deferred: bool,
}

// ==================== Enhanced TableRef and ColumnRef ====================

/// A table's name, columns, and schema metadata, generated by the table
/// macros.
///
/// SQL rendering only reads `schema`, `name` and `column_names` (through
/// [`TableSqlRef`]). Migrations read the rest.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct TableRef {
    // SQL rendering fields
    pub name: &'static str,
    pub column_names: &'static [&'static str],

    // Schema metadata
    pub schema: Option<&'static str>,
    pub qualified_name: &'static str,
    pub columns: &'static [ColumnRef],
    pub primary_key: Option<PrimaryKeyRef>,
    pub foreign_keys: &'static [ForeignKeyRef],
    pub constraints: &'static [ConstraintRef],
    pub dependency_names: &'static [&'static str],

    // Dialect-specific
    pub dialect: TableDialect,
}

impl TableRef {
    /// Creates a `TableRef` with just a name and column names, for SQL
    /// rendering.
    ///
    /// Metadata fields are empty. Use a struct literal when metadata
    /// matters.
    #[must_use]
    pub const fn sql(name: &'static str, column_names: &'static [&'static str]) -> Self {
        Self {
            name,
            column_names,
            schema: None,
            qualified_name: "",
            columns: &[],
            primary_key: None,
            foreign_keys: &[],
            constraints: &[],
            dependency_names: &[],
            dialect: TableDialect::PostgreSQL {
                is_unlogged: false,
                is_temporary: false,
                inherits: None,
                tablespace: None,
                is_rls_enabled: false,
                comment: None,
            },
        }
    }
}

/// The parts of a [`TableRef`] that SQL rendering needs.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct TableSqlRef {
    pub schema: Option<&'static str>,
    pub name: &'static str,
    pub column_names: &'static [&'static str],
}

impl TableSqlRef {
    /// Copies the rendering fields out of a [`TableRef`].
    #[inline]
    #[must_use]
    pub const fn from_table_ref(table: TableRef) -> Self {
        Self {
            schema: table.schema,
            name: table.name,
            column_names: table.column_names,
        }
    }

    /// Copies the rendering fields out of a borrowed [`TableRef`].
    #[inline]
    #[must_use]
    pub const fn from_table_ref_ref(table: &TableRef) -> Self {
        Self {
            schema: table.schema,
            name: table.name,
            column_names: table.column_names,
        }
    }
}

impl From<&TableRef> for TableSqlRef {
    #[inline]
    fn from(value: &TableRef) -> Self {
        Self::from_table_ref_ref(value)
    }
}

impl From<TableRef> for TableSqlRef {
    #[inline]
    fn from(value: TableRef) -> Self {
        Self::from_table_ref(value)
    }
}

/// Column metadata flags for [`ColumnRef`], packed into one byte.
///
/// Holds the `NOT NULL`, primary-key, `UNIQUE` and has-default bits.
///
/// # Examples
///
/// ```
/// use drizzle_core::ColumnFlags;
///
/// let flags = ColumnFlags::NOT_NULL | ColumnFlags::PRIMARY_KEY;
/// assert!(flags.contains(ColumnFlags::NOT_NULL));
/// assert!(!flags.contains(ColumnFlags::UNIQUE));
/// ```
#[derive(Clone, Copy, Debug, PartialEq, Eq, Default)]
pub struct ColumnFlags(u8);

impl ColumnFlags {
    /// Column is declared `NOT NULL`.
    pub const NOT_NULL: Self = Self(1 << 0);
    /// Column participates in the table's primary key.
    pub const PRIMARY_KEY: Self = Self(1 << 1);
    /// Column has a `UNIQUE` constraint.
    pub const UNIQUE: Self = Self(1 << 2);
    /// Column has a `DEFAULT` clause.
    pub const HAS_DEFAULT: Self = Self(1 << 3);

    /// Returns a flag set with no bits set.
    #[must_use]
    pub const fn empty() -> Self {
        Self(0)
    }

    /// Reconstructs a flag set from its raw byte representation.
    #[must_use]
    pub const fn from_bits(bits: u8) -> Self {
        Self(bits)
    }

    /// Returns the raw byte representation.
    #[must_use]
    pub const fn bits(self) -> u8 {
        self.0
    }

    /// Returns `true` when every bit in `other` is set in `self`.
    #[must_use]
    pub const fn contains(self, other: Self) -> bool {
        (self.0 & other.0) == other.0
    }

    /// Returns the union of two flag sets.
    #[must_use]
    pub const fn union(self, other: Self) -> Self {
        Self(self.0 | other.0)
    }
}

impl core::ops::BitOr for ColumnFlags {
    type Output = Self;
    fn bitor(self, rhs: Self) -> Self {
        self.union(rhs)
    }
}

impl core::ops::BitOrAssign for ColumnFlags {
    fn bitor_assign(&mut self, rhs: Self) {
        *self = self.union(rhs);
    }
}

/// A column's table, name, and schema metadata, generated by the table
/// macros.
///
/// SQL rendering only reads `table` and `name` (through [`ColumnSqlRef`]).
/// Migrations read the rest.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct ColumnRef {
    // SQL rendering fields
    pub table: &'static str,
    pub name: &'static str,

    // Schema metadata
    pub sql_type: &'static str,
    pub flags: ColumnFlags,

    // Dialect-specific
    pub dialect: ColumnDialect,
}

impl ColumnRef {
    /// Creates a `ColumnRef` with just a table and column name, for SQL
    /// rendering.
    ///
    /// Metadata fields are empty. Use a struct literal when metadata
    /// matters.
    #[must_use]
    pub const fn sql(table: &'static str, name: &'static str) -> Self {
        Self {
            table,
            name,
            sql_type: "",
            flags: ColumnFlags::empty(),
            dialect: ColumnDialect::SQLite {
                autoincrement: false,
                default: None,
                generated_expression: None,
                generated_stored: false,
                collate: None,
                enum_variants: None,
            },
        }
    }

    /// Returns `true` if this column is declared `NOT NULL`.
    #[must_use]
    pub const fn not_null(&self) -> bool {
        self.flags.contains(ColumnFlags::NOT_NULL)
    }

    /// Returns `true` if this column participates in the primary key.
    #[must_use]
    pub const fn primary_key(&self) -> bool {
        self.flags.contains(ColumnFlags::PRIMARY_KEY)
    }

    /// Returns `true` if this column has a `UNIQUE` constraint.
    #[must_use]
    pub const fn unique(&self) -> bool {
        self.flags.contains(ColumnFlags::UNIQUE)
    }

    /// Returns `true` if this column has a `DEFAULT` clause.
    #[must_use]
    pub const fn has_default(&self) -> bool {
        self.flags.contains(ColumnFlags::HAS_DEFAULT)
    }
}

/// The parts of a [`ColumnRef`] that SQL rendering needs.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct ColumnSqlRef {
    pub table: &'static str,
    pub name: &'static str,
}

impl ColumnSqlRef {
    /// Copies the rendering fields out of a [`ColumnRef`].
    #[inline]
    #[must_use]
    pub const fn from_column_ref(column: ColumnRef) -> Self {
        Self {
            table: column.table,
            name: column.name,
        }
    }

    /// Copies the rendering fields out of a borrowed [`ColumnRef`].
    #[inline]
    #[must_use]
    pub const fn from_column_ref_ref(column: &ColumnRef) -> Self {
        Self {
            table: column.table,
            name: column.name,
        }
    }
}

impl From<&ColumnRef> for ColumnSqlRef {
    #[inline]
    fn from(value: &ColumnRef) -> Self {
        Self::from_column_ref_ref(value)
    }
}

impl From<ColumnRef> for ColumnSqlRef {
    #[inline]
    fn from(value: ColumnRef) -> Self {
        Self::from_column_ref(value)
    }
}

// ==================== Identifier quoting ====================

/// Writes `name` as a double-quoted SQL identifier, doubling any embedded
/// `"`.
///
/// Always uses double quotes, so it is wrong for MySQL. Rendering a [`SQL`](crate::SQL)
/// fragment picks the right quote for its dialect on its own.
///
/// # Examples
///
/// ```
/// let mut buf = String::new();
/// drizzle_core::sql::write_quoted_ident(&mut buf, r#"odd"name"#);
/// assert_eq!(buf, r#""odd""name""#);
/// ```
#[inline]
pub fn write_quoted_ident(buf: &mut impl core::fmt::Write, name: &str) {
    write_dialect_quoted_ident(Dialect::SQLite, buf, name);
}

/// Writes `name` as a quoted identifier for `dialect`.
///
/// PostgreSQL and SQLite use `"..."`; MySQL uses backticks. Embedded quote
/// characters are doubled so a runtime name cannot end the identifier
/// early.
#[inline]
pub(crate) fn write_dialect_quoted_ident(
    dialect: Dialect,
    buf: &mut impl core::fmt::Write,
    name: &str,
) {
    let delimiter = match dialect {
        Dialect::MySQL => '`',
        Dialect::SQLite | Dialect::PostgreSQL => '"',
    };

    let _ = buf.write_char(delimiter);
    if name.contains(delimiter) {
        for ch in name.chars() {
            if ch == delimiter {
                let _ = buf.write_char(delimiter);
                let _ = buf.write_char(delimiter);
            } else {
                let _ = buf.write_char(ch);
            }
        }
    } else {
        let _ = buf.write_str(name);
    }
    let _ = buf.write_char(delimiter);
}

// ==================== SQLChunk ====================

/// One piece of a [`SQL`](crate::SQL) fragment.
///
/// The renderer puts spaces between chunks where SQL needs them, so chunks
/// never carry their own surrounding whitespace.
#[derive(Clone)]
pub enum SQLChunk<'a, V: SQLParam> {
    /// A keyword, operator or punctuation: `SELECT`, `=`, `(`, ...
    Token(Token),

    /// A quoted identifier (table, column or alias name): `"name"`, or
    /// `` `name` `` for MySQL.
    Ident(Cow<'a, str>),

    /// Raw SQL text, written as-is with no quoting: function names,
    /// expressions.
    Raw(Cow<'a, str>),

    /// An unsigned integer literal written into the SQL text, for clauses
    /// such as `LIMIT` and `OFFSET` that are not parameterized.
    Number(usize),

    /// A placeholder and its bound value: `?`, `$1` or `:name`, depending on
    /// the dialect and the placeholder.
    Param(Param<'a, V>),

    /// A table, rendered as `"schema"."table"` or `"table"`. Its column names
    /// are used to expand an empty `SELECT` projection.
    Table(TableSqlRef),

    /// A column, rendered as `"table"."column"`.
    Column(ColumnSqlRef),
}

impl<'a, V: SQLParam> SQLChunk<'a, V> {
    // ==================== const constructors ====================

    /// Creates a token chunk.
    #[inline]
    #[must_use]
    pub const fn token(t: Token) -> Self {
        Self::Token(t)
    }

    /// Creates a quoted identifier from a static string.
    #[inline]
    #[must_use]
    pub const fn ident_static(name: &'static str) -> Self {
        Self::Ident(Cow::Borrowed(name))
    }

    /// Creates raw SQL text from a static string.
    #[inline]
    #[must_use]
    pub const fn raw_static(text: &'static str) -> Self {
        Self::Raw(Cow::Borrowed(text))
    }

    /// Creates a table chunk.
    #[inline]
    #[must_use]
    pub const fn table(table: TableRef) -> Self {
        Self::Table(TableSqlRef::from_table_ref(table))
    }

    /// Creates a column chunk.
    #[inline]
    #[must_use]
    pub const fn column(column: ColumnRef) -> Self {
        Self::Column(ColumnSqlRef::from_column_ref(column))
    }

    /// Creates a parameter chunk that borrows its value.
    #[inline]
    pub const fn param_borrowed(value: &'a V, placeholder: Placeholder) -> Self {
        Self::Param(Param {
            value: Some(Cow::Borrowed(value)),
            placeholder,
        })
    }

    // ==================== non-const constructors ====================

    /// Creates a quoted identifier.
    #[inline]
    pub fn ident(name: impl Into<Cow<'a, str>>) -> Self {
        Self::Ident(name.into())
    }

    /// Creates raw SQL text.
    #[inline]
    pub fn raw(text: impl Into<Cow<'a, str>>) -> Self {
        Self::Raw(text.into())
    }

    /// Creates an unsigned integer literal chunk.
    #[inline]
    #[must_use]
    pub const fn number(value: usize) -> Self {
        Self::Number(value)
    }

    /// Creates a parameter chunk with a value and placeholder.
    #[inline]
    pub fn param(value: impl Into<Cow<'a, V>>, placeholder: Placeholder) -> Self {
        Self::Param(Param {
            value: Some(value.into()),
            placeholder,
        })
    }

    // ==================== write implementation ====================

    /// Writes the chunk's text to `buf`, without surrounding spaces.
    #[inline]
    pub(crate) fn write(&self, buf: &mut impl core::fmt::Write) {
        match self {
            SQLChunk::Token(token) => {
                let _ = buf.write_str(token.as_str());
            }
            SQLChunk::Ident(name) => {
                write_dialect_quoted_ident(V::DIALECT, buf, name);
            }
            SQLChunk::Raw(text) => {
                let _ = buf.write_str(text);
            }
            SQLChunk::Number(value) => {
                let _ = write!(buf, "{value}");
            }
            SQLChunk::Param(Param { placeholder, .. }) => {
                let _ = write!(buf, "{placeholder}");
            }
            SQLChunk::Table(t) => {
                if let Some(schema) = t.schema {
                    write_dialect_quoted_ident(V::DIALECT, buf, schema);
                    let _ = buf.write_char('.');
                }
                write_dialect_quoted_ident(V::DIALECT, buf, t.name);
            }
            SQLChunk::Column(c) => {
                write_dialect_quoted_ident(V::DIALECT, buf, c.table);
                let _ = buf.write_char('.');
                write_dialect_quoted_ident(V::DIALECT, buf, c.name);
            }
        }
    }

    /// Whether the chunk is word-like, needing a space before another
    /// word-like chunk.
    #[inline]
    pub(crate) const fn is_word_like(&self) -> bool {
        match self {
            SQLChunk::Token(t) => !t.is_punctuation() && !t.is_operator(),
            SQLChunk::Ident(_)
            | SQLChunk::Raw(_)
            | SQLChunk::Number(_)
            | SQLChunk::Param(_)
            | SQLChunk::Table(_)
            | SQLChunk::Column(_) => true,
        }
    }
}

impl<V: SQLParam + core::fmt::Debug> core::fmt::Debug for SQLChunk<'_, V> {
    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
        match self {
            SQLChunk::Token(token) => f.debug_tuple("Token").field(token).finish(),
            SQLChunk::Ident(name) => f.debug_tuple("Ident").field(name).finish(),
            SQLChunk::Raw(text) => f.debug_tuple("Raw").field(text).finish(),
            SQLChunk::Number(value) => f.debug_tuple("Number").field(value).finish(),
            SQLChunk::Param(param) => f.debug_tuple("Param").field(param).finish(),
            SQLChunk::Table(t) => f
                .debug_tuple("Table")
                .field(&t.schema)
                .field(&t.name)
                .finish(),
            SQLChunk::Column(c) => f
                .debug_tuple("Column")
                .field(&format!("{}.{}", c.table, c.name))
                .finish(),
        }
    }
}

// ==================== From implementations ====================

impl<V: SQLParam> From<Token> for SQLChunk<'_, V> {
    #[inline]
    fn from(value: Token) -> Self {
        Self::Token(value)
    }
}

impl<V: SQLParam> From<TableRef> for SQLChunk<'_, V> {
    #[inline]
    fn from(value: TableRef) -> Self {
        Self::Table(value.into())
    }
}

impl<V: SQLParam> From<ColumnRef> for SQLChunk<'_, V> {
    #[inline]
    fn from(value: ColumnRef) -> Self {
        Self::Column(value.into())
    }
}

impl<'a, V: SQLParam> From<Param<'a, V>> for SQLChunk<'a, V> {
    #[inline]
    fn from(value: Param<'a, V>) -> Self {
        Self::Param(value)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::dialect::{Dialect, MySQLDialect, SQLiteDialect};
    use core::mem::size_of;

    #[allow(dead_code)]
    #[derive(Clone, Debug)]
    struct TestParam([usize; 4]);

    impl SQLParam for TestParam {
        const DIALECT: Dialect = Dialect::SQLite;
        type DialectMarker = SQLiteDialect;
    }

    #[derive(Clone, Debug)]
    struct MySQLTestParam;

    impl SQLParam for MySQLTestParam {
        const DIALECT: Dialect = Dialect::MySQL;
        type DialectMarker = MySQLDialect;
    }

    #[test]
    fn sql_chunk_stays_slim() {
        // Param is the dominant variant for this 32-byte test parameter.
        assert!(size_of::<SQLChunk<'static, TestParam>>() <= 64);
    }

    #[test]
    fn quoted_ident_uses_the_dialect_delimiter() {
        let mut sqlite = String::new();
        write_dialect_quoted_ident(Dialect::SQLite, &mut sqlite, "account\"owner");
        assert_eq!(sqlite, "\"account\"\"owner\"");

        let mut postgres = String::new();
        write_dialect_quoted_ident(Dialect::PostgreSQL, &mut postgres, "account\"owner");
        assert_eq!(postgres, "\"account\"\"owner\"");

        let mut mysql = String::new();
        write_dialect_quoted_ident(Dialect::MySQL, &mut mysql, "account`owner");
        assert_eq!(mysql, "`account``owner`");
    }

    #[test]
    fn quoted_ident_keeps_injection_text_inside_the_identifier() {
        let mut mysql = String::new();
        write_dialect_quoted_ident(Dialect::MySQL, &mut mysql, "users`; DROP TABLE audit; --");
        assert_eq!(mysql, "`users``; DROP TABLE audit; --`");
    }

    #[test]
    fn table_chunk_preserves_structured_mysql_database_qualification() {
        let table = TableRef {
            schema: Some("tenant`db"),
            ..TableRef::sql("user`accounts", &["id"])
        };
        let chunk = SQLChunk::<MySQLTestParam>::table(table);
        let mut sql = String::new();

        chunk.write(&mut sql);

        assert_eq!(sql, "`tenant``db`.`user``accounts`");
    }
}