gluesql_core/ast_builder/
update.rs

1use {
2    super::{AssignmentNode, Build, ExprNode},
3    crate::{
4        ast::{Assignment, Expr, Statement},
5        result::Result,
6    },
7};
8
9#[derive(Clone, Debug)]
10pub struct UpdateNode {
11    table_name: String,
12}
13
14impl UpdateNode {
15    pub fn new(table_name: String) -> Self {
16        Self { table_name }
17    }
18
19    pub fn filter<'a, T: Into<ExprNode<'a>>>(self, expr: T) -> UpdateFilterNode<'a> {
20        UpdateFilterNode::new(self.table_name, expr)
21    }
22
23    pub fn set<'a, T: Into<ExprNode<'a>>>(self, id: &str, value: T) -> UpdateSetNode<'a> {
24        UpdateSetNode::new(self.table_name, None, id, value)
25    }
26}
27
28#[derive(Clone, Debug)]
29pub struct UpdateFilterNode<'a> {
30    table_name: String,
31    selection: ExprNode<'a>,
32}
33
34impl<'a> UpdateFilterNode<'a> {
35    pub fn new<T: Into<ExprNode<'a>>>(table_name: String, expr: T) -> Self {
36        Self {
37            table_name,
38            selection: expr.into(),
39        }
40    }
41
42    pub fn filter<T: Into<ExprNode<'a>>>(mut self, expr: T) -> Self {
43        self.selection = self.selection.and(expr.into());
44        self
45    }
46
47    pub fn set<T: Into<ExprNode<'a>>>(self, id: &str, value: T) -> UpdateSetNode<'a> {
48        UpdateSetNode::new(self.table_name, Some(self.selection), id, value)
49    }
50}
51
52#[derive(Clone, Debug)]
53pub struct UpdateSetNode<'a> {
54    table_name: String,
55    selection: Option<ExprNode<'a>>,
56    assignments: Vec<AssignmentNode<'a>>,
57}
58
59impl<'a> UpdateSetNode<'a> {
60    pub fn new<T: Into<ExprNode<'a>>>(
61        table_name: String,
62        selection: Option<ExprNode<'a>>,
63        id: &str,
64        value: T,
65    ) -> Self {
66        let assignments = vec![AssignmentNode::Expr(id.to_owned(), value.into())];
67
68        Self {
69            table_name,
70            selection,
71            assignments,
72        }
73    }
74
75    pub fn set<T: Into<ExprNode<'a>>>(mut self, id: &str, value: T) -> Self {
76        self.assignments
77            .push(AssignmentNode::Expr(id.to_owned(), value.into()));
78        self
79    }
80}
81
82impl<'a> Build for UpdateSetNode<'a> {
83    fn build(self) -> Result<Statement> {
84        let table_name = self.table_name;
85        let selection = self.selection.map(Expr::try_from).transpose()?;
86        let assignments = self
87            .assignments
88            .into_iter()
89            .map(Assignment::try_from)
90            .collect::<Result<Vec<_>>>()?;
91        Ok(Statement::Update {
92            table_name,
93            assignments,
94            selection,
95        })
96    }
97}
98
99#[cfg(test)]
100mod tests {
101    use crate::ast_builder::{Build, col, num, table, test, text};
102
103    #[test]
104    fn update() {
105        let actual = table("Foo").update().set("id", "2").build();
106        let expected = "UPDATE Foo SET id = 2";
107        test(actual, expected);
108
109        let actual = table("Foo")
110            .update()
111            .set("id", "2")
112            .set("name", "Bar")
113            .build();
114        let expected = "UPDATE Foo SET id = 2, name=Bar";
115        test(actual, expected);
116
117        let actual = table("Foo")
118            .update()
119            .filter("Bar = 1")
120            .set("id", "2")
121            .set("name", "americano")
122            .build();
123        let expected = "UPDATE Foo SET id = 2, name = americano WHERE Bar = 1";
124        test(actual, expected);
125
126        let actual = table("Foo")
127            .update()
128            .filter(col("id").gt(num(1)))
129            .filter("name = 'americano'")
130            .set("name", text("espresso"))
131            .build();
132        let expected = "
133            UPDATE Foo
134            SET name = 'espresso'
135            WHERE id > 1 AND name = 'americano'";
136        test(actual, expected);
137
138        let actual = table("Foo")
139            .update()
140            .filter("body_item = 1")
141            .set("id", "2")
142            .set(
143                "head_item",
144                "(SELECT id FROM head_item WHERE level = 3 LIMIT 1)",
145            )
146            .build();
147        let expected = "UPDATE Foo SET id = 2, head_item = (SELECT id FROM head_item WHERE level = 3 LIMIT 1) WHERE body_item = 1";
148        test(actual, expected);
149    }
150}