Skip to main content

floz_orm/query/
insert.rs

1//! INSERT query builder.
2//!
3//! ```ignore
4//! let (sql, params) = InsertQuery::new("users")
5//!     .set("name", Value::String("Alice".into()))
6//!     .set("age", Value::Short(30))
7//!     .returning_all()
8//!     .to_sql();
9//! // "INSERT INTO users (name, age) VALUES ($1, $2) RETURNING *"
10//! ```
11
12use std::fmt::Write;
13
14use crate::error::FlozError;
15use crate::value::Value;
16use super::quote_table;
17
18/// A single-row INSERT query builder.
19#[derive(Debug, Clone)]
20pub struct InsertQuery {
21    table: String,
22    columns: Vec<String>,
23    values: Vec<Value>,
24    returning: ReturningClause,
25    on_conflict: Option<OnConflict>,
26}
27
28/// What to return after INSERT.
29#[derive(Debug, Clone)]
30enum ReturningClause {
31    None,
32    All,               // RETURNING *
33    Cols(Vec<String>), // RETURNING col1, col2
34}
35
36/// ON CONFLICT configuration.
37#[derive(Debug, Clone)]
38struct OnConflict {
39    column: String,
40    updates: Vec<(String, Value)>,
41}
42
43impl InsertQuery {
44    /// Create a new INSERT query for the given table.
45    pub fn new(table: &str) -> Self {
46        Self {
47            table: table.to_string(),
48            columns: Vec::new(),
49            values: Vec::new(),
50            returning: ReturningClause::None,
51            on_conflict: None,
52        }
53    }
54
55    /// Set a column value. Chainable.
56    pub fn set(mut self, column: &str, value: impl Into<Value>) -> Self {
57        self.columns.push(column.to_string());
58        self.values.push(value.into());
59        self
60    }
61
62    /// Add RETURNING * to get the full inserted row back.
63    pub fn returning_all(mut self) -> Self {
64        self.returning = ReturningClause::All;
65        self
66    }
67
68    /// Add RETURNING for a single column.
69    pub fn returning_col(mut self, column: &str) -> Self {
70        self.returning = ReturningClause::Cols(vec![column.to_string()]);
71        self
72    }
73
74    /// Add RETURNING for multiple columns.
75    pub fn returning_cols(mut self, columns: &[&str]) -> Self {
76        self.returning = ReturningClause::Cols(
77            columns.iter().map(|c| c.to_string()).collect(),
78        );
79        self
80    }
81
82    /// Add ON CONFLICT (column) DO UPDATE SET ...
83    pub fn on_conflict(mut self, column: &str) -> Self {
84        self.on_conflict = Some(OnConflict {
85            column: column.to_string(),
86            updates: Vec::new(),
87        });
88        self
89    }
90
91    /// Add a column update for the ON CONFLICT clause. Must call `on_conflict()` first.
92    pub fn do_update(mut self, column: &str, value: impl Into<Value>) -> Self {
93        if let Some(ref mut oc) = self.on_conflict {
94            oc.updates.push((column.to_string(), value.into()));
95        }
96        self
97    }
98
99    /// Render to SQL string + parameters.
100    pub fn to_sql(&self) -> (String, Vec<Value>) {
101        let mut sql = String::with_capacity(128);
102        let mut params = Vec::new();
103        let mut idx = 0usize;
104
105        // INSERT INTO table
106        write!(sql, "INSERT INTO {} (", quote_table(&self.table)).unwrap();
107
108        // Column list
109        for (i, col) in self.columns.iter().enumerate() {
110            if i > 0 { sql.push_str(", "); }
111            sql.push_str(col);
112        }
113
114        // VALUES
115        sql.push_str(") VALUES (");
116        for (i, val) in self.values.iter().enumerate() {
117            if i > 0 { sql.push_str(", "); }
118            idx += 1;
119            write!(sql, "${}", idx).unwrap();
120            params.push(val.clone());
121        }
122        sql.push(')');
123
124        // ON CONFLICT
125        if let Some(ref oc) = self.on_conflict {
126            write!(sql, " ON CONFLICT ({}) DO UPDATE SET ", oc.column).unwrap();
127            for (i, (col, val)) in oc.updates.iter().enumerate() {
128                if i > 0 { sql.push_str(", "); }
129                idx += 1;
130                write!(sql, "{} = ${}", col, idx).unwrap();
131                params.push(val.clone());
132            }
133        }
134
135        // RETURNING
136        match &self.returning {
137            ReturningClause::None => {}
138            ReturningClause::All => sql.push_str(" RETURNING *"),
139            ReturningClause::Cols(cols) => {
140                sql.push_str(" RETURNING ");
141                for (i, col) in cols.iter().enumerate() {
142                    if i > 0 { sql.push_str(", "); }
143                    sql.push_str(col);
144                }
145            }
146        }
147
148        (sql, params)
149    }
150}
151
152/// A bulk INSERT query builder.
153///
154/// All rows must have the exact same columns. If any row is missing a column
155/// that the first row defined, `to_sql()` returns `FlozError::MismatchedBulkInsertColumns`.
156#[derive(Debug, Clone)]
157pub struct InsertManyQuery {
158    table: String,
159    /// Each row is a list of (column_name, value) pairs.
160    rows: Vec<Vec<(String, Value)>>,
161    returning: ReturningClause,
162}
163
164impl InsertManyQuery {
165    /// Create a new bulk INSERT query.
166    pub fn new(table: &str) -> Self {
167        Self {
168            table: table.to_string(),
169            rows: Vec::new(),
170            returning: ReturningClause::None,
171        }
172    }
173
174    /// Add a row. Each row is a slice of (column_name, value) pairs.
175    pub fn row(mut self, fields: Vec<(String, Value)>) -> Self {
176        self.rows.push(fields);
177        self
178    }
179
180    /// Add multiple rows at once.
181    pub fn rows(mut self, rows: Vec<Vec<(String, Value)>>) -> Self {
182        self.rows = rows;
183        self
184    }
185
186    /// Add RETURNING *.
187    pub fn returning_all(mut self) -> Self {
188        self.returning = ReturningClause::All;
189        self
190    }
191
192    /// Render to SQL string + parameters, or error if rows have mismatched columns.
193    pub fn to_sql(&self) -> Result<(String, Vec<Value>), FlozError> {
194        if self.rows.is_empty() {
195            // Empty insert — just produce a valid but no-op statement
196            let sql = format!("INSERT INTO {} DEFAULT VALUES", quote_table(&self.table));
197            return Ok((sql, Vec::new()));
198        }
199
200        // Get column order from first row
201        let first_row = &self.rows[0];
202        let columns: Vec<&str> = first_row.iter().map(|(c, _)| c.as_str()).collect();
203        let col_count = columns.len();
204
205        // Validate all rows have the same columns
206        for (row_idx, row) in self.rows.iter().enumerate().skip(1) {
207            if row.len() != col_count {
208                // Find which column is missing
209                let row_cols: Vec<&str> = row.iter().map(|(c, _)| c.as_str()).collect();
210                for expected in &columns {
211                    if !row_cols.contains(expected) {
212                        return Err(FlozError::MismatchedBulkInsertColumns {
213                            row: row_idx,
214                            column: expected.to_string(),
215                        });
216                    }
217                }
218                // If we get here, row has extra columns
219                return Err(FlozError::MismatchedBulkInsertColumns {
220                    row: row_idx,
221                    column: format!("row has {} columns, expected {}", row.len(), col_count),
222                });
223            }
224        }
225
226        let mut sql = String::with_capacity(128);
227        let mut params = Vec::new();
228        let mut idx = 0usize;
229
230        // INSERT INTO table (col1, col2, ...)
231        write!(sql, "INSERT INTO {} (", quote_table(&self.table)).unwrap();
232        for (i, col) in columns.iter().enumerate() {
233            if i > 0 { sql.push_str(", "); }
234            sql.push_str(col);
235        }
236        sql.push_str(") VALUES ");
237
238        // VALUES ($1, $2), ($3, $4), ...
239        for (row_idx, row) in self.rows.iter().enumerate() {
240            if row_idx > 0 { sql.push_str(", "); }
241            sql.push('(');
242            for (i, (_, val)) in row.iter().enumerate() {
243                if i > 0 { sql.push_str(", "); }
244                idx += 1;
245                write!(sql, "${}", idx).unwrap();
246                params.push(val.clone());
247            }
248            sql.push(')');
249        }
250
251        // RETURNING
252        match &self.returning {
253            ReturningClause::None => {}
254            ReturningClause::All => sql.push_str(" RETURNING *"),
255            ReturningClause::Cols(cols) => {
256                sql.push_str(" RETURNING ");
257                for (i, col) in cols.iter().enumerate() {
258                    if i > 0 { sql.push_str(", "); }
259                    sql.push_str(col);
260                }
261            }
262        }
263
264        Ok((sql, params))
265    }
266}
267
268#[cfg(test)]
269mod tests {
270    use super::*;
271
272    #[test]
273    fn insert_single_column() {
274        let (sql, params) = InsertQuery::new("users")
275            .set("name", "Alice")
276            .to_sql();
277        assert_eq!(sql, "INSERT INTO users (name) VALUES ($1)");
278        assert_eq!(params, vec![Value::String("Alice".into())]);
279    }
280
281    #[test]
282    fn insert_multiple_columns() {
283        let (sql, params) = InsertQuery::new("users")
284            .set("name", "Alice")
285            .set("age", 30i16)
286            .to_sql();
287        assert_eq!(sql, "INSERT INTO users (name, age) VALUES ($1, $2)");
288        assert_eq!(params, vec![
289            Value::String("Alice".into()),
290            Value::Short(30),
291        ]);
292    }
293
294    #[test]
295    fn insert_returning_all() {
296        let (sql, _) = InsertQuery::new("users")
297            .set("name", "Alice")
298            .returning_all()
299            .to_sql();
300        assert!(sql.ends_with("RETURNING *"));
301    }
302
303    #[test]
304    fn insert_returning_col() {
305        let (sql, _) = InsertQuery::new("users")
306            .set("name", "Alice")
307            .returning_col("id")
308            .to_sql();
309        assert!(sql.ends_with("RETURNING id"));
310    }
311
312    #[test]
313    fn insert_returning_cols() {
314        let (sql, _) = InsertQuery::new("users")
315            .set("name", "Alice")
316            .returning_cols(&["id", "created_at"])
317            .to_sql();
318        assert!(sql.ends_with("RETURNING id, created_at"));
319    }
320
321    #[test]
322    fn insert_on_conflict() {
323        let (sql, params) = InsertQuery::new("users")
324            .set("name", "Alice")
325            .set("age", 30i16)
326            .on_conflict("name")
327            .do_update("age", 31i16)
328            .to_sql();
329        assert_eq!(
330            sql,
331            "INSERT INTO users (name, age) VALUES ($1, $2) ON CONFLICT (name) DO UPDATE SET age = $3"
332        );
333        assert_eq!(params.len(), 3);
334    }
335
336    #[test]
337    fn insert_on_conflict_multiple_updates() {
338        let (sql, params) = InsertQuery::new("users")
339            .set("name", "Alice")
340            .set("age", 30i16)
341            .on_conflict("name")
342            .do_update("age", 31i16)
343            .do_update("name", "Alice Updated")
344            .to_sql();
345        assert_eq!(
346            sql,
347            "INSERT INTO users (name, age) VALUES ($1, $2) ON CONFLICT (name) DO UPDATE SET age = $3, name = $4"
348        );
349        assert_eq!(params.len(), 4);
350    }
351
352    #[test]
353    fn insert_schema_qualified() {
354        let (sql, _) = InsertQuery::new("auth.users")
355            .set("name", "Alice")
356            .to_sql();
357        assert!(sql.starts_with("INSERT INTO \"auth\".\"users\""));
358    }
359
360    // ── InsertMany ──
361
362    #[test]
363    fn insert_many_basic() {
364        let (sql, params) = InsertManyQuery::new("users")
365            .row(vec![
366                ("name".into(), Value::String("Alice".into())),
367                ("age".into(), Value::Short(30)),
368            ])
369            .row(vec![
370                ("name".into(), Value::String("Bob".into())),
371                ("age".into(), Value::Short(25)),
372            ])
373            .to_sql()
374            .unwrap();
375        assert_eq!(
376            sql,
377            "INSERT INTO users (name, age) VALUES ($1, $2), ($3, $4)"
378        );
379        assert_eq!(params.len(), 4);
380    }
381
382    #[test]
383    fn insert_many_single_row() {
384        let (sql, params) = InsertManyQuery::new("users")
385            .row(vec![
386                ("name".into(), Value::String("Alice".into())),
387            ])
388            .to_sql()
389            .unwrap();
390        assert_eq!(sql, "INSERT INTO users (name) VALUES ($1)");
391        assert_eq!(params.len(), 1);
392    }
393
394    #[test]
395    fn insert_many_returning() {
396        let (sql, _) = InsertManyQuery::new("users")
397            .row(vec![("name".into(), Value::String("A".into()))])
398            .row(vec![("name".into(), Value::String("B".into()))])
399            .returning_all()
400            .to_sql()
401            .unwrap();
402        assert!(sql.ends_with("RETURNING *"));
403    }
404
405    #[test]
406    fn insert_many_ragged_rows_error() {
407        let result = InsertManyQuery::new("users")
408            .row(vec![
409                ("name".into(), Value::String("Alice".into())),
410                ("age".into(), Value::Short(30)),
411            ])
412            .row(vec![
413                ("name".into(), Value::String("Bob".into())),
414                // Missing "age"!
415            ])
416            .to_sql();
417        assert!(result.is_err());
418        match result {
419            Err(FlozError::MismatchedBulkInsertColumns { row, column }) => {
420                assert_eq!(row, 1);
421                assert_eq!(column, "age");
422            }
423            _ => panic!("Expected MismatchedBulkInsertColumns"),
424        }
425    }
426
427    #[test]
428    fn insert_many_empty() {
429        let (sql, params) = InsertManyQuery::new("users")
430            .to_sql()
431            .unwrap();
432        assert_eq!(sql, "INSERT INTO users DEFAULT VALUES");
433        assert!(params.is_empty());
434    }
435
436    #[test]
437    fn insert_many_three_rows() {
438        let (sql, params) = InsertManyQuery::new("users")
439            .row(vec![("x".into(), Value::Int(1))])
440            .row(vec![("x".into(), Value::Int(2))])
441            .row(vec![("x".into(), Value::Int(3))])
442            .to_sql()
443            .unwrap();
444        assert_eq!(sql, "INSERT INTO users (x) VALUES ($1), ($2), ($3)");
445        assert_eq!(params.len(), 3);
446    }
447}