Skip to main content

a3s_orm/query/
insert.rs

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    /// Add a value to the first row. This preserves the original single-row API.
62    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}