Skip to main content

postrust_sql/
update.rs

1//! UPDATE statement builder.
2
3use crate::{
4    builder::SqlFragment,
5    expr::Expr,
6    identifier::{escape_ident, from_qi, QualifiedIdentifier},
7};
8
9/// Builder for UPDATE statements.
10#[derive(Clone, Debug, Default)]
11pub struct UpdateBuilder {
12    table: Option<SqlFragment>,
13    set: Vec<(String, SqlFragment)>,
14    where_clauses: Vec<SqlFragment>,
15    returning: Vec<SqlFragment>,
16}
17
18impl UpdateBuilder {
19    /// Create a new UPDATE builder.
20    pub fn new() -> Self {
21        Self::default()
22    }
23
24    /// Set the target table.
25    pub fn table(mut self, qi: &QualifiedIdentifier) -> Self {
26        self.table = Some(SqlFragment::raw(from_qi(qi)));
27        self
28    }
29
30    /// Set the target table with alias.
31    pub fn table_as(mut self, qi: &QualifiedIdentifier, alias: &str) -> Self {
32        self.table = Some(SqlFragment::raw(format!(
33            "{} AS {}",
34            from_qi(qi),
35            escape_ident(alias)
36        )));
37        self
38    }
39
40    /// Add a SET clause with parameterized value.
41    pub fn set<V: Into<crate::param::SqlParam>>(mut self, column: &str, value: V) -> Self {
42        let mut frag = SqlFragment::new();
43        frag.push_param(value);
44        self.set.push((column.to_string(), frag));
45        self
46    }
47
48    /// Add a SET clause with raw SQL.
49    pub fn set_raw(mut self, column: &str, value: SqlFragment) -> Self {
50        self.set.push((column.to_string(), value));
51        self
52    }
53
54    /// Add a WHERE clause.
55    pub fn where_expr(mut self, expr: Expr) -> Self {
56        self.where_clauses.push(expr.into_fragment());
57        self
58    }
59
60    /// Add a raw WHERE clause.
61    pub fn where_raw(mut self, sql: SqlFragment) -> Self {
62        self.where_clauses.push(sql);
63        self
64    }
65
66    /// Add RETURNING clause.
67    pub fn returning(mut self, column: &str) -> Self {
68        self.returning.push(SqlFragment::raw(escape_ident(column)));
69        self
70    }
71
72    /// Add RETURNING * clause.
73    pub fn returning_all(mut self) -> Self {
74        self.returning.push(SqlFragment::raw("*"));
75        self
76    }
77
78    /// Build the UPDATE statement.
79    pub fn build(self) -> SqlFragment {
80        let mut result = SqlFragment::new();
81
82        result.push("UPDATE ");
83
84        if let Some(table) = self.table {
85            result.append(table);
86        }
87
88        // SET
89        if !self.set.is_empty() {
90            result.push(" SET ");
91            for (i, (col, val)) in self.set.into_iter().enumerate() {
92                if i > 0 {
93                    result.push(", ");
94                }
95                result.push(&escape_ident(&col));
96                result.push(" = ");
97                result.append(val);
98            }
99        }
100
101        // WHERE
102        if !self.where_clauses.is_empty() {
103            result.push(" WHERE ");
104            for (i, clause) in self.where_clauses.into_iter().enumerate() {
105                if i > 0 {
106                    result.push(" AND ");
107                }
108                result.append(clause);
109            }
110        }
111
112        // RETURNING
113        if !self.returning.is_empty() {
114            result.push(" RETURNING ");
115            for (i, ret) in self.returning.into_iter().enumerate() {
116                if i > 0 {
117                    result.push(", ");
118                }
119                result.append(ret);
120            }
121        }
122
123        result
124    }
125}
126
127#[cfg(test)]
128mod tests {
129    use super::*;
130    use crate::param::SqlParam;
131
132    #[test]
133    fn test_simple_update() {
134        let qi = QualifiedIdentifier::new("public", "users");
135        let sql = UpdateBuilder::new()
136            .table(&qi)
137            .set("name", SqlParam::text("Jane"))
138            .where_expr(Expr::eq("id", 1i64))
139            .build();
140
141        assert!(sql.sql().contains("UPDATE"));
142        assert!(sql.sql().contains("SET"));
143        assert!(sql.sql().contains("WHERE"));
144        assert_eq!(sql.params().len(), 2);
145    }
146
147    #[test]
148    fn test_update_returning() {
149        let qi = QualifiedIdentifier::unqualified("users");
150        let sql = UpdateBuilder::new()
151            .table(&qi)
152            .set("status", SqlParam::text("active"))
153            .returning_all()
154            .build();
155
156        assert!(sql.sql().contains("RETURNING *"));
157    }
158
159    #[test]
160    fn test_update_multiple_sets() {
161        let qi = QualifiedIdentifier::unqualified("users");
162        let sql = UpdateBuilder::new()
163            .table(&qi)
164            .set("name", SqlParam::text("John"))
165            .set("email", SqlParam::text("john@new.com"))
166            .set("updated_at", SqlParam::text("now()"))
167            .where_expr(Expr::eq("id", 5i64))
168            .build();
169
170        assert_eq!(sql.params().len(), 4); // 3 sets + 1 where
171    }
172}