pg_worm/query/
update.rs

1use std::{
2    future::{Future, IntoFuture},
3    marker::PhantomData,
4    pin::Pin,
5};
6
7use tokio_postgres::types::ToSql;
8
9use crate::TypedColumn;
10
11use super::{push_all_with_sep, Executable, PushChunk, Query, SqlChunk, ToQuery, Where};
12
13/// State representing that an UPDATE
14/// has been set.
15///
16/// `UPDATE` queries in this state cannot be executed.
17#[doc(hidden)]
18pub struct NoneSet;
19/// State representing that an UDPATE
20/// has been set.
21#[doc(hidden)]
22pub struct SomeSet;
23
24/// A struct for building `UPDATE` queries.
25///
26/// The query can only be executed once at least one
27/// update has been made.
28pub struct Update<'a, State = NoneSet> {
29    table: &'static str,
30    updates: Vec<SqlChunk<'a>>,
31    where_: Where<'a>,
32    state: PhantomData<State>,
33}
34
35impl<'a> ToQuery<'a, u64> for Update<'a, SomeSet> {}
36
37impl<'a, T> Update<'a, T> {
38    /// Begin building a new `UPDATE` query.
39    pub fn new(table: &'static str) -> Update<'a, NoneSet> {
40        Update {
41            table,
42            updates: vec![],
43            where_: Where::Empty,
44            state: PhantomData::<NoneSet>,
45        }
46    }
47
48    /// Add a `WHERE` to the query.
49    ///
50    /// If called multiple times, the conditions are
51    /// joined using `AND`.
52    pub fn where_(mut self, where_: Where<'a>) -> Update<'a, T> {
53        self.where_ = self.where_.and(where_);
54
55        self
56    }
57
58    /// Add a raw `WHERE` clause to your query.
59    ///
60    /// You can reference the `params` by using the `?` placeholder in your statement.
61    ///
62    /// Note: you need to pass the exact types Postgres is expecting.
63    /// Failure to do so will result in (sometimes confusing) runtime errors.
64    ///
65    /// Otherwise this behaves exactly like `where_`.
66    ///
67    /// # Example
68    ///
69    /// ```ignore
70    /// Book::select()
71    ///     .where_(Book::id.neq(&3))
72    ///     .where_raw("complex_function(book.title, ?, ?)", vec![&true, &"Foobar"])
73    ///     .await?;
74    /// ```
75    pub fn where_raw(
76        self,
77        statement: impl Into<String>,
78        params: Vec<&'a (dyn ToSql + Sync)>,
79    ) -> Update<'a, T> {
80        let where_ = Where::new(statement.into(), params);
81
82        self.where_(where_)
83    }
84
85    /// Add a `SET` instruction to your `UPDATE` query.
86    ///
87    /// This function has to be called at least once before
88    /// you can execute the query.
89    pub fn set<U: ToSql + Sync>(
90        mut self,
91        col: TypedColumn<U>,
92        value: &'a U,
93    ) -> Update<'a, SomeSet> {
94        self.updates
95            .push(SqlChunk(format!("{} = ?", col.column_name), vec![value]));
96
97        Update {
98            state: PhantomData::<SomeSet>,
99            updates: self.updates,
100            where_: self.where_,
101            table: self.table,
102        }
103    }
104}
105
106impl<'a> PushChunk<'a> for Update<'a, SomeSet> {
107    fn push_to_buffer<T>(&mut self, buffer: &mut super::Query<'a, T>) {
108        // Which table to update
109        buffer.0.push_str("UPDATE ");
110        buffer.0.push_str(self.table);
111
112        // Which updates to make
113        buffer.0.push_str(" SET ");
114        push_all_with_sep(&mut self.updates, buffer, ", ");
115
116        // Which rows to update
117        if !self.where_.is_empty() {
118            buffer.0.push_str(" WHERE ");
119            self.where_.push_to_buffer(buffer);
120        }
121    }
122}
123
124impl<'a> IntoFuture for Update<'a, SomeSet>
125where
126    Update<'a, SomeSet>: ToQuery<'a, u64>,
127    Query<'a, u64>: Executable<Output = u64>,
128{
129    type Output = Result<u64, crate::Error>;
130
131    type IntoFuture = Pin<Box<dyn Future<Output = Self::Output> + 'a>>;
132
133    fn into_future(mut self) -> Self::IntoFuture {
134        let query = self.to_query();
135
136        Box::pin(async move { query.exec().await })
137    }
138}