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 {
59        Self {
60            _marker: std::marker::PhantomData,
61        }
62    }
63
64    pub fn select(self, _cols: crate::expr::ColumnSet<T>) -> Select<T> {
65        Select::bare()
66    }
67
68    /// Select an explicit projection list (`SELECT <fields…> FROM table`).
69    pub fn project(self, fields: Vec<Projection>) -> Select<T> {
70        let mut s = Select::bare();
71        s.projections = fields;
72        s
73    }
74
75    /// `SELECT count() FROM table GROUP ALL`
76    pub fn count(self, _field: &str) -> Select<T> {
77        let mut s = Select::bare();
78        s.count = true;
79        s.group_all = true;
80        s
81    }
82
83    pub fn insert(self) -> Insert<T> {
84        Insert {
85            data: Vec::new(),
86            return_fields: vec![],
87        }
88    }
89    pub fn create(self) -> Create<T> {
90        Create::for_table()
91    }
92    pub fn update(self) -> Update<T> {
93        Update::for_table()
94    }
95    pub fn delete(self) -> Delete<T> {
96        Delete::for_table()
97    }
98}
99
100impl<T: SurrealRecord> Default for Table<T> {
101    fn default() -> Self {
102        Self::new()
103    }
104}
105
106// ═══════════════════════════════════════════════════════════════════════════════
107// SELECT
108// ═══════════════════════════════════════════════════════════════════════════════
109
110pub struct Select<T: SurrealRecord> {
111    _marker: std::marker::PhantomData<T>,
112    projections: Vec<Projection>,
113    filter: Option<Box<dyn DynExpr>>,
114    order: Vec<(String, Order)>,
115    limit: Option<u32>,
116    start: u32,
117    fetch: Vec<String>,
118    group_by: Vec<String>,
119    group_all: bool,
120    count: bool,
121    count_alias: Option<&'static str>,
122}
123
124impl<T: SurrealRecord> Select<T> {
125    fn bare() -> Self {
126        Select {
127            _marker: std::marker::PhantomData,
128            projections: Vec::new(),
129            filter: None,
130            order: Vec::new(),
131            limit: None,
132            start: 0,
133            fetch: Vec::new(),
134            group_by: Vec::new(),
135            group_all: false,
136            count: false,
137            count_alias: None,
138        }
139    }
140
141    pub fn filter(mut self, expr: impl DynExpr + 'static) -> Self {
142        self.filter = Some(Box::new(expr));
143        self
144    }
145    pub fn limit(mut self, n: u32) -> Self {
146        self.limit = Some(n);
147        self
148    }
149    pub fn start(mut self, n: u32) -> Self {
150        self.start = n;
151        self
152    }
153    pub fn fetch(mut self, field: &str) -> Self {
154        self.fetch.push(field.to_string());
155        self
156    }
157    pub fn group_by<C: DynExpr>(mut self, col: C) -> Self {
158        let mut buf = String::new();
159        col.render_dyn(&mut buf);
160        self.group_by.push(buf);
161        self
162    }
163    /// `GROUP ALL` (whole-table aggregate, e.g. with `count()`).
164    pub fn group_all(mut self) -> Self {
165        self.group_all = true;
166        self
167    }
168    /// Alias for the `count()` projection: `SELECT count() AS <alias>`.
169    pub fn count_as(mut self, alias: &'static str) -> Self {
170        self.count = true;
171        self.count_alias = Some(alias);
172        self
173    }
174
175    pub fn order_by<C: DynExpr>(mut self, col: C, dir: Order) -> Self {
176        let mut buf = String::new();
177        col.render_dyn(&mut buf);
178        self.order.push((buf, dir));
179        self
180    }
181
182    pub fn order_asc<C: DynExpr>(self, col: C) -> Self {
183        self.order_by(col, Order::Asc)
184    }
185    pub fn order_desc<C: DynExpr>(self, col: C) -> Self {
186        self.order_by(col, Order::Desc)
187    }
188
189    fn render_select_list(&self, q: &mut String) {
190        if self.count {
191            q.push_str("count()");
192            if let Some(a) = self.count_alias {
193                q.push_str(" AS ");
194                q.push_str(a);
195            }
196        } else if self.projections.is_empty() {
197            q.push('*');
198        } else {
199            for (i, p) in self.projections.iter().enumerate() {
200                if i > 0 {
201                    q.push_str(", ");
202                }
203                p.render(q);
204            }
205        }
206    }
207
208    pub fn to_surrealql(&self) -> String {
209        let mut q = String::from("SELECT ");
210        self.render_select_list(&mut q);
211        q.push_str(" FROM ");
212        q.push_str(T::table_name());
213        if let Some(ref f) = self.filter {
214            q.push_str(" WHERE ");
215            f.render_dyn(&mut q);
216        }
217        for (i, (col, dir)) in self.order.iter().enumerate() {
218            if i == 0 {
219                q.push_str(" ORDER BY ");
220            } else {
221                q.push_str(", ");
222            }
223            q.push_str(&format!("{col} {dir}"));
224        }
225        for (i, g) in self.group_by.iter().enumerate() {
226            if i == 0 {
227                q.push_str(" GROUP BY ");
228            } else {
229                q.push_str(", ");
230            }
231            q.push_str(g);
232        }
233        if self.group_all {
234            q.push_str(" GROUP ALL");
235        }
236        if self.start > 0 {
237            q.push_str(&format!(" START {}", self.start));
238        }
239        if let Some(n) = self.limit {
240            q.push_str(&format!(" LIMIT {n}"));
241        }
242        for f in &self.fetch {
243            q.push_str(&format!(" FETCH {f}"));
244        }
245        q
246    }
247}
248
249impl<T: SurrealRecord> std::fmt::Display for Select<T> {
250    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
251        write!(f, "{}", self.to_surrealql())
252    }
253}
254
255// ═══════════════════════════════════════════════════════════════════════════════
256// INSERT
257// ═══════════════════════════════════════════════════════════════════════════════
258
259pub struct Insert<T: SurrealRecord> {
260    data: Vec<T>,
261    return_fields: Vec<&'static str>,
262}
263
264impl<T: SurrealRecord> Insert<T> {
265    pub fn content(mut self, record: T) -> Self {
266        self.data.push(record);
267        self
268    }
269    pub fn return_field(mut self, field: &'static str) -> Self {
270        self.return_fields.push(field);
271        self
272    }
273    pub fn data(&self) -> &[T] {
274        &self.data
275    }
276
277    pub fn to_surrealql(&self) -> String {
278        let returning = if self.return_fields.is_empty() {
279            ""
280        } else {
281            " RETURN AFTER"
282        };
283        format!("INSERT INTO {} $data{}", T::table_name(), returning)
284    }
285}
286
287// ═══════════════════════════════════════════════════════════════════════════════
288// UPDATE
289// ═══════════════════════════════════════════════════════════════════════════════
290
291enum SetVal {
292    /// `SET k = v` where v is a rendered expression
293    Assign(String, String),
294    /// `MERGE <expr>`
295    Merge(String),
296    /// `CONTENT <expr>` (full replace)
297    Content(String),
298}
299
300pub struct Update<T: SurrealRecord> {
301    _marker: std::marker::PhantomData<T>,
302    target: Target,
303    filter: Option<Box<dyn DynExpr>>,
304    sets: Vec<SetVal>,
305    returning: Returning,
306}
307
308impl<T: SurrealRecord> Update<T> {
309    pub(crate) fn for_table() -> Self {
310        Self {
311            _marker: std::marker::PhantomData,
312            target: Target::Table(T::table_name()),
313            filter: None,
314            sets: Vec::new(),
315            returning: Returning::None,
316        }
317    }
318
319    /// Target a single record: `UPDATE type::record('table', <id>)`.
320    pub fn record<V: SurrealQL>(mut self, id: V) -> Self {
321        self.target = Target::Record(RecordLink::new(T::table_name(), id));
322        self
323    }
324
325    pub fn filter(mut self, expr: impl DynExpr + 'static) -> Self {
326        self.filter = Some(Box::new(expr));
327        self
328    }
329
330    /// `SET col = <literal>`.
331    pub fn set<C: SurrealQL>(mut self, col: Column<T, C>, value: C) -> Self {
332        let mut buf = String::new();
333        C::render_literal(&value, &mut buf);
334        self.sets.push(SetVal::Assign(col.name.to_string(), buf));
335        self
336    }
337    /// `SET col = <literal>` by raw column name.
338    pub fn set_lit<C: SurrealQL>(mut self, col: &str, value: C) -> Self {
339        let mut buf = String::new();
340        C::render_literal(&value, &mut buf);
341        self.sets.push(SetVal::Assign(col.to_string(), buf));
342        self
343    }
344    /// `SET col = <expr>` — e.g. a record link, `time::now()`, NONE, `use_count + 1`.
345    pub fn set_expr(mut self, col: &str, expr: impl DynExpr) -> Self {
346        let mut buf = String::new();
347        expr.render_dyn(&mut buf);
348        self.sets.push(SetVal::Assign(col.to_string(), buf));
349        self
350    }
351    /// `SET col = <raw SurrealQL>`.
352    pub fn set_raw(mut self, col: &str, raw: impl Into<String>) -> Self {
353        self.sets.push(SetVal::Assign(col.to_string(), raw.into()));
354        self
355    }
356    /// `MERGE <expr>` — deep-merge the given object into the record.
357    pub fn merge(mut self, expr: impl DynExpr) -> Self {
358        let mut buf = String::new();
359        expr.render_dyn(&mut buf);
360        self.sets.push(SetVal::Merge(buf));
361        self
362    }
363    /// `CONTENT <expr>` — full-replace the record's content (upsert by record id).
364    pub fn content(mut self, expr: impl DynExpr) -> Self {
365        let mut buf = String::new();
366        expr.render_dyn(&mut buf);
367        self.sets.push(SetVal::Content(buf));
368        self
369    }
370    pub fn returning(mut self, r: Returning) -> Self {
371        self.returning = r;
372        self
373    }
374
375    pub fn to_surrealql(&self) -> String {
376        let mut q = String::from("UPDATE ");
377        self.target.render(&mut q);
378        // SurrealQL order: SET/MERGE/CONTENT first, then WHERE, then RETURN.
379        let mut set_pairs = Vec::new();
380        let mut merge_clause = None;
381        let mut content_clause = None;
382        for s in &self.sets {
383            match s {
384                SetVal::Assign(k, v) => set_pairs.push(format!("{k} = {v}")),
385                SetVal::Merge(v) => merge_clause = Some(v.clone()),
386                SetVal::Content(v) => content_clause = Some(v.clone()),
387            }
388        }
389        if let Some(c) = content_clause {
390            q.push_str(" CONTENT ");
391            q.push_str(&c);
392        } else if let Some(m) = merge_clause {
393            q.push_str(" MERGE ");
394            q.push_str(&m);
395        } else if !set_pairs.is_empty() {
396            q.push_str(" SET ");
397            q.push_str(&set_pairs.join(", "));
398        }
399        if let Some(ref f) = self.filter {
400            q.push_str(" WHERE ");
401            f.render_dyn(&mut q);
402        }
403        self.returning.render(&mut q);
404        q
405    }
406}
407
408impl<T: SurrealRecord> std::fmt::Display for Update<T> {
409    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
410        write!(f, "{}", self.to_surrealql())
411    }
412}
413
414// ═══════════════════════════════════════════════════════════════════════════════
415// CREATE
416// ═══════════════════════════════════════════════════════════════════════════════
417
418enum CreateBody {
419    /// `CONTENT <expr>`
420    Content(String),
421    /// `SET a = x, b = y`
422    Set(Vec<(String, String)>),
423}
424
425/// `CREATE <target> [CONTENT … | SET …] [RETURN …]`.
426pub struct Create<T: SurrealRecord> {
427    _marker: std::marker::PhantomData<T>,
428    target: Target,
429    body: CreateBody,
430    returning: Returning,
431}
432
433impl<T: SurrealRecord> Create<T> {
434    pub(crate) fn for_table() -> Self {
435        Self {
436            _marker: std::marker::PhantomData,
437            target: Target::Table(T::table_name()),
438            body: CreateBody::Set(Vec::new()),
439            returning: Returning::None,
440        }
441    }
442
443    /// Target a single record id: `CREATE type::record('table', <id>)`.
444    pub fn record<V: SurrealQL>(mut self, id: V) -> Self {
445        self.target = Target::Record(RecordLink::new(T::table_name(), id));
446        self
447    }
448
449    /// `CONTENT <expr>` — replaces any accumulated SET pairs.
450    pub fn content(mut self, expr: impl DynExpr) -> Self {
451        let mut buf = String::new();
452        expr.render_dyn(&mut buf);
453        self.body = CreateBody::Content(buf);
454        self
455    }
456
457    /// `SET col = <literal>`.
458    pub fn set_lit<C: SurrealQL>(mut self, col: &str, value: C) -> Self {
459        let mut buf = String::new();
460        C::render_literal(&value, &mut buf);
461        self.push_set(col, buf);
462        self
463    }
464    /// `SET col = <expr>`.
465    pub fn set_expr(mut self, col: &str, expr: impl DynExpr) -> Self {
466        let mut buf = String::new();
467        expr.render_dyn(&mut buf);
468        self.push_set(col, buf);
469        self
470    }
471    /// `SET col = <raw SurrealQL>`.
472    pub fn set_raw(mut self, col: &str, raw: impl Into<String>) -> Self {
473        self.push_set(col, raw.into());
474        self
475    }
476
477    fn push_set(&mut self, col: &str, rendered: String) {
478        match &mut self.body {
479            CreateBody::Set(v) => v.push((col.to_string(), rendered)),
480            CreateBody::Content(_) => {
481                self.body = CreateBody::Set(vec![(col.to_string(), rendered)]);
482            }
483        }
484    }
485
486    pub fn returning(mut self, r: Returning) -> Self {
487        self.returning = r;
488        self
489    }
490
491    pub fn to_surrealql(&self) -> String {
492        let mut q = String::from("CREATE ");
493        self.target.render(&mut q);
494        match &self.body {
495            CreateBody::Content(c) => {
496                q.push_str(" CONTENT ");
497                q.push_str(c);
498            }
499            CreateBody::Set(pairs) if !pairs.is_empty() => {
500                q.push_str(" SET ");
501                q.push_str(
502                    &pairs
503                        .iter()
504                        .map(|(k, v)| format!("{k} = {v}"))
505                        .collect::<Vec<_>>()
506                        .join(", "),
507                );
508            }
509            CreateBody::Set(_) => {}
510        }
511        self.returning.render(&mut q);
512        q
513    }
514}
515
516impl<T: SurrealRecord> std::fmt::Display for Create<T> {
517    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
518        write!(f, "{}", self.to_surrealql())
519    }
520}
521
522// ═══════════════════════════════════════════════════════════════════════════════
523// DELETE
524// ═══════════════════════════════════════════════════════════════════════════════
525
526pub struct Delete<T: SurrealRecord> {
527    _marker: std::marker::PhantomData<T>,
528    target: Target,
529    filter: Option<Box<dyn DynExpr>>,
530    returning: Returning,
531}
532
533impl<T: SurrealRecord> Delete<T> {
534    pub(crate) fn for_table() -> Self {
535        Self {
536            _marker: std::marker::PhantomData,
537            target: Target::Table(T::table_name()),
538            filter: None,
539            returning: Returning::None,
540        }
541    }
542    /// Target a single record: `DELETE type::record('table', <id>)`.
543    pub fn record<V: SurrealQL>(mut self, id: V) -> Self {
544        self.target = Target::Record(RecordLink::new(T::table_name(), id));
545        self
546    }
547    pub fn filter(mut self, expr: impl DynExpr + 'static) -> Self {
548        self.filter = Some(Box::new(expr));
549        self
550    }
551    pub fn returning(mut self, r: Returning) -> Self {
552        self.returning = r;
553        self
554    }
555    pub fn to_surrealql(&self) -> String {
556        let mut q = String::from("DELETE ");
557        self.target.render(&mut q);
558        if let Some(ref f) = self.filter {
559            q.push_str(" WHERE ");
560            f.render_dyn(&mut q);
561        }
562        self.returning.render(&mut q);
563        q
564    }
565}
566
567impl<T: SurrealRecord> std::fmt::Display for Delete<T> {
568    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
569        write!(f, "{}", self.to_surrealql())
570    }
571}
572
573// ═══════════════════════════════════════════════════════════════════════════════
574// Batch — multiple statements joined by `;` (mutate-then-reselect pattern)
575// ═══════════════════════════════════════════════════════════════════════════════
576
577/// Concatenates SurrealQL statements with `;` separators. The store's typical
578/// pattern is a mutation followed by a SELECT that re-projects the row.
579#[derive(Default)]
580pub struct Batch {
581    statements: Vec<String>,
582}
583
584impl Batch {
585    pub fn new() -> Self {
586        Self {
587            statements: Vec::new(),
588        }
589    }
590    pub fn push(mut self, stmt: impl ToString) -> Self {
591        self.statements.push(stmt.to_string());
592        self
593    }
594    pub fn to_surrealql(&self) -> String {
595        self.statements.join(";\n")
596    }
597    /// Number of statements (useful for `.take(n)` indexing on the response).
598    pub fn len(&self) -> usize {
599        self.statements.len()
600    }
601    pub fn is_empty(&self) -> bool {
602        self.statements.is_empty()
603    }
604}
605
606impl std::fmt::Display for Batch {
607    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
608        write!(f, "{}", self.to_surrealql())
609    }
610}
611
612// ═══════════════════════════════════════════════════════════════════════════════
613// RELATE — graph edges
614// ═══════════════════════════════════════════════════════════════════════════════
615
616pub struct Relate<E: SurrealEdge> {
617    _marker: std::marker::PhantomData<E>,
618}
619
620impl<E: SurrealEdge> Relate<E> {
621    pub fn new() -> Self {
622        Self {
623            _marker: std::marker::PhantomData,
624        }
625    }
626
627    pub fn to_surrealql(
628        from: &Thing<impl SurrealRecord>,
629        to: &Thing<impl SurrealRecord>,
630    ) -> String {
631        format!(
632            "RELATE {}:{} -> {} -> {}:{}",
633            from.table(),
634            from.key,
635            E::edge_name(),
636            to.table(),
637            to.key
638        )
639    }
640}
641
642impl<E: SurrealEdge> Default for Relate<E> {
643    fn default() -> Self {
644        Self::new()
645    }
646}
647
648// ═══════════════════════════════════════════════════════════════════════════════
649// RELATE with content
650// ═══════════════════════════════════════════════════════════════════════════════
651
652/// Build a RELATE query with edge content.
653///
654/// ```ignore
655/// RelateEdge::<Follows>::from(user).to(other).content(Follows { since: now }).build()
656/// ```
657pub struct RelateEdge<E: SurrealEdge> {
658    _marker: std::marker::PhantomData<E>,
659    from_label: String,
660    to_label: String,
661    content_json: Option<serde_json::Value>,
662}
663
664impl<E: SurrealEdge> RelateEdge<E> {
665    pub fn from(from: &Thing<impl SurrealRecord>) -> Self {
666        Self {
667            _marker: std::marker::PhantomData,
668            from_label: format!("{}:{}", from.table(), from.key),
669            to_label: String::new(),
670            content_json: None,
671        }
672    }
673
674    pub fn to(mut self, to: &Thing<impl SurrealRecord>) -> Self {
675        self.to_label = format!("{}:{}", to.table(), to.key);
676        self
677    }
678
679    /// Attach content to the edge record.
680    pub fn content(mut self, edge: &impl serde::Serialize) -> Self {
681        self.content_json = serde_json::to_value(edge).ok();
682        self
683    }
684
685    pub fn build(&self) -> String {
686        let mut q = format!(
687            "RELATE {} -> {} -> {}",
688            self.from_label,
689            E::edge_name(),
690            self.to_label
691        );
692        if let Some(ref c) = self.content_json {
693            q.push_str(&format!(
694                " CONTENT {}",
695                serde_json::to_string(c).unwrap_or_default()
696            ));
697        }
698        q
699    }
700}