floz_orm/query/

update.rs

//! UPDATE query builder.
//!
//! ```ignore
//! let (sql, params) = UpdateQuery::new("users")
//!     .set("age", Value::Short(31))
//!     .where_(col_id.eq(1))
//!     .to_sql()?;
//! // "UPDATE users SET age = $1 WHERE id = $2"
//! ```

use std::fmt::Write;

use crate::error::FlozError;
use crate::expr::Expr;
use crate::value::Value;
use super::quote_table;

/// An UPDATE query builder.
///
/// Safety: `to_sql()` returns `FlozError::UnguardedUpdate` if no WHERE clause is set,
/// to prevent accidental table-wide updates.
#[derive(Debug, Clone)]
pub struct UpdateQuery {
    table: String,
    sets: Vec<(String, Value)>,
    where_clause: Option<Expr>,
    returning: ReturningClause,
}

#[derive(Debug, Clone)]
enum ReturningClause {
    None,
    All,
    Cols(Vec<String>),
}

impl UpdateQuery {
    /// Create a new UPDATE query for the given table.
    pub fn new(table: &str) -> Self {
        Self {
            table: table.to_string(),
            sets: Vec::new(),
            where_clause: None,
            returning: ReturningClause::None,
        }
    }

    /// Set a column to a value. Chainable.
    pub fn set(mut self, column: &str, value: impl Into<Value>) -> Self {
        self.sets.push((column.to_string(), value.into()));
        self
    }

    /// Add a WHERE clause.
    pub fn where_(mut self, expr: Expr) -> Self {
        self.where_clause = Some(expr);
        self
    }

    /// Add RETURNING *.
    pub fn returning_all(mut self) -> Self {
        self.returning = ReturningClause::All;
        self
    }

    /// Add RETURNING for specific columns.
    pub fn returning_cols(mut self, columns: &[&str]) -> Self {
        self.returning = ReturningClause::Cols(
            columns.iter().map(|c| c.to_string()).collect(),
        );
        self
    }

    /// Render to SQL string + parameters.
    ///
    /// Returns `FlozError::UnguardedUpdate` if no WHERE clause is set.
    pub fn to_sql(&self) -> Result<(String, Vec<Value>), FlozError> {
        if self.where_clause.is_none() {
            return Err(FlozError::UnguardedUpdate);
        }

        let mut sql = String::with_capacity(128);
        let mut params = Vec::new();
        let mut idx = 0usize;

        // UPDATE table SET
        write!(sql, "UPDATE {} SET ", quote_table(&self.table)).unwrap();

        // SET col = $N, ...
        for (i, (col, val)) in self.sets.iter().enumerate() {
            if i > 0 { sql.push_str(", "); }
            idx += 1;
            write!(sql, "{} = ${}", col, idx).unwrap();
            params.push(val.clone());
        }

        // WHERE
        if let Some(ref expr) = self.where_clause {
            sql.push_str(" WHERE ");
            expr.to_sql(&mut sql, &mut params, &mut idx);
        }

        // RETURNING
        match &self.returning {
            ReturningClause::None => {}
            ReturningClause::All => sql.push_str(" RETURNING *"),
            ReturningClause::Cols(cols) => {
                sql.push_str(" RETURNING ");
                for (i, col) in cols.iter().enumerate() {
                    if i > 0 { sql.push_str(", "); }
                    sql.push_str(col);
                }
            }
        }

        Ok((sql, params))
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::column::Column;

    fn col_id() -> Column<i32> { Column::new("id", "users") }
    fn col_age() -> Column<i32> { Column::new("age", "users") }

    #[test]
    fn update_basic() {
        let (sql, params) = UpdateQuery::new("users")
            .set("age", 31i16)
            .where_(col_id().eq(1))
            .to_sql()
            .unwrap();
        assert_eq!(sql, "UPDATE users SET age = $1 WHERE id = $2");
        assert_eq!(params, vec![Value::Short(31), Value::Int(1)]);
    }

    #[test]
    fn update_multiple_sets() {
        let (sql, params) = UpdateQuery::new("users")
            .set("name", "Updated")
            .set("age", 31i16)
            .where_(col_id().eq(1))
            .to_sql()
            .unwrap();
        assert_eq!(sql, "UPDATE users SET name = $1, age = $2 WHERE id = $3");
        assert_eq!(params.len(), 3);
    }

    #[test]
    fn update_returning() {
        let (sql, _) = UpdateQuery::new("users")
            .set("age", 31i16)
            .where_(col_id().eq(1))
            .returning_all()
            .to_sql()
            .unwrap();
        assert!(sql.ends_with("RETURNING *"));
    }

    #[test]
    fn update_returning_cols() {
        let (sql, _) = UpdateQuery::new("users")
            .set("age", 31i16)
            .where_(col_id().eq(1))
            .returning_cols(&["id", "age"])
            .to_sql()
            .unwrap();
        assert!(sql.ends_with("RETURNING id, age"));
    }

    #[test]
    fn update_unguarded_error() {
        let result = UpdateQuery::new("users")
            .set("age", 31i16)
            .to_sql();
        assert!(matches!(result, Err(FlozError::UnguardedUpdate)));
    }

    #[test]
    fn update_complex_where() {
        let (sql, params) = UpdateQuery::new("users")
            .set("is_active", false)
            .where_(col_age().lt(18).or(col_age().gt(65)))
            .to_sql()
            .unwrap();
        assert_eq!(
            sql,
            "UPDATE users SET is_active = $1 WHERE (age < $2 OR age > $3)"
        );
        assert_eq!(params.len(), 3);
    }

    #[test]
    fn update_schema_qualified() {
        let (sql, _) = UpdateQuery::new("auth.users")
            .set("age", 31i16)
            .where_(col_id().eq(1))
            .to_sql()
            .unwrap();
        assert!(sql.starts_with("UPDATE \"auth\".\"users\""));
    }
}