Skip to main content

somnia_core/
query.rs

1use crate::{
2    expr::{Column, DynExpr, Order, Projection, RecordLink, SurrealQL},
3    types::{SurrealEdge, SurrealRecord, Thing},
4};
5
6/// How a mutating statement should return its affected rows.
7#[derive(Debug, Clone, Copy, PartialEq, Eq)]
8pub enum Returning {
9    /// no RETURN clause
10    None,
11    /// `RETURN NONE`
12    Nothing,
13    /// `RETURN BEFORE`
14    Before,
15    /// `RETURN AFTER`
16    After,
17    /// `RETURN DIFF`
18    Diff,
19}
20
21impl Returning {
22    fn render(self, buf: &mut String) {
23        match self {
24            Returning::None => {}
25            Returning::Nothing => buf.push_str(" RETURN NONE"),
26            Returning::Before => buf.push_str(" RETURN BEFORE"),
27            Returning::After => buf.push_str(" RETURN AFTER"),
28            Returning::Diff => buf.push_str(" RETURN DIFF"),
29        }
30    }
31}
32
33/// A statement target: either a whole table (`asset`) or a single record link
34/// (`type::record('asset', '<id>')`).
35enum Target {
36    Table(&'static str),
37    Record(RecordLink),
38}
39
40impl Target {
41    fn render(&self, buf: &mut String) {
42        match self {
43            Target::Table(t) => buf.push_str(t),
44            Target::Record(r) => r.render_dyn(buf),
45        }
46    }
47}
48
49// ═══════════════════════════════════════════════════════════════════════════════
50// Table
51// ═══════════════════════════════════════════════════════════════════════════════
52
53pub struct Table<T: SurrealRecord> {
54    _marker: std::marker::PhantomData<T>,
55}
56
57impl<T: SurrealRecord> Table<T> {
58    pub fn new() -> Self { Self { _marker: std::marker::PhantomData } }
59
60    pub fn select(self, _cols: crate::expr::ColumnSet<T>) -> Select<T> {
61        Select::bare()
62    }
63
64    /// Select an explicit projection list (`SELECT <fields…> FROM table`).
65    pub fn project(self, fields: Vec<Projection>) -> Select<T> {
66        let mut s = Select::bare();
67        s.projections = fields;
68        s
69    }
70
71    /// `SELECT count() FROM table GROUP ALL`
72    pub fn count(self, _field: &str) -> Select<T> {
73        let mut s = Select::bare();
74        s.count = true;
75        s.group_all = true;
76        s
77    }
78
79    pub fn insert(self) -> Insert<T> { Insert { data: Vec::new(), return_fields: vec![] } }
80    pub fn create(self) -> Create<T> { Create::for_table() }
81    pub fn update(self) -> Update<T> { Update::for_table() }
82    pub fn delete(self) -> Delete<T> { Delete::for_table() }
83}
84
85impl<T: SurrealRecord> Default for Table<T> {
86    fn default() -> Self { Self::new() }
87}
88
89// ═══════════════════════════════════════════════════════════════════════════════
90// SELECT
91// ═══════════════════════════════════════════════════════════════════════════════
92
93pub struct Select<T: SurrealRecord> {
94    _marker: std::marker::PhantomData<T>,
95    projections: Vec<Projection>,
96    filter: Option<Box<dyn DynExpr>>,
97    order: Vec<(String, Order)>,
98    limit: Option<u32>,
99    start: u32,
100    fetch: Vec<String>,
101    group_by: Vec<String>,
102    group_all: bool,
103    count: bool,
104    count_alias: Option<&'static str>,
105}
106
107impl<T: SurrealRecord> Select<T> {
108    fn bare() -> Self {
109        Select {
110            _marker: std::marker::PhantomData,
111            projections: Vec::new(),
112            filter: None,
113            order: Vec::new(),
114            limit: None,
115            start: 0,
116            fetch: Vec::new(),
117            group_by: Vec::new(),
118            group_all: false,
119            count: false,
120            count_alias: None,
121        }
122    }
123
124    pub fn filter(mut self, expr: impl DynExpr + 'static) -> Self { self.filter = Some(Box::new(expr)); self }
125    pub fn limit(mut self, n: u32) -> Self { self.limit = Some(n); self }
126    pub fn start(mut self, n: u32) -> Self { self.start = n; self }
127    pub fn fetch(mut self, field: &str) -> Self { self.fetch.push(field.to_string()); self }
128    pub fn group_by<C: DynExpr>(mut self, col: C) -> Self {
129        let mut buf = String::new(); col.render_dyn(&mut buf); self.group_by.push(buf); self
130    }
131    /// `GROUP ALL` (whole-table aggregate, e.g. with `count()`).
132    pub fn group_all(mut self) -> Self { self.group_all = true; self }
133    /// Alias for the `count()` projection: `SELECT count() AS <alias>`.
134    pub fn count_as(mut self, alias: &'static str) -> Self { self.count = true; self.count_alias = Some(alias); self }
135
136    pub fn order_by<C: DynExpr>(mut self, col: C, dir: Order) -> Self {
137        let mut buf = String::new(); col.render_dyn(&mut buf); self.order.push((buf, dir)); self
138    }
139
140    pub fn order_asc<C: DynExpr>(self, col: C) -> Self { self.order_by(col, Order::Asc) }
141    pub fn order_desc<C: DynExpr>(self, col: C) -> Self { self.order_by(col, Order::Desc) }
142
143    fn render_select_list(&self, q: &mut String) {
144        if self.count {
145            q.push_str("count()");
146            if let Some(a) = self.count_alias { q.push_str(" AS "); q.push_str(a); }
147        } else if self.projections.is_empty() {
148            q.push('*');
149        } else {
150            for (i, p) in self.projections.iter().enumerate() {
151                if i > 0 { q.push_str(", "); }
152                p.render(q);
153            }
154        }
155    }
156
157    pub fn to_surrealql(&self) -> String {
158        let mut q = String::from("SELECT ");
159        self.render_select_list(&mut q);
160        q.push_str(" FROM ");
161        q.push_str(T::table_name());
162        if let Some(ref f) = self.filter { q.push_str(" WHERE "); f.render_dyn(&mut q); }
163        for (i, (col, dir)) in self.order.iter().enumerate() {
164            if i == 0 { q.push_str(" ORDER BY "); } else { q.push_str(", "); }
165            q.push_str(&format!("{col} {dir}"));
166        }
167        for (i, g) in self.group_by.iter().enumerate() {
168            if i == 0 { q.push_str(" GROUP BY "); } else { q.push_str(", "); }
169            q.push_str(g);
170        }
171        if self.group_all { q.push_str(" GROUP ALL"); }
172        if self.start > 0 { q.push_str(&format!(" START {}", self.start)); }
173        if let Some(n) = self.limit { q.push_str(&format!(" LIMIT {n}")); }
174        for f in &self.fetch { q.push_str(&format!(" FETCH {f}")); }
175        q
176    }
177}
178
179impl<T: SurrealRecord> std::fmt::Display for Select<T> {
180    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
181        write!(f, "{}", self.to_surrealql())
182    }
183}
184
185// ═══════════════════════════════════════════════════════════════════════════════
186// INSERT
187// ═══════════════════════════════════════════════════════════════════════════════
188
189pub struct Insert<T: SurrealRecord> {
190    data: Vec<T>,
191    return_fields: Vec<&'static str>,
192}
193
194impl<T: SurrealRecord> Insert<T> {
195    pub fn content(mut self, record: T) -> Self { self.data.push(record); self }
196    pub fn return_field(mut self, field: &'static str) -> Self { self.return_fields.push(field); self }
197    pub fn data(&self) -> &[T] { &self.data }
198
199    pub fn to_surrealql(&self) -> String {
200        let returning = if self.return_fields.is_empty() { "" } else { " RETURN AFTER" };
201        format!("INSERT INTO {} $data{}", T::table_name(), returning)
202    }
203}
204
205// ═══════════════════════════════════════════════════════════════════════════════
206// UPDATE
207// ═══════════════════════════════════════════════════════════════════════════════
208
209enum SetVal {
210    /// `SET k = v` where v is a rendered expression
211    Assign(String, String),
212    /// `MERGE <expr>`
213    Merge(String),
214    /// `CONTENT <expr>` (full replace)
215    Content(String),
216}
217
218pub struct Update<T: SurrealRecord> {
219    _marker: std::marker::PhantomData<T>,
220    target: Target,
221    filter: Option<Box<dyn DynExpr>>,
222    sets: Vec<SetVal>,
223    returning: Returning,
224}
225
226impl<T: SurrealRecord> Update<T> {
227    pub(crate) fn for_table() -> Self {
228        Self { _marker: std::marker::PhantomData, target: Target::Table(T::table_name()),
229            filter: None, sets: Vec::new(), returning: Returning::None }
230    }
231
232    /// Target a single record: `UPDATE type::record('table', <id>)`.
233    pub fn record<V: SurrealQL>(mut self, id: V) -> Self {
234        self.target = Target::Record(RecordLink::new(T::table_name(), id));
235        self
236    }
237
238    pub fn filter(mut self, expr: impl DynExpr + 'static) -> Self { self.filter = Some(Box::new(expr)); self }
239
240    /// `SET col = <literal>`.
241    pub fn set<C: SurrealQL>(mut self, col: Column<T, C>, value: C) -> Self {
242        let mut buf = String::new(); C::render_literal(&value, &mut buf);
243        self.sets.push(SetVal::Assign(col.name.to_string(), buf)); self
244    }
245    /// `SET col = <literal>` by raw column name.
246    pub fn set_lit<C: SurrealQL>(mut self, col: &str, value: C) -> Self {
247        let mut buf = String::new(); C::render_literal(&value, &mut buf);
248        self.sets.push(SetVal::Assign(col.to_string(), buf)); self
249    }
250    /// `SET col = <expr>` — e.g. a record link, `time::now()`, NONE, `use_count + 1`.
251    pub fn set_expr(mut self, col: &str, expr: impl DynExpr) -> Self {
252        let mut buf = String::new(); expr.render_dyn(&mut buf);
253        self.sets.push(SetVal::Assign(col.to_string(), buf)); self
254    }
255    /// `SET col = <raw SurrealQL>`.
256    pub fn set_raw(mut self, col: &str, raw: impl Into<String>) -> Self {
257        self.sets.push(SetVal::Assign(col.to_string(), raw.into())); self
258    }
259    /// `MERGE <expr>` — deep-merge the given object into the record.
260    pub fn merge(mut self, expr: impl DynExpr) -> Self {
261        let mut buf = String::new(); expr.render_dyn(&mut buf);
262        self.sets.push(SetVal::Merge(buf)); self
263    }
264    /// `CONTENT <expr>` — full-replace the record's content (upsert by record id).
265    pub fn content(mut self, expr: impl DynExpr) -> Self {
266        let mut buf = String::new(); expr.render_dyn(&mut buf);
267        self.sets.push(SetVal::Content(buf)); self
268    }
269    pub fn returning(mut self, r: Returning) -> Self { self.returning = r; self }
270
271    pub fn to_surrealql(&self) -> String {
272        let mut q = String::from("UPDATE ");
273        self.target.render(&mut q);
274        // SurrealQL order: SET/MERGE/CONTENT first, then WHERE, then RETURN.
275        let mut set_pairs = Vec::new();
276        let mut merge_clause = None;
277        let mut content_clause = None;
278        for s in &self.sets {
279            match s {
280                SetVal::Assign(k, v) => set_pairs.push(format!("{k} = {v}")),
281                SetVal::Merge(v) => merge_clause = Some(v.clone()),
282                SetVal::Content(v) => content_clause = Some(v.clone()),
283            }
284        }
285        if let Some(c) = content_clause {
286            q.push_str(" CONTENT ");
287            q.push_str(&c);
288        } else if let Some(m) = merge_clause {
289            q.push_str(" MERGE ");
290            q.push_str(&m);
291        } else if !set_pairs.is_empty() {
292            q.push_str(" SET ");
293            q.push_str(&set_pairs.join(", "));
294        }
295        if let Some(ref f) = self.filter { q.push_str(" WHERE "); f.render_dyn(&mut q); }
296        self.returning.render(&mut q);
297        q
298    }
299}
300
301impl<T: SurrealRecord> std::fmt::Display for Update<T> {
302    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
303        write!(f, "{}", self.to_surrealql())
304    }
305}
306
307// ═══════════════════════════════════════════════════════════════════════════════
308// CREATE
309// ═══════════════════════════════════════════════════════════════════════════════
310
311enum CreateBody {
312    /// `CONTENT <expr>`
313    Content(String),
314    /// `SET a = x, b = y`
315    Set(Vec<(String, String)>),
316}
317
318/// `CREATE <target> [CONTENT … | SET …] [RETURN …]`.
319pub struct Create<T: SurrealRecord> {
320    _marker: std::marker::PhantomData<T>,
321    target: Target,
322    body: CreateBody,
323    returning: Returning,
324}
325
326impl<T: SurrealRecord> Create<T> {
327    pub(crate) fn for_table() -> Self {
328        Self { _marker: std::marker::PhantomData, target: Target::Table(T::table_name()),
329            body: CreateBody::Set(Vec::new()), returning: Returning::None }
330    }
331
332    /// Target a single record id: `CREATE type::record('table', <id>)`.
333    pub fn record<V: SurrealQL>(mut self, id: V) -> Self {
334        self.target = Target::Record(RecordLink::new(T::table_name(), id));
335        self
336    }
337
338    /// `CONTENT <expr>` — replaces any accumulated SET pairs.
339    pub fn content(mut self, expr: impl DynExpr) -> Self {
340        let mut buf = String::new(); expr.render_dyn(&mut buf);
341        self.body = CreateBody::Content(buf); self
342    }
343
344    /// `SET col = <literal>`.
345    pub fn set_lit<C: SurrealQL>(mut self, col: &str, value: C) -> Self {
346        let mut buf = String::new(); C::render_literal(&value, &mut buf);
347        self.push_set(col, buf); self
348    }
349    /// `SET col = <expr>`.
350    pub fn set_expr(mut self, col: &str, expr: impl DynExpr) -> Self {
351        let mut buf = String::new(); expr.render_dyn(&mut buf);
352        self.push_set(col, buf); self
353    }
354    /// `SET col = <raw SurrealQL>`.
355    pub fn set_raw(mut self, col: &str, raw: impl Into<String>) -> Self {
356        self.push_set(col, raw.into()); self
357    }
358
359    fn push_set(&mut self, col: &str, rendered: String) {
360        match &mut self.body {
361            CreateBody::Set(v) => v.push((col.to_string(), rendered)),
362            CreateBody::Content(_) => {
363                self.body = CreateBody::Set(vec![(col.to_string(), rendered)]);
364            }
365        }
366    }
367
368    pub fn returning(mut self, r: Returning) -> Self { self.returning = r; self }
369
370    pub fn to_surrealql(&self) -> String {
371        let mut q = String::from("CREATE ");
372        self.target.render(&mut q);
373        match &self.body {
374            CreateBody::Content(c) => { q.push_str(" CONTENT "); q.push_str(c); }
375            CreateBody::Set(pairs) if !pairs.is_empty() => {
376                q.push_str(" SET ");
377                q.push_str(&pairs.iter().map(|(k, v)| format!("{k} = {v}")).collect::<Vec<_>>().join(", "));
378            }
379            CreateBody::Set(_) => {}
380        }
381        self.returning.render(&mut q);
382        q
383    }
384}
385
386impl<T: SurrealRecord> std::fmt::Display for Create<T> {
387    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
388        write!(f, "{}", self.to_surrealql())
389    }
390}
391
392// ═══════════════════════════════════════════════════════════════════════════════
393// DELETE
394// ═══════════════════════════════════════════════════════════════════════════════
395
396pub struct Delete<T: SurrealRecord> {
397    _marker: std::marker::PhantomData<T>,
398    target: Target,
399    filter: Option<Box<dyn DynExpr>>,
400    returning: Returning,
401}
402
403impl<T: SurrealRecord> Delete<T> {
404    pub(crate) fn for_table() -> Self {
405        Self { _marker: std::marker::PhantomData, target: Target::Table(T::table_name()),
406            filter: None, returning: Returning::None }
407    }
408    /// Target a single record: `DELETE type::record('table', <id>)`.
409    pub fn record<V: SurrealQL>(mut self, id: V) -> Self {
410        self.target = Target::Record(RecordLink::new(T::table_name(), id));
411        self
412    }
413    pub fn filter(mut self, expr: impl DynExpr + 'static) -> Self { self.filter = Some(Box::new(expr)); self }
414    pub fn returning(mut self, r: Returning) -> Self { self.returning = r; self }
415    pub fn to_surrealql(&self) -> String {
416        let mut q = String::from("DELETE ");
417        self.target.render(&mut q);
418        if let Some(ref f) = self.filter { q.push_str(" WHERE "); f.render_dyn(&mut q); }
419        self.returning.render(&mut q);
420        q
421    }
422}
423
424impl<T: SurrealRecord> std::fmt::Display for Delete<T> {
425    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
426        write!(f, "{}", self.to_surrealql())
427    }
428}
429
430// ═══════════════════════════════════════════════════════════════════════════════
431// Batch — multiple statements joined by `;` (mutate-then-reselect pattern)
432// ═══════════════════════════════════════════════════════════════════════════════
433
434/// Concatenates SurrealQL statements with `;` separators. The store's typical
435/// pattern is a mutation followed by a SELECT that re-projects the row.
436#[derive(Default)]
437pub struct Batch {
438    statements: Vec<String>,
439}
440
441impl Batch {
442    pub fn new() -> Self { Self { statements: Vec::new() } }
443    pub fn push(mut self, stmt: impl ToString) -> Self {
444        self.statements.push(stmt.to_string()); self
445    }
446    pub fn to_surrealql(&self) -> String {
447        self.statements.join(";\n")
448    }
449    /// Number of statements (useful for `.take(n)` indexing on the response).
450    pub fn len(&self) -> usize { self.statements.len() }
451    pub fn is_empty(&self) -> bool { self.statements.is_empty() }
452}
453
454impl std::fmt::Display for Batch {
455    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
456        write!(f, "{}", self.to_surrealql())
457    }
458}
459
460// ═══════════════════════════════════════════════════════════════════════════════
461// RELATE — graph edges
462// ═══════════════════════════════════════════════════════════════════════════════
463
464pub struct Relate<E: SurrealEdge> {
465    _marker: std::marker::PhantomData<E>,
466}
467
468impl<E: SurrealEdge> Relate<E> {
469    pub fn new() -> Self { Self { _marker: std::marker::PhantomData } }
470
471    pub fn to_surrealql(
472        from: &Thing<impl SurrealRecord>,
473        to: &Thing<impl SurrealRecord>,
474    ) -> String {
475        format!("RELATE {}:{} -> {} -> {}:{}",
476            from.table(), from.key, E::edge_name(), to.table(), to.key)
477    }
478}
479
480impl<E: SurrealEdge> Default for Relate<E> {
481    fn default() -> Self { Self::new() }
482}
483
484// ═══════════════════════════════════════════════════════════════════════════════
485// RELATE with content
486// ═══════════════════════════════════════════════════════════════════════════════
487
488/// Build a RELATE query with edge content.
489///
490/// ```ignore
491/// RelateEdge::<Follows>::from(user).to(other).content(Follows { since: now }).build()
492/// ```
493pub struct RelateEdge<E: SurrealEdge> {
494    _marker: std::marker::PhantomData<E>,
495    from_label: String,
496    to_label: String,
497    content_json: Option<serde_json::Value>,
498}
499
500impl<E: SurrealEdge> RelateEdge<E> {
501    pub fn from(from: &Thing<impl SurrealRecord>) -> Self {
502        Self { _marker: std::marker::PhantomData,
503            from_label: format!("{}:{}", from.table(), from.key),
504            to_label: String::new(), content_json: None }
505    }
506
507    pub fn to(mut self, to: &Thing<impl SurrealRecord>) -> Self {
508        self.to_label = format!("{}:{}", to.table(), to.key); self
509    }
510
511    /// Attach content to the edge record.
512    pub fn content(mut self, edge: &impl serde::Serialize) -> Self {
513        self.content_json = serde_json::to_value(edge).ok(); self
514    }
515
516    pub fn build(&self) -> String {
517        let mut q = format!("RELATE {} -> {} -> {}",
518            self.from_label, E::edge_name(), self.to_label);
519        if let Some(ref c) = self.content_json {
520            q.push_str(&format!(" CONTENT {}", serde_json::to_string(c).unwrap_or_default()));
521        }
522        q
523    }
524}