pg_worm/query/
select.rs

1use std::{
2    future::{Future, IntoFuture},
3    marker::PhantomData,
4    ops::Deref,
5    pin::Pin,
6};
7
8use tokio_postgres::{types::ToSql, Row};
9
10use super::{Executable, PushChunk, Query, ToQuery, Where};
11use crate::Column;
12
13/// A struct which holds the information needed to build
14/// a `SELECT` query.
15pub struct Select<'a, T = Vec<Row>> {
16    cols: Vec<Column>,
17    from: &'static str,
18    where_: Where<'a>,
19    marker: PhantomData<T>,
20    limit: Option<u64>,
21    offset: Option<u64>,
22}
23
24impl<'a, T> ToQuery<'a, T> for Select<'a, T> {}
25
26impl<'a, T> Select<'a, T> {
27    #[doc(hidden)]
28    pub fn new(cols: &[&dyn Deref<Target = Column>], from: &'static str) -> Select<'a, T> {
29        Select {
30            cols: cols.iter().map(|i| (***i)).collect(),
31            from,
32            where_: Where::Empty,
33            marker: PhantomData::<T>,
34            limit: None,
35            offset: None,
36        }
37    }
38
39    /// Add a `WHERE` clause to your query.
40    ///
41    /// If used multiple time, the conditions are joined
42    /// using `AND`.
43    pub fn where_(mut self, where_: Where<'a>) -> Select<'a, T> {
44        if self.where_.is_empty() {
45            self.where_ = where_;
46        } else {
47            self.where_ = self.where_.and(where_);
48        }
49
50        self
51    }
52
53    /// Add a raw `WHERE` clause to your query.
54    ///
55    /// You can reference the `params` by using the `?` placeholder in your statement.
56    ///
57    /// Note: you need to pass the exact types Postgres is expecting.
58    /// Failure to do so will result in (sometimes confusing) runtime errors.
59    ///
60    /// Otherwise this behaves exactly like `where_`.
61    ///
62    /// # Example
63    ///
64    /// ```ignore
65    /// Book::select()
66    ///     .where_(Book::id.neq(&3))
67    ///     .where_raw("complex_function(book.title, ?, ?)", vec![&true, &"Foobar"])
68    ///     .await?;
69    /// ```
70    pub fn where_raw(
71        self,
72        statement: impl Into<String>,
73        params: Vec<&'a (dyn ToSql + Sync)>,
74    ) -> Select<'a, T> {
75        let where_ = Where::new(statement.into(), params);
76
77        self.where_(where_)
78    }
79
80    /// Add a `LIMIT` to your query.
81    pub fn limit(mut self, limit: u64) -> Select<'a, T> {
82        self.limit = Some(limit);
83
84        self
85    }
86
87    /// Add an `OFFSET` to your query.
88    pub fn offset(mut self, offset: u64) -> Select<'a, T> {
89        self.offset = Some(offset);
90
91        self
92    }
93}
94
95impl<'a, T> PushChunk<'a> for Select<'a, T> {
96    fn push_to_buffer<B>(&mut self, buffer: &mut Query<'a, B>) {
97        buffer.0.push_str("SELECT ");
98
99        // Push the selected columns
100        let cols = self
101            .cols
102            .iter()
103            .map(|i| i.full_name())
104            .collect::<Vec<_>>()
105            .join(", ");
106        buffer.0.push_str(&cols);
107
108        // Push the table from which the columns
109        // are selected
110        buffer.0.push_str(" FROM ");
111        buffer.0.push_str(self.from);
112
113        // If it exists, push the WHERE clause
114        if !self.where_.is_empty() {
115            buffer.0.push_str(" WHERE ");
116            self.where_.push_to_buffer(buffer);
117        }
118
119        // If set, add a LIMIT
120        if let Some(limit) = self.limit {
121            buffer.0.push_str(" LIMIT ");
122            buffer.0.push_str(&limit.to_string());
123        }
124
125        // If set, add an OFFSET
126        if let Some(offset) = self.offset {
127            buffer.0.push_str(" OFFSET ");
128            buffer.0.push_str(&offset.to_string())
129        }
130    }
131}
132
133impl<'a, T: Sync + Send + 'a> IntoFuture for Select<'a, T>
134where
135    Select<'a, T>: ToQuery<'a, T>,
136    Query<'a, T>: Executable<Output = T>,
137{
138    type IntoFuture = Pin<Box<dyn Future<Output = Self::Output> + 'a>>;
139    type Output = Result<T, crate::Error>;
140
141    fn into_future(mut self) -> Self::IntoFuture {
142        let query = self.to_query();
143        Box::pin(async move { query.exec().await })
144    }
145}
146
147#[cfg(test)]
148mod test {
149    #![allow(dead_code)]
150
151    use crate::prelude::*;
152
153    #[derive(Model)]
154    struct Book {
155        #[column(primary_key, auto)]
156        id: i64,
157        title: String,
158    }
159
160    #[test]
161    fn select_limit() {
162        let query = Book::select().limit(3).to_query().0;
163        assert_eq!(query, "SELECT book.id, book.title FROM book LIMIT 3");
164    }
165
166    #[test]
167    fn select_offset() {
168        let query = Book::select().offset(4).to_query().0;
169        assert_eq!(query, "SELECT book.id, book.title FROM book OFFSET 4");
170    }
171}