1use std::marker::PhantomData;
2
3use crate::ast::{
4 Assignment, ConflictAction, ConflictAssignment, ConflictNode, ConflictValue, InsertNode,
5 QueryNode, TableNode,
6};
7use crate::expression::{Column, Selection};
8use crate::schema::Table;
9use crate::value::IntoSqlValue;
10
11use super::Query;
12
13#[derive(Clone, Debug)]
14pub struct InsertRow<T: Table> {
15 assignments: Vec<Assignment>,
16 marker: PhantomData<fn() -> T>,
17}
18
19impl<T: Table> Default for InsertRow<T> {
20 fn default() -> Self {
21 Self::new()
22 }
23}
24
25impl<T: Table> InsertRow<T> {
26 pub const fn new() -> Self {
27 Self {
28 assignments: Vec::new(),
29 marker: PhantomData,
30 }
31 }
32
33 pub fn value<V>(mut self, column: Column<T, V>, value: impl IntoSqlValue<V>) -> Self {
34 self.assignments.push(assignment(column, value));
35 self
36 }
37}
38
39#[derive(Clone, Debug)]
40pub struct InsertQuery<T: Table, O = ()> {
41 node: InsertNode,
42 marker: PhantomData<fn() -> (T, O)>,
43}
44
45pub fn insert_into<T: Table>() -> InsertQuery<T> {
46 InsertQuery {
47 node: InsertNode {
48 table: TableNode {
49 name: T::NAME,
50 alias: None,
51 },
52 rows: vec![Vec::new()],
53 conflict: None,
54 returning: Vec::new(),
55 },
56 marker: PhantomData,
57 }
58}
59
60impl<T: Table, O> InsertQuery<T, O> {
61 pub fn value<V>(mut self, column: Column<T, V>, value: impl IntoSqlValue<V>) -> Self {
63 if self.node.rows.is_empty() {
64 self.node.rows.push(Vec::new());
65 }
66 if let Some(row) = self.node.rows.first_mut() {
67 row.push(assignment(column, value));
68 }
69 self
70 }
71
72 pub fn row(mut self, row: InsertRow<T>) -> Self {
73 if self.node.rows.len() == 1 && self.node.rows.first().is_some_and(Vec::is_empty) {
74 if let Some(first) = self.node.rows.first_mut() {
75 *first = row.assignments;
76 }
77 } else {
78 self.node.rows.push(row.assignments);
79 }
80 self
81 }
82
83 pub fn rows(mut self, rows: impl IntoIterator<Item = InsertRow<T>>) -> Self {
84 for row in rows {
85 self = self.row(row);
86 }
87 self
88 }
89
90 pub fn on_conflict<C: ConflictTarget<T>>(mut self, target: C) -> Self {
91 self.node.conflict = Some(ConflictNode {
92 target: target.columns(),
93 action: None,
94 });
95 self
96 }
97
98 pub fn do_nothing(mut self) -> Self {
99 self.conflict_mut().action = Some(ConflictAction::DoNothing);
100 self
101 }
102
103 pub fn do_update<V>(mut self, column: Column<T, V>, value: impl IntoSqlValue<V>) -> Self {
104 let assignment = ConflictAssignment {
105 table: column.table_name(),
106 column: column.name(),
107 value: ConflictValue::Bound(value.into_sql_value()),
108 };
109 self.push_conflict_update(assignment);
110 self
111 }
112
113 pub fn do_update_from_excluded<V>(mut self, column: Column<T, V>) -> Self {
114 let assignment = ConflictAssignment {
115 table: column.table_name(),
116 column: column.name(),
117 value: ConflictValue::Excluded {
118 table: column.table_name(),
119 column: column.name(),
120 },
121 };
122 self.push_conflict_update(assignment);
123 self
124 }
125
126 pub fn returning<S: Selection>(mut self, selection: S) -> InsertQuery<T, S::Output> {
127 self.node.returning = selection.expressions();
128 InsertQuery {
129 node: self.node,
130 marker: PhantomData,
131 }
132 }
133
134 fn conflict_mut(&mut self) -> &mut ConflictNode {
135 self.node.conflict.get_or_insert_with(|| ConflictNode {
136 target: Vec::new(),
137 action: None,
138 })
139 }
140
141 fn push_conflict_update(&mut self, assignment: ConflictAssignment) {
142 let conflict = self.conflict_mut();
143 match &mut conflict.action {
144 Some(ConflictAction::DoUpdate(assignments)) => assignments.push(assignment),
145 _ => conflict.action = Some(ConflictAction::DoUpdate(vec![assignment])),
146 }
147 }
148}
149
150impl<T: Table, O> Query for InsertQuery<T, O> {
151 type Output = O;
152
153 fn compile(self, dialect: &impl crate::Dialect) -> crate::Result<crate::CompiledQuery> {
154 crate::compiler::compile(QueryNode::Insert(self.node), dialect)
155 }
156}
157
158pub trait ConflictTarget<T: Table> {
159 fn columns(self) -> Vec<&'static str>;
160}
161
162impl<T: Table, V> ConflictTarget<T> for Column<T, V> {
163 fn columns(self) -> Vec<&'static str> {
164 vec![self.name()]
165 }
166}
167
168macro_rules! conflict_target_tuple {
169 ($($column:ident),+ $(,)?) => {
170 impl<TableType: Table, $($column),+> ConflictTarget<TableType>
171 for ($(Column<TableType, $column>,)+)
172 {
173 #[allow(non_snake_case)]
174 fn columns(self) -> Vec<&'static str> {
175 let ($($column,)+) = self;
176 vec![$($column.name(),)+]
177 }
178 }
179 };
180}
181
182conflict_target_tuple!(A, B);
183conflict_target_tuple!(A, B, C);
184conflict_target_tuple!(A, B, C, D);
185
186fn assignment<T: Table, V>(column: Column<T, V>, value: impl IntoSqlValue<V>) -> Assignment {
187 Assignment {
188 table: column.table_name(),
189 column: column.name(),
190 value: value.into_sql_value(),
191 }
192}