halo/
insert.rs

1//! InsertBuilder：构建 INSERT 语句（对齐 go-sqlbuilder `insert.go` 的核心行为）。
2
3use crate::args::Args;
4use crate::flavor::Flavor;
5use crate::injection::{Injection, InjectionMarker};
6use crate::modifiers::{Arg, Builder, escape, escape_all};
7use crate::select::SelectBuilder;
8use crate::string_builder::StringBuilder;
9use std::cell::RefCell;
10use std::rc::Rc;
11
12const INSERT_MARKER_INIT: InjectionMarker = 0;
13const INSERT_MARKER_AFTER_INSERT_INTO: InjectionMarker = 1;
14const INSERT_MARKER_AFTER_COLS: InjectionMarker = 2;
15const INSERT_MARKER_AFTER_VALUES: InjectionMarker = 3;
16const INSERT_MARKER_AFTER_SELECT: InjectionMarker = 4;
17const INSERT_MARKER_AFTER_RETURNING: InjectionMarker = 5;
18
19#[derive(Debug, Clone)]
20pub struct InsertBuilder {
21    verb: &'static str,
22    table: Option<String>,
23    cols: Vec<String>,
24    values: Vec<Vec<String>>,
25    returning: Vec<String>,
26
27    args: Rc<RefCell<Args>>,
28
29    injection: Injection,
30    marker: InjectionMarker,
31
32    // Insert-Select holder
33    sb_holder: Option<String>,
34}
35
36impl Default for InsertBuilder {
37    fn default() -> Self {
38        Self::new()
39    }
40}
41
42impl InsertBuilder {
43    pub fn new() -> Self {
44        Self {
45            verb: "INSERT",
46            table: None,
47            cols: Vec::new(),
48            values: Vec::new(),
49            returning: Vec::new(),
50            args: Rc::new(RefCell::new(Args::default())),
51            injection: Injection::new(),
52            marker: INSERT_MARKER_INIT,
53            sb_holder: None,
54        }
55    }
56
57    pub fn set_flavor(&mut self, flavor: Flavor) -> Flavor {
58        let mut a = self.args.borrow_mut();
59        let old = a.flavor;
60        a.flavor = flavor;
61        old
62    }
63
64    pub fn flavor(&self) -> Flavor {
65        self.args.borrow().flavor
66    }
67
68    pub fn clone_builder(&self) -> Self {
69        let mut cloned = self.clone();
70
71        // 深拷贝 Args（避免 Rc 共享）
72        let args = Rc::new(RefCell::new(self.args.borrow().clone()));
73        cloned.args = args;
74
75        cloned
76    }
77
78    fn var(&self, v: impl Into<Arg>) -> String {
79        self.args.borrow_mut().add(v)
80    }
81
82    pub fn insert_into(&mut self, table: &str) -> &mut Self {
83        self.verb = "INSERT";
84        self.table = Some(escape(table));
85        self.marker = INSERT_MARKER_AFTER_INSERT_INTO;
86        self
87    }
88
89    pub fn insert_ignore_into(&mut self, table: &str) -> &mut Self {
90        let flavor = self.flavor();
91        self.verb = flavor.prepare_insert_ignore();
92        self.table = Some(escape(table));
93        self.marker = INSERT_MARKER_AFTER_INSERT_INTO;
94
95        // PostgreSQL: ON CONFLICT DO NOTHING 需要在 VALUES 后插入
96        if flavor == Flavor::PostgreSQL {
97            self.marker = INSERT_MARKER_AFTER_VALUES;
98            self.sql("ON CONFLICT DO NOTHING");
99        }
100        self
101    }
102
103    pub fn replace_into(&mut self, table: &str) -> &mut Self {
104        self.verb = "REPLACE";
105        self.table = Some(escape(table));
106        self.marker = INSERT_MARKER_AFTER_INSERT_INTO;
107        self
108    }
109
110    pub fn cols(&mut self, cols: impl IntoIterator<Item = impl Into<String>>) -> &mut Self {
111        self.cols = escape_all(cols.into_iter().map(Into::into));
112        self.marker = INSERT_MARKER_AFTER_COLS;
113        self
114    }
115
116    /// Insert-Select：返回一个 SelectBuilder 来构建 SELECT 部分。
117    pub fn select(&mut self, cols: impl IntoIterator<Item = impl Into<String>>) -> SelectBuilder {
118        let mut sb = SelectBuilder::new();
119        sb.select(cols);
120        sb.set_flavor(self.flavor());
121        let ph = self.var(Arg::Builder(Box::new(sb.clone_builder())));
122        self.sb_holder = Some(ph);
123        sb
124    }
125
126    pub fn values(&mut self, values: impl IntoIterator<Item = impl Into<Arg>>) -> &mut Self {
127        let placeholders: Vec<String> = values.into_iter().map(|v| self.var(v.into())).collect();
128        self.values.push(placeholders);
129        self.marker = INSERT_MARKER_AFTER_VALUES;
130        self
131    }
132
133    pub fn returning(&mut self, cols: impl IntoIterator<Item = impl Into<String>>) -> &mut Self {
134        self.returning = cols.into_iter().map(Into::into).collect();
135        self.marker = INSERT_MARKER_AFTER_RETURNING;
136        self
137    }
138
139    pub fn sql(&mut self, sql: impl Into<String>) -> &mut Self {
140        self.injection.sql(self.marker, sql);
141        self
142    }
143}
144
145impl Builder for InsertBuilder {
146    fn build_with_flavor(&self, flavor: Flavor, initial_arg: &[Arg]) -> (String, Vec<Arg>) {
147        let mut buf = StringBuilder::new();
148        write_injection(&mut buf, &self.injection, INSERT_MARKER_INIT);
149
150        // Oracle multi-values: INSERT ALL ... INTO ... VALUES (...) ... SELECT 1 from DUAL
151        if flavor == Flavor::Oracle && self.values.len() > 1 {
152            buf.write_leading(self.verb);
153            buf.write_str(" ALL");
154
155            for row in &self.values {
156                if let Some(t) = &self.table {
157                    buf.write_str(" INTO ");
158                    buf.write_str(t);
159                }
160                write_injection(&mut buf, &self.injection, INSERT_MARKER_AFTER_INSERT_INTO);
161
162                if !self.cols.is_empty() {
163                    buf.write_str(" (");
164                    buf.write_str(&self.cols.join(", "));
165                    buf.write_str(")");
166                    write_injection(&mut buf, &self.injection, INSERT_MARKER_AFTER_COLS);
167                }
168
169                buf.write_str(" VALUES (");
170                buf.write_str(&row.join(", "));
171                buf.write_str(")");
172            }
173
174            buf.write_str(" SELECT 1 from DUAL");
175            write_injection(&mut buf, &self.injection, INSERT_MARKER_AFTER_VALUES);
176            return self
177                .args
178                .borrow()
179                .compile_with_flavor(&buf.into_string(), flavor, initial_arg);
180        }
181
182        if let Some(t) = &self.table {
183            buf.write_leading(self.verb);
184            buf.write_str(" INTO ");
185            buf.write_str(t);
186        }
187        write_injection(&mut buf, &self.injection, INSERT_MARKER_AFTER_INSERT_INTO);
188
189        if !self.cols.is_empty() {
190            buf.write_str(" (");
191            buf.write_str(&self.cols.join(", "));
192            buf.write_str(")");
193            write_injection(&mut buf, &self.injection, INSERT_MARKER_AFTER_COLS);
194        }
195
196        if flavor == Flavor::SQLServer && !self.returning.is_empty() {
197            buf.write_str(" OUTPUT ");
198            let prefixed: Vec<String> = self
199                .returning
200                .iter()
201                .map(|c| format!("INSERTED.{c}"))
202                .collect();
203            buf.write_str(&prefixed.join(", "));
204            write_injection(&mut buf, &self.injection, INSERT_MARKER_AFTER_RETURNING);
205        }
206
207        if let Some(sb) = &self.sb_holder {
208            buf.write_str(" ");
209            buf.write_str(sb);
210            write_injection(&mut buf, &self.injection, INSERT_MARKER_AFTER_SELECT);
211        } else if !self.values.is_empty() {
212            buf.write_leading("VALUES");
213            buf.write_str(" ");
214            let rows: Vec<String> = self
215                .values
216                .iter()
217                .map(|r| format!("({})", r.join(", ")))
218                .collect();
219            buf.write_str(&rows.join(", "));
220        }
221
222        write_injection(&mut buf, &self.injection, INSERT_MARKER_AFTER_VALUES);
223
224        if (flavor == Flavor::PostgreSQL || flavor == Flavor::SQLite) && !self.returning.is_empty()
225        {
226            buf.write_leading("RETURNING");
227            buf.write_str(" ");
228            buf.write_str(&self.returning.join(", "));
229            write_injection(&mut buf, &self.injection, INSERT_MARKER_AFTER_RETURNING);
230        }
231
232        self.args
233            .borrow()
234            .compile_with_flavor(&buf.into_string(), flavor, initial_arg)
235    }
236
237    fn flavor(&self) -> Flavor {
238        self.flavor()
239    }
240}
241
242fn write_injection(buf: &mut StringBuilder, inj: &Injection, marker: InjectionMarker) {
243    let sqls = inj.at(marker);
244    if sqls.is_empty() {
245        return;
246    }
247    buf.write_leading("");
248    buf.write_str(&sqls.join(" "));
249}
halo/insert.rs

halo/
insert.rs
