halo/
cte_query.rs

1//! CTEQueryBuilder：构建一个 CTE 表（对齐 go-sqlbuilder `ctequery.go`）。
2
3use crate::args::Args;
4use crate::flavor::Flavor;
5use crate::injection::{Injection, InjectionMarker};
6use crate::modifiers::{Arg, Builder};
7use crate::string_builder::StringBuilder;
8use std::cell::RefCell;
9use std::rc::Rc;
10
11const CTE_QUERY_MARKER_INIT: InjectionMarker = 0;
12const CTE_QUERY_MARKER_AFTER_TABLE: InjectionMarker = 1;
13const CTE_QUERY_MARKER_AFTER_AS: InjectionMarker = 2;
14
15#[derive(Default)]
16pub struct CTEQueryBuilder {
17    name: Option<String>,
18    cols: Vec<String>,
19    builder_var: Option<String>,
20    #[allow(clippy::type_complexity)]
21    builder: Option<Box<dyn Builder>>,
22    auto_add_to_table_list: bool,
23
24    args: Rc<RefCell<Args>>,
25    injection: Injection,
26    marker: InjectionMarker,
27}
28
29// Box<dyn Builder> 无法自动 Debug；避免 Debug 派生失败。
30impl std::fmt::Debug for CTEQueryBuilder {
31    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
32        f.debug_struct("CTEQueryBuilder")
33            .field("name", &self.name)
34            .field("cols", &self.cols)
35            .field("builder_var", &self.builder_var)
36            .field("auto_add_to_table_list", &self.auto_add_to_table_list)
37            .finish()
38    }
39}
40
41impl Clone for CTEQueryBuilder {
42    fn clone(&self) -> Self {
43        self.clone_builder()
44    }
45}
46
47impl CTEQueryBuilder {
48    pub fn new() -> Self {
49        Self {
50            name: None,
51            cols: Vec::new(),
52            builder_var: None,
53            builder: None,
54            auto_add_to_table_list: false,
55            args: Rc::new(RefCell::new(Args::default())),
56            injection: Injection::new(),
57            marker: CTE_QUERY_MARKER_INIT,
58        }
59    }
60
61    pub fn set_flavor(&mut self, flavor: Flavor) -> Flavor {
62        let mut a = self.args.borrow_mut();
63        let old = a.flavor;
64        a.flavor = flavor;
65        old
66    }
67
68    pub fn flavor(&self) -> Flavor {
69        self.args.borrow().flavor
70    }
71
72    pub fn clone_builder(&self) -> Self {
73        let cloned = Self {
74            name: self.name.clone(),
75            cols: self.cols.clone(),
76            builder_var: self.builder_var.clone(),
77            builder: self
78                .builder
79                .as_ref()
80                .map(|b| dyn_clone::clone_box(b.as_ref())),
81            auto_add_to_table_list: self.auto_add_to_table_list,
82            args: Rc::new(RefCell::new(self.args.borrow().clone())),
83            injection: self.injection.clone(),
84            marker: self.marker,
85        };
86
87        if let (Some(ph), Some(b)) = (&self.builder_var, &self.builder) {
88            cloned
89                .args
90                .borrow_mut()
91                .replace(ph, Arg::Builder(dyn_clone::clone_box(b.as_ref())));
92        }
93
94        cloned
95    }
96
97    fn var(&self, v: impl Into<Arg>) -> String {
98        self.args.borrow_mut().add(v)
99    }
100
101    pub fn table(
102        &mut self,
103        name: impl Into<String>,
104        cols: impl IntoIterator<Item = impl Into<String>>,
105    ) -> &mut Self {
106        self.name = Some(name.into());
107        self.cols = cols.into_iter().map(Into::into).collect();
108        self.marker = CTE_QUERY_MARKER_AFTER_TABLE;
109        self
110    }
111
112    pub fn as_(&mut self, builder: impl Builder + 'static) -> &mut Self {
113        let b: Box<dyn Builder> = Box::new(builder);
114        let ph = self.var(Arg::Builder(dyn_clone::clone_box(b.as_ref())));
115        self.builder = Some(b);
116        self.builder_var = Some(ph);
117        self.marker = CTE_QUERY_MARKER_AFTER_AS;
118        self
119    }
120
121    pub fn add_to_table_list(&mut self) -> &mut Self {
122        self.auto_add_to_table_list = true;
123        self
124    }
125
126    pub fn should_add_to_table_list(&self) -> bool {
127        self.auto_add_to_table_list
128    }
129
130    pub fn table_name(&self) -> Option<&str> {
131        self.name.as_deref()
132    }
133
134    pub fn sql(&mut self, sql: impl Into<String>) -> &mut Self {
135        self.injection.sql(self.marker, sql);
136        self
137    }
138}
139
140impl Builder for CTEQueryBuilder {
141    fn build_with_flavor(&self, flavor: Flavor, initial_arg: &[Arg]) -> (String, Vec<Arg>) {
142        let mut buf = StringBuilder::new();
143        write_injection(&mut buf, &self.injection, CTE_QUERY_MARKER_INIT);
144
145        if let Some(name) = &self.name {
146            buf.write_leading(name);
147            if !self.cols.is_empty() {
148                buf.write_str(" (");
149                buf.write_str(&self.cols.join(", "));
150                buf.write_str(")");
151            }
152            write_injection(&mut buf, &self.injection, CTE_QUERY_MARKER_AFTER_TABLE);
153        }
154
155        if let Some(ph) = &self.builder_var {
156            buf.write_leading("AS (");
157            buf.write_str(ph);
158            buf.write_str(")");
159            write_injection(&mut buf, &self.injection, CTE_QUERY_MARKER_AFTER_AS);
160        }
161
162        self.args
163            .borrow()
164            .compile_with_flavor(&buf.into_string(), flavor, initial_arg)
165    }
166
167    fn flavor(&self) -> Flavor {
168        self.flavor()
169    }
170}
171
172fn write_injection(buf: &mut StringBuilder, inj: &Injection, marker: InjectionMarker) {
173    let sqls = inj.at(marker);
174    if sqls.is_empty() {
175        return;
176    }
177    buf.write_leading("");
178    buf.write_str(&sqls.join(" "));
179}
halo/cte_query.rs

halo/
cte_query.rs
