halo_space/
insert.rs

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

halo_space/
insert.rs
