postrust_sql/
insert.rs

1//! INSERT statement builder.
2
3use crate::{
4    builder::SqlFragment,
5    identifier::{escape_ident, from_qi, QualifiedIdentifier},
6    param::SqlParam,
7};
8
9/// Builder for INSERT statements.
10#[derive(Clone, Debug, Default)]
11pub struct InsertBuilder {
12    table: Option<SqlFragment>,
13    columns: Vec<String>,
14    values: Vec<Vec<SqlFragment>>,
15    on_conflict: Option<OnConflict>,
16    returning: Vec<SqlFragment>,
17}
18
19#[derive(Clone, Debug)]
20pub enum OnConflict {
21    DoNothing,
22    DoUpdate {
23        columns: Vec<String>,
24        set: Vec<(String, SqlFragment)>,
25        where_clause: Option<SqlFragment>,
26    },
27}
28
29impl InsertBuilder {
30    /// Create a new INSERT builder.
31    pub fn new() -> Self {
32        Self::default()
33    }
34
35    /// Set the target table.
36    pub fn into_table(mut self, qi: &QualifiedIdentifier) -> Self {
37        self.table = Some(SqlFragment::raw(from_qi(qi)));
38        self
39    }
40
41    /// Set the columns to insert.
42    pub fn columns(mut self, cols: Vec<String>) -> Self {
43        self.columns = cols;
44        self
45    }
46
47    /// Add a row of values.
48    pub fn values(mut self, vals: Vec<SqlParam>) -> Self {
49        let row: Vec<SqlFragment> = vals
50            .into_iter()
51            .map(|v| {
52                let mut frag = SqlFragment::new();
53                frag.push_param(v);
54                frag
55            })
56            .collect();
57        self.values.push(row);
58        self
59    }
60
61    /// Add a row of raw SQL values.
62    pub fn values_raw(mut self, vals: Vec<SqlFragment>) -> Self {
63        self.values.push(vals);
64        self
65    }
66
67    /// Set ON CONFLICT DO NOTHING.
68    pub fn on_conflict_do_nothing(mut self) -> Self {
69        self.on_conflict = Some(OnConflict::DoNothing);
70        self
71    }
72
73    /// Set ON CONFLICT DO UPDATE.
74    pub fn on_conflict_do_update(
75        mut self,
76        conflict_columns: Vec<String>,
77        set: Vec<(String, SqlFragment)>,
78    ) -> Self {
79        self.on_conflict = Some(OnConflict::DoUpdate {
80            columns: conflict_columns,
81            set,
82            where_clause: None,
83        });
84        self
85    }
86
87    /// Add RETURNING clause.
88    pub fn returning(mut self, column: &str) -> Self {
89        self.returning.push(SqlFragment::raw(escape_ident(column)));
90        self
91    }
92
93    /// Add RETURNING * clause.
94    pub fn returning_all(mut self) -> Self {
95        self.returning.push(SqlFragment::raw("*"));
96        self
97    }
98
99    /// Build the INSERT statement.
100    pub fn build(self) -> SqlFragment {
101        let mut result = SqlFragment::new();
102
103        result.push("INSERT INTO ");
104
105        if let Some(table) = self.table {
106            result.append(table);
107        }
108
109        // Columns
110        if !self.columns.is_empty() {
111            result.push(" (");
112            for (i, col) in self.columns.iter().enumerate() {
113                if i > 0 {
114                    result.push(", ");
115                }
116                result.push(&escape_ident(col));
117            }
118            result.push(")");
119        }
120
121        // VALUES
122        if !self.values.is_empty() {
123            result.push(" VALUES ");
124            for (i, row) in self.values.into_iter().enumerate() {
125                if i > 0 {
126                    result.push(", ");
127                }
128                result.push("(");
129                for (j, val) in row.into_iter().enumerate() {
130                    if j > 0 {
131                        result.push(", ");
132                    }
133                    result.append(val);
134                }
135                result.push(")");
136            }
137        } else {
138            result.push(" DEFAULT VALUES");
139        }
140
141        // ON CONFLICT
142        if let Some(conflict) = self.on_conflict {
143            match conflict {
144                OnConflict::DoNothing => {
145                    result.push(" ON CONFLICT DO NOTHING");
146                }
147                OnConflict::DoUpdate {
148                    columns,
149                    set,
150                    where_clause,
151                } => {
152                    result.push(" ON CONFLICT (");
153                    for (i, col) in columns.iter().enumerate() {
154                        if i > 0 {
155                            result.push(", ");
156                        }
157                        result.push(&escape_ident(col));
158                    }
159                    result.push(") DO UPDATE SET ");
160                    for (i, (col, val)) in set.into_iter().enumerate() {
161                        if i > 0 {
162                            result.push(", ");
163                        }
164                        result.push(&escape_ident(&col));
165                        result.push(" = ");
166                        result.append(val);
167                    }
168                    if let Some(where_sql) = where_clause {
169                        result.push(" WHERE ");
170                        result.append(where_sql);
171                    }
172                }
173            }
174        }
175
176        // RETURNING
177        if !self.returning.is_empty() {
178            result.push(" RETURNING ");
179            for (i, ret) in self.returning.into_iter().enumerate() {
180                if i > 0 {
181                    result.push(", ");
182                }
183                result.append(ret);
184            }
185        }
186
187        result
188    }
189}
190
191#[cfg(test)]
192mod tests {
193    use super::*;
194
195    #[test]
196    fn test_simple_insert() {
197        let qi = QualifiedIdentifier::new("public", "users");
198        let sql = InsertBuilder::new()
199            .into_table(&qi)
200            .columns(vec!["name".into(), "email".into()])
201            .values(vec![
202                SqlParam::text("John"),
203                SqlParam::text("john@example.com"),
204            ])
205            .build();
206
207        assert!(sql.sql().contains("INSERT INTO"));
208        assert!(sql.sql().contains("VALUES"));
209        assert_eq!(sql.params().len(), 2);
210    }
211
212    #[test]
213    fn test_insert_returning() {
214        let qi = QualifiedIdentifier::unqualified("users");
215        let sql = InsertBuilder::new()
216            .into_table(&qi)
217            .columns(vec!["name".into()])
218            .values(vec![SqlParam::text("John")])
219            .returning("id")
220            .build();
221
222        assert!(sql.sql().contains("RETURNING"));
223    }
224
225    #[test]
226    fn test_insert_on_conflict_nothing() {
227        let qi = QualifiedIdentifier::unqualified("users");
228        let sql = InsertBuilder::new()
229            .into_table(&qi)
230            .columns(vec!["email".into()])
231            .values(vec![SqlParam::text("john@example.com")])
232            .on_conflict_do_nothing()
233            .build();
234
235        assert!(sql.sql().contains("ON CONFLICT DO NOTHING"));
236    }
237
238    #[test]
239    fn test_insert_upsert() {
240        let qi = QualifiedIdentifier::unqualified("users");
241        let mut name_val = SqlFragment::new();
242        name_val.push("EXCLUDED.\"name\"");
243
244        let sql = InsertBuilder::new()
245            .into_table(&qi)
246            .columns(vec!["id".into(), "name".into()])
247            .values(vec![SqlParam::Int(1), SqlParam::text("John")])
248            .on_conflict_do_update(vec!["id".into()], vec![("name".into(), name_val)])
249            .build();
250
251        assert!(sql.sql().contains("ON CONFLICT"));
252        assert!(sql.sql().contains("DO UPDATE SET"));
253    }
254}
postrust_sql/insert.rs

postrust_sql/
insert.rs
