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    /// Render `INSERT INTO <table> <object|array> [RETURN AFTER]`, serializing the
278    /// queued record(s) inline as SurrealQL object literals (JSON is a valid
279    /// subset). A single record renders as `{ … }`, multiple as `[ {…}, {…} ]`.
280    pub fn to_surrealql(&self) -> String
281    where
282        T: serde::Serialize,
283    {
284        let body = match self.data.as_slice() {
285            [] => "[]".to_string(),
286            [one] => serde_json::to_string(one).unwrap_or_else(|_| "{}".to_string()),
287            many => serde_json::to_string(many).unwrap_or_else(|_| "[]".to_string()),
288        };
289        let returning = if self.return_fields.is_empty() {
290            ""
291        } else {
292            " RETURN AFTER"
293        };
294        format!("INSERT INTO {} {}{}", T::table_name(), body, returning)
295    }
296}
297
298// ═══════════════════════════════════════════════════════════════════════════════
299// UPDATE
300// ═══════════════════════════════════════════════════════════════════════════════
301
302enum SetVal {
303    /// `SET k = v` where v is a rendered expression
304    Assign(String, String),
305    /// `MERGE <expr>`
306    Merge(String),
307    /// `CONTENT <expr>` (full replace)
308    Content(String),
309}
310
311pub struct Update<T: SurrealRecord> {
312    _marker: std::marker::PhantomData<T>,
313    target: Target,
314    filter: Option<Box<dyn DynExpr>>,
315    sets: Vec<SetVal>,
316    returning: Returning,
317}
318
319impl<T: SurrealRecord> Update<T> {
320    pub(crate) fn for_table() -> Self {
321        Self {
322            _marker: std::marker::PhantomData,
323            target: Target::Table(T::table_name()),
324            filter: None,
325            sets: Vec::new(),
326            returning: Returning::None,
327        }
328    }
329
330    /// Target a single record: `UPDATE type::record('table', <id>)`.
331    pub fn record<V: SurrealQL>(mut self, id: V) -> Self {
332        self.target = Target::Record(RecordLink::new(T::table_name(), id));
333        self
334    }
335
336    pub fn filter(mut self, expr: impl DynExpr + 'static) -> Self {
337        self.filter = Some(Box::new(expr));
338        self
339    }
340
341    /// `SET col = <literal>`.
342    pub fn set<C: SurrealQL>(mut self, col: Column<T, C>, value: C) -> Self {
343        let mut buf = String::new();
344        C::render_literal(&value, &mut buf);
345        self.sets.push(SetVal::Assign(col.name.to_string(), buf));
346        self
347    }
348    /// `SET col = <literal>` by raw column name.
349    pub fn set_lit<C: SurrealQL>(mut self, col: &str, value: C) -> Self {
350        let mut buf = String::new();
351        C::render_literal(&value, &mut buf);
352        self.sets.push(SetVal::Assign(col.to_string(), buf));
353        self
354    }
355    /// `SET col = <expr>` — e.g. a record link, `time::now()`, NONE, `use_count + 1`.
356    pub fn set_expr(mut self, col: &str, expr: impl DynExpr) -> Self {
357        let mut buf = String::new();
358        expr.render_dyn(&mut buf);
359        self.sets.push(SetVal::Assign(col.to_string(), buf));
360        self
361    }
362    /// `SET col = <raw SurrealQL>`.
363    pub fn set_raw(mut self, col: &str, raw: impl Into<String>) -> Self {
364        self.sets.push(SetVal::Assign(col.to_string(), raw.into()));
365        self
366    }
367    /// `MERGE <expr>` — deep-merge the given object into the record.
368    pub fn merge(mut self, expr: impl DynExpr) -> Self {
369        let mut buf = String::new();
370        expr.render_dyn(&mut buf);
371        self.sets.push(SetVal::Merge(buf));
372        self
373    }
374    /// `CONTENT <expr>` — full-replace the record's content (upsert by record id).
375    pub fn content(mut self, expr: impl DynExpr) -> Self {
376        let mut buf = String::new();
377        expr.render_dyn(&mut buf);
378        self.sets.push(SetVal::Content(buf));
379        self
380    }
381    pub fn returning(mut self, r: Returning) -> Self {
382        self.returning = r;
383        self
384    }
385
386    /// Follow this `UPDATE` with a reselecting [`Select`], joined as a `;`-separated
387    /// batch. See [`Create::then_select`] for motivation.
388    pub fn then_select(self, select: Select<T>) -> String {
389        format!("{};\n{}", self.to_surrealql(), select.to_surrealql())
390    }
391
392    pub fn to_surrealql(&self) -> String {
393        let mut q = String::from("UPDATE ");
394        self.target.render(&mut q);
395        // SurrealQL order: SET/MERGE/CONTENT first, then WHERE, then RETURN.
396        let mut set_pairs = Vec::new();
397        let mut merge_clause = None;
398        let mut content_clause = None;
399        for s in &self.sets {
400            match s {
401                SetVal::Assign(k, v) => set_pairs.push(format!("{k} = {v}")),
402                SetVal::Merge(v) => merge_clause = Some(v.clone()),
403                SetVal::Content(v) => content_clause = Some(v.clone()),
404            }
405        }
406        if let Some(c) = content_clause {
407            q.push_str(" CONTENT ");
408            q.push_str(&c);
409        } else if let Some(m) = merge_clause {
410            q.push_str(" MERGE ");
411            q.push_str(&m);
412        } else if !set_pairs.is_empty() {
413            q.push_str(" SET ");
414            q.push_str(&set_pairs.join(", "));
415        }
416        if let Some(ref f) = self.filter {
417            q.push_str(" WHERE ");
418            f.render_dyn(&mut q);
419        }
420        self.returning.render(&mut q);
421        q
422    }
423}
424
425impl<T: SurrealRecord> std::fmt::Display for Update<T> {
426    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
427        write!(f, "{}", self.to_surrealql())
428    }
429}
430
431// ═══════════════════════════════════════════════════════════════════════════════
432// CREATE
433// ═══════════════════════════════════════════════════════════════════════════════
434
435enum CreateBody {
436    /// `CONTENT <expr>`
437    Content(String),
438    /// `SET a = x, b = y`
439    Set(Vec<(String, String)>),
440}
441
442/// `CREATE <target> [CONTENT … | SET …] [RETURN …]`.
443pub struct Create<T: SurrealRecord> {
444    _marker: std::marker::PhantomData<T>,
445    target: Target,
446    body: CreateBody,
447    returning: Returning,
448}
449
450impl<T: SurrealRecord> Create<T> {
451    pub(crate) fn for_table() -> Self {
452        Self {
453            _marker: std::marker::PhantomData,
454            target: Target::Table(T::table_name()),
455            body: CreateBody::Set(Vec::new()),
456            returning: Returning::None,
457        }
458    }
459
460    /// Target a single record id: `CREATE type::record('table', <id>)`.
461    pub fn record<V: SurrealQL>(mut self, id: V) -> Self {
462        self.target = Target::Record(RecordLink::new(T::table_name(), id));
463        self
464    }
465
466    /// `CONTENT <expr>` — replaces any accumulated SET pairs.
467    pub fn content(mut self, expr: impl DynExpr) -> Self {
468        let mut buf = String::new();
469        expr.render_dyn(&mut buf);
470        self.body = CreateBody::Content(buf);
471        self
472    }
473
474    /// `SET col = <literal>`.
475    pub fn set_lit<C: SurrealQL>(mut self, col: &str, value: C) -> Self {
476        let mut buf = String::new();
477        C::render_literal(&value, &mut buf);
478        self.push_set(col, buf);
479        self
480    }
481    /// `SET col = <expr>`.
482    pub fn set_expr(mut self, col: &str, expr: impl DynExpr) -> Self {
483        let mut buf = String::new();
484        expr.render_dyn(&mut buf);
485        self.push_set(col, buf);
486        self
487    }
488    /// `SET col = <raw SurrealQL>`.
489    pub fn set_raw(mut self, col: &str, raw: impl Into<String>) -> Self {
490        self.push_set(col, raw.into());
491        self
492    }
493
494    fn push_set(&mut self, col: &str, rendered: String) {
495        match &mut self.body {
496            CreateBody::Set(v) => v.push((col.to_string(), rendered)),
497            CreateBody::Content(_) => {
498                self.body = CreateBody::Set(vec![(col.to_string(), rendered)]);
499            }
500        }
501    }
502
503    pub fn returning(mut self, r: Returning) -> Self {
504        self.returning = r;
505        self
506    }
507
508    /// Follow this `CREATE` with a reselecting [`Select`], joined as a `;`-separated
509    /// batch. The select is rendered immediately, producing a complete SurrealQL
510    /// string ready for `db.query()`.
511    ///
512    /// This replaces the manual `Batch::new().push(create).push(select).to_surrealql()`
513    /// pattern for mutate-then-reselect workflows.
514    pub fn then_select(self, select: Select<T>) -> String {
515        format!("{};\n{}", self.to_surrealql(), select.to_surrealql())
516    }
517
518    pub fn to_surrealql(&self) -> String {
519        let mut q = String::from("CREATE ");
520        self.target.render(&mut q);
521        match &self.body {
522            CreateBody::Content(c) => {
523                q.push_str(" CONTENT ");
524                q.push_str(c);
525            }
526            CreateBody::Set(pairs) if !pairs.is_empty() => {
527                q.push_str(" SET ");
528                q.push_str(
529                    &pairs
530                        .iter()
531                        .map(|(k, v)| format!("{k} = {v}"))
532                        .collect::<Vec<_>>()
533                        .join(", "),
534                );
535            }
536            CreateBody::Set(_) => {}
537        }
538        self.returning.render(&mut q);
539        q
540    }
541}
542
543impl<T: SurrealRecord> std::fmt::Display for Create<T> {
544    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
545        write!(f, "{}", self.to_surrealql())
546    }
547}
548
549// ═══════════════════════════════════════════════════════════════════════════════
550// DELETE
551// ═══════════════════════════════════════════════════════════════════════════════
552
553pub struct Delete<T: SurrealRecord> {
554    _marker: std::marker::PhantomData<T>,
555    target: Target,
556    filter: Option<Box<dyn DynExpr>>,
557    returning: Returning,
558}
559
560impl<T: SurrealRecord> Delete<T> {
561    pub(crate) fn for_table() -> Self {
562        Self {
563            _marker: std::marker::PhantomData,
564            target: Target::Table(T::table_name()),
565            filter: None,
566            returning: Returning::None,
567        }
568    }
569    /// Target a single record: `DELETE type::record('table', <id>)`.
570    pub fn record<V: SurrealQL>(mut self, id: V) -> Self {
571        self.target = Target::Record(RecordLink::new(T::table_name(), id));
572        self
573    }
574    pub fn filter(mut self, expr: impl DynExpr + 'static) -> Self {
575        self.filter = Some(Box::new(expr));
576        self
577    }
578    pub fn returning(mut self, r: Returning) -> Self {
579        self.returning = r;
580        self
581    }
582
583    /// Follow this `DELETE` with a reselecting [`Select`], joined as a `;`-separated
584    /// batch. See [`Create::then_select`] for motivation.
585    pub fn then_select(self, select: Select<T>) -> String {
586        format!("{};\n{}", self.to_surrealql(), select.to_surrealql())
587    }
588
589    pub fn to_surrealql(&self) -> String {
590        let mut q = String::from("DELETE ");
591        self.target.render(&mut q);
592        if let Some(ref f) = self.filter {
593            q.push_str(" WHERE ");
594            f.render_dyn(&mut q);
595        }
596        self.returning.render(&mut q);
597        q
598    }
599}
600
601impl<T: SurrealRecord> std::fmt::Display for Delete<T> {
602    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
603        write!(f, "{}", self.to_surrealql())
604    }
605}
606
607// ═══════════════════════════════════════════════════════════════════════════════
608// Batch — multiple statements joined by `;` (mutate-then-reselect pattern)
609// ═══════════════════════════════════════════════════════════════════════════════
610
611/// Concatenates SurrealQL statements with `;` separators. The store's typical
612/// pattern is a mutation followed by a SELECT that re-projects the row.
613#[derive(Default)]
614pub struct Batch {
615    statements: Vec<String>,
616}
617
618impl Batch {
619    pub fn new() -> Self {
620        Self {
621            statements: Vec::new(),
622        }
623    }
624    pub fn push(mut self, stmt: impl ToString) -> Self {
625        self.statements.push(stmt.to_string());
626        self
627    }
628    pub fn to_surrealql(&self) -> String {
629        self.statements.join(";\n")
630    }
631    /// Number of statements (useful for `.take(n)` indexing on the response).
632    pub fn len(&self) -> usize {
633        self.statements.len()
634    }
635    pub fn is_empty(&self) -> bool {
636        self.statements.is_empty()
637    }
638}
639
640impl std::fmt::Display for Batch {
641    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
642        write!(f, "{}", self.to_surrealql())
643    }
644}
645
646// ═══════════════════════════════════════════════════════════════════════════════
647// RELATE — graph edges
648// ═══════════════════════════════════════════════════════════════════════════════
649
650/// Render a record's id as `table:<escaped-key>` into `buf`.
651fn record_id(thing: &Thing<impl SurrealRecord>, buf: &mut String) {
652    buf.push_str(thing.table());
653    buf.push(':');
654    thing.key.render_id(buf);
655}
656
657/// Return a record's id as a `table:<escaped-key>` string.
658fn record_id_string(thing: &Thing<impl SurrealRecord>) -> String {
659    let mut s = String::new();
660    record_id(thing, &mut s);
661    s
662}
663
664pub struct Relate<E: SurrealEdge> {
665    _marker: std::marker::PhantomData<E>,
666}
667
668impl<E: SurrealEdge> Relate<E> {
669    pub fn new() -> Self {
670        Self {
671            _marker: std::marker::PhantomData,
672        }
673    }
674
675    pub fn to_surrealql(
676        from: &Thing<impl SurrealRecord>,
677        to: &Thing<impl SurrealRecord>,
678    ) -> String {
679        let mut q = String::from("RELATE ");
680        record_id(from, &mut q);
681        q.push_str(" -> ");
682        q.push_str(E::edge_name());
683        q.push_str(" -> ");
684        record_id(to, &mut q);
685        q
686    }
687}
688
689impl<E: SurrealEdge> Default for Relate<E> {
690    fn default() -> Self {
691        Self::new()
692    }
693}
694
695// ═══════════════════════════════════════════════════════════════════════════════
696// RELATE with content
697// ═══════════════════════════════════════════════════════════════════════════════
698
699/// Build a RELATE query with edge content.
700///
701/// ```ignore
702/// RelateEdge::<Follows>::from(user).to(other).content(Follows { since: now }).build()
703/// ```
704pub struct RelateEdge<E: SurrealEdge> {
705    _marker: std::marker::PhantomData<E>,
706    from_label: String,
707    to_label: String,
708    content_json: Option<serde_json::Value>,
709}
710
711impl<E: SurrealEdge> RelateEdge<E> {
712    pub fn from(from: &Thing<impl SurrealRecord>) -> Self {
713        Self {
714            _marker: std::marker::PhantomData,
715            from_label: record_id_string(from),
716            to_label: String::new(),
717            content_json: None,
718        }
719    }
720
721    pub fn to(mut self, to: &Thing<impl SurrealRecord>) -> Self {
722        self.to_label = record_id_string(to);
723        self
724    }
725
726    /// Attach content to the edge record.
727    pub fn content(mut self, edge: &impl serde::Serialize) -> Self {
728        self.content_json = serde_json::to_value(edge).ok();
729        self
730    }
731
732    pub fn build(&self) -> String {
733        let mut q = format!(
734            "RELATE {} -> {} -> {}",
735            self.from_label,
736            E::edge_name(),
737            self.to_label
738        );
739        if let Some(ref c) = self.content_json {
740            q.push_str(&format!(
741                " CONTENT {}",
742                serde_json::to_string(c).unwrap_or_default()
743            ));
744        }
745        q
746    }
747}