Skip to main content

kali/builder/
mod.rs

1use super::{builder::expr::Expr, builder::ordering::ColumnOrdering};
2use crate::column::Column;
3use std::marker::PhantomData;
4use value::Value;
5
6pub mod expr;
7pub mod ordering;
8pub mod value;
9
10pub struct Select;
11pub struct Insert;
12pub struct Update;
13pub struct Delete;
14pub struct OnConflicted;
15
16#[derive(PartialEq, Eq)]
17enum QueryKind {
18    Select,
19    Insert,
20    Update,
21    Delete,
22}
23
24macro_rules! generic_builder {
25    ($name:ident, $kind:expr) => {
26        pub fn $name(table: &'a str) -> Self {
27            Self {
28                table,
29                kind: $kind,
30                columns: None,
31                filter: None,
32                limit: None,
33                offset: None,
34                values: None,
35                set: None,
36                on_conflict: None,
37                returning: None,
38                order_by: Vec::new(),
39                _type: PhantomData,
40            }
41        }
42    };
43}
44pub trait Veccable<T> {
45    fn to_vec(self) -> Vec<T>;
46}
47
48impl<T> Veccable<T> for T {
49    fn to_vec(self) -> Vec<T> {
50        vec![self]
51    }
52}
53
54impl<T: Clone> Veccable<T> for (T, T) {
55    fn to_vec(self) -> Vec<T> {
56        vec![self.0.clone(), self.1.clone()]
57    }
58}
59
60impl<T: Clone> Veccable<T> for (T, T, T) {
61    fn to_vec(self) -> Vec<T> {
62        vec![self.0.clone(), self.1.clone(), self.2.clone()]
63    }
64}
65
66pub enum OnConflict {
67    Ignore,
68    Update,
69}
70
71pub struct QueryBuilder<'a, S, C: Column> {
72    table: &'a str,
73    kind: QueryKind,
74    columns: Option<&'a [C]>,
75    values: Option<Vec<(C, Value)>>,
76    set: Option<Vec<(C, Value)>>,
77    on_conflict: Option<(Vec<C>, OnConflict)>,
78    returning: Option<&'a [C]>,
79    filter: Option<Expr<'a, C>>,
80    limit: Option<i64>,
81    offset: Option<i64>,
82    order_by: Vec<ColumnOrdering<C>>,
83    _type: PhantomData<S>,
84}
85
86impl<'a, C: Column> QueryBuilder<'a, Select, C> {
87    generic_builder!(select_from, QueryKind::Select);
88
89    pub fn columns(mut self, columns: &'a [C]) -> Self {
90        self.columns = Some(columns);
91        self
92    }
93
94    pub fn order_by(mut self, ordering: ColumnOrdering<C>) -> Self {
95        self.order_by.push(ordering);
96        self
97    }
98
99    pub fn limit(mut self, limit: i64) -> Self {
100        self.limit = Some(limit);
101        self
102    }
103
104    pub fn offset(mut self, offset: i64) -> Self {
105        self.offset = Some(offset);
106        self
107    }
108}
109
110impl<'a, C: Column> QueryBuilder<'a, Insert, C> {
111    generic_builder!(insert_into, QueryKind::Insert);
112
113    pub fn value<V: Into<Value> + 'a>(mut self, col: C, value: V) -> Self {
114        if let Some(values) = &mut self.values {
115            values.push((col, value.into()));
116        } else {
117            self.values = Some(vec![(col, value.into())]);
118        }
119
120        self
121    }
122
123    pub fn on_conflict<CV: Veccable<C> + 'a>(
124        mut self,
125        columns: CV,
126        on_conflict: OnConflict,
127    ) -> QueryBuilder<'a, OnConflicted, C> {
128        self.on_conflict = Some((columns.to_vec(), on_conflict));
129        QueryBuilder {
130            table: self.table,
131            kind: self.kind,
132            columns: self.columns,
133            filter: self.filter,
134            limit: self.limit,
135            offset: self.offset,
136            values: self.values,
137            set: self.set,
138            on_conflict: self.on_conflict,
139            returning: self.returning,
140            order_by: self.order_by,
141            _type: PhantomData,
142        }
143    }
144
145    pub fn returning(mut self, columns: &'a [C]) -> Self {
146        self.returning = Some(columns);
147        self
148    }
149}
150
151impl<'a, C: Column> QueryBuilder<'a, OnConflicted, C> {
152    /// Only valid after `on_conflict`
153    pub fn set<V: Into<Value>>(mut self, column: C, value: V) -> Self {
154        if !self.on_conflict.is_some() {
155            // todo: should use typestate pattern to prevent this at compile time
156            panic!("Cannot set value without an ON CONFLICT clause");
157        }
158
159        let value = value.into();
160        if let Some(set) = &mut self.set {
161            set.push((column, value));
162        } else {
163            self.set = Some(vec![(column, value)]);
164        }
165        self
166    }
167
168    pub fn returning(mut self, columns: &'a [C]) -> Self {
169        self.returning = Some(columns);
170        self
171    }
172}
173
174impl<'a, C: Column> QueryBuilder<'a, Update, C> {
175    generic_builder!(update, QueryKind::Update);
176
177    pub fn set<V: Into<Value>>(mut self, column: C, value: V) -> Self {
178        let value = value.into();
179        if let Some(set) = &mut self.set {
180            set.push((column, value));
181        } else {
182            self.set = Some(vec![(column, value)]);
183        }
184        self
185    }
186
187    pub fn returning(mut self, columns: &'a [C]) -> Self {
188        self.returning = Some(columns);
189        self
190    }
191}
192
193impl<'a, C: Column> QueryBuilder<'a, Delete, C> {
194    generic_builder!(delete_from, QueryKind::Delete);
195
196    pub fn returning(mut self, columns: &'a [C]) -> Self {
197        self.returning = Some(columns);
198        self
199    }
200}
201
202macro_rules! assert_kind {
203    ($self:ident, $kind:pat) => {
204        match $self.kind {
205            $kind => {}
206            _ => panic!("Wrong query kind, expected {}", stringify!($kind)),
207        }
208    };
209}
210
211impl<'a, T, C: Column> QueryBuilder<'a, T, C> {
212    pub fn filter(mut self, expr: Expr<'a, C>) -> Self {
213        if let Some(where_clause) = self.filter {
214            self.filter = Some(where_clause.and(expr));
215        } else {
216            self.filter = Some(expr);
217        }
218
219        self
220    }
221
222    pub fn to_sql(self) -> (String, Vec<Value>) {
223        let mut values = Vec::new();
224        let mut query = match self.kind {
225            QueryKind::Select => "SELECT ".to_string(),
226            QueryKind::Insert => "INSERT INTO ".to_string(),
227            QueryKind::Update => "UPDATE ".to_string(),
228            QueryKind::Delete => "DELETE ".to_string(),
229        };
230
231        if let Some(columns) = self.columns {
232            assert_kind!(self, QueryKind::Select | QueryKind::Insert);
233            push_separated(&mut query, columns.iter(), |query, column| {
234                column.write(query);
235            });
236        } else if matches!(self.kind, QueryKind::Select) {
237            query.push_str("*");
238        }
239
240        match self.kind {
241            QueryKind::Select => {
242                query.push_str(" FROM ");
243                query.push_str(self.table);
244            }
245            QueryKind::Insert => {
246                query.push_str(self.table);
247            }
248            QueryKind::Update => {
249                query.push_str(self.table);
250            }
251            QueryKind::Delete => {
252                query.push_str(" FROM ");
253                query.push_str(self.table);
254            }
255        }
256
257        if let Some(sql_values) = self.values {
258            assert_kind!(self, QueryKind::Insert);
259            query.push_str(" (");
260            push_separated(&mut query, sql_values.iter(), |query, (column, _)| {
261                column.write(query);
262            });
263            query.push_str(") VALUES (");
264            push_separated(&mut query, sql_values.into_iter(), |query, (_, value)| {
265                query.push_str("?");
266                values.push(value);
267            });
268            query.push_str(")");
269        }
270
271        // match is necessary to prevent .set from being consumed
272        match self.kind {
273            QueryKind::Update => {
274                if let Some(set) = self.set {
275                    assert_kind!(self, QueryKind::Update);
276                    query.push_str(" SET ");
277                    push_separated(&mut query, set.into_iter(), |query, (column, value)| {
278                        column.write(query);
279                        query.push_str(" = ?");
280                        values.push(value);
281                    });
282                }
283            }
284            QueryKind::Insert => {
285                if let Some(on_conflict) = self.on_conflict {
286                    assert_kind!(self, QueryKind::Insert);
287                    query.push_str(" ON CONFLICT (");
288                    push_separated(&mut query, on_conflict.0.iter(), |query, column| {
289                        column.write(query);
290                    });
291                    query.push_str(") ");
292                    match on_conflict.1 {
293                        OnConflict::Ignore => query.push_str("DO NOTHING"),
294                        OnConflict::Update => {
295                            query.push_str("DO UPDATE SET ");
296                            // todo: should be enforced with typeset
297                            let set = self.set.expect("SET clause is required for UPDATE");
298                            push_separated(
299                                &mut query,
300                                set.into_iter(),
301                                |query, (column, value)| {
302                                    column.write(query);
303                                    query.push_str(" = ?");
304                                    values.push(value);
305                                },
306                            );
307                        }
308                    }
309                }
310            }
311            _ => {
312                assert!(
313                    self.set.is_none(),
314                    "SET clause is not allowed in this query kind"
315                );
316            }
317        }
318
319        if let Some(where_clause) = self.filter {
320            assert_kind!(
321                self,
322                QueryKind::Select | QueryKind::Update | QueryKind::Delete
323            );
324            query.push_str(" WHERE ");
325            where_clause.write(&mut query, &mut values);
326        }
327
328        if !self.order_by.is_empty() {
329            assert_kind!(self, QueryKind::Select);
330            query.push_str(" ORDER BY ");
331            push_separated(&mut query, self.order_by.iter(), |query, ordering| {
332                ordering.write(query);
333            });
334        }
335
336        if let Some(limit) = self.limit {
337            assert_kind!(self, QueryKind::Select);
338            query.push_str(" LIMIT ");
339            query.push_str(&limit.to_string());
340        }
341
342        if let Some(offset) = self.offset {
343            assert_kind!(self, QueryKind::Select);
344            query.push_str(" OFFSET ");
345            query.push_str(&offset.to_string());
346        }
347
348        if let Some(returning) = self.returning {
349            assert_kind!(
350                self,
351                QueryKind::Insert | QueryKind::Update | QueryKind::Delete
352            );
353            query.push_str(" RETURNING ");
354            push_separated(&mut query, returning.iter(), |query, column| {
355                column.write(query)
356            });
357        }
358
359        (query, values)
360    }
361
362    pub async fn fetch_one<'e, E, S>(mut self, executor: E) -> Result<S, sqlx::Error>
363    where
364        E: 'e + sqlx::Executor<'e, Database = sqlx::Sqlite>,
365        S: for<'r> sqlx::FromRow<'r, sqlx::sqlite::SqliteRow>,
366    {
367        if self.limit.is_none() && self.kind == QueryKind::Select {
368            self.limit = Some(1);
369        }
370
371        let (query, values) = self.to_sql();
372        let mut query = sqlx::query(&query);
373        for value in values.into_iter() {
374            query = value.bind_to(query);
375        }
376
377        query
378            .fetch_one(executor)
379            .await
380            .and_then(|row| S::from_row(&row))
381    }
382
383    pub async fn fetch_optional<'e, 'c: 'e, E, S>(
384        mut self,
385        executor: E,
386    ) -> Result<Option<S>, sqlx::Error>
387    where
388        E: 'e + sqlx::Executor<'c, Database = sqlx::Sqlite>,
389        S: for<'r> sqlx::FromRow<'r, sqlx::sqlite::SqliteRow>,
390    {
391        if self.limit.is_none() && self.kind == QueryKind::Select {
392            self.limit = Some(1);
393        }
394
395        let (query, values) = self.to_sql();
396        let mut query = sqlx::query(&query);
397        for value in values.into_iter() {
398            query = value.bind_to(query);
399        }
400
401        query
402            .fetch_optional(executor)
403            .await
404            .and_then(|row| match row {
405                Some(row) => S::from_row(&row).map(Some),
406                None => Ok(None),
407            })
408    }
409
410    pub async fn fetch_all<'e, 'c: 'e, E, S>(self, executor: E) -> Result<Vec<S>, sqlx::Error>
411    where
412        E: 'e + sqlx::Executor<'c, Database = sqlx::Sqlite>,
413        S: for<'r> sqlx::FromRow<'r, sqlx::sqlite::SqliteRow>,
414    {
415        let (query, values) = self.to_sql();
416        let mut query = sqlx::query(&query);
417        for value in values.into_iter() {
418            query = value.bind_to(query);
419        }
420
421        query.fetch_all(executor).await.and_then(|rows| {
422            rows.into_iter()
423                .map(|row| S::from_row(&row))
424                .collect::<Result<Vec<_>, _>>()
425        })
426    }
427
428    pub async fn execute<'e, 'c: 'e, E>(
429        self,
430        executor: E,
431    ) -> Result<sqlx::sqlite::SqliteQueryResult, sqlx::Error>
432    where
433        E: 'e + sqlx::Executor<'c, Database = sqlx::Sqlite>,
434    {
435        let (query, values) = self.to_sql();
436        let mut query = sqlx::query(&query);
437        for value in values.into_iter() {
438            query = value.bind_to(query);
439        }
440
441        query.execute(executor).await
442    }
443}
444
445fn push_separated<I, F>(query: &mut String, iter: I, mut cb: F)
446where
447    I: Iterator,
448    F: FnMut(&mut String, I::Item),
449{
450    let mut first = true;
451    for item in iter {
452        if !first {
453            query.push_str(", ");
454        } else {
455            first = false;
456        }
457        cb(query, item);
458    }
459}