1use std::fmt::Write;
13
14use crate::error::FlozError;
15use crate::value::Value;
16use super::quote_table;
17
18#[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#[derive(Debug, Clone)]
30enum ReturningClause {
31 None,
32 All, Cols(Vec<String>), }
35
36#[derive(Debug, Clone)]
38struct OnConflict {
39 column: String,
40 updates: Vec<(String, Value)>,
41}
42
43impl InsertQuery {
44 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 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 pub fn returning_all(mut self) -> Self {
64 self.returning = ReturningClause::All;
65 self
66 }
67
68 pub fn returning_col(mut self, column: &str) -> Self {
70 self.returning = ReturningClause::Cols(vec![column.to_string()]);
71 self
72 }
73
74 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 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 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 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 write!(sql, "INSERT INTO {} (", quote_table(&self.table)).unwrap();
107
108 for (i, col) in self.columns.iter().enumerate() {
110 if i > 0 { sql.push_str(", "); }
111 sql.push_str(col);
112 }
113
114 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 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 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#[derive(Debug, Clone)]
157pub struct InsertManyQuery {
158 table: String,
159 rows: Vec<Vec<(String, Value)>>,
161 returning: ReturningClause,
162}
163
164impl InsertManyQuery {
165 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 pub fn row(mut self, fields: Vec<(String, Value)>) -> Self {
176 self.rows.push(fields);
177 self
178 }
179
180 pub fn rows(mut self, rows: Vec<Vec<(String, Value)>>) -> Self {
182 self.rows = rows;
183 self
184 }
185
186 pub fn returning_all(mut self) -> Self {
188 self.returning = ReturningClause::All;
189 self
190 }
191
192 pub fn to_sql(&self) -> Result<(String, Vec<Value>), FlozError> {
194 if self.rows.is_empty() {
195 let sql = format!("INSERT INTO {} DEFAULT VALUES", quote_table(&self.table));
197 return Ok((sql, Vec::new()));
198 }
199
200 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 for (row_idx, row) in self.rows.iter().enumerate().skip(1) {
207 if row.len() != col_count {
208 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 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 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 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 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 #[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 ])
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}