Skip to main content

drizzle_postgres/builder/
update.rs

1use crate::common::PostgresSchemaType;
2use crate::values::PostgresValue;
3use core::marker::PhantomData;
4use drizzle_core::{SQLTable, ToSQL};
5
6// Import the ExecutableState trait
7use super::ExecutableState;
8
9//------------------------------------------------------------------------------
10// Type State Markers
11//------------------------------------------------------------------------------
12
13pub use drizzle_core::builder::{
14    UpdateInitial, UpdateReturningSet, UpdateSetClauseSet, UpdateWhereSet,
15};
16
17/// Marker for the state after FROM clause
18#[derive(Debug, Clone, Copy, Default)]
19pub struct UpdateFromSet;
20
21// Mark states that can execute update queries
22impl ExecutableState for UpdateFromSet {}
23
24//------------------------------------------------------------------------------
25// UpdateBuilder Definition
26//------------------------------------------------------------------------------
27
28/// Builds an UPDATE query specifically for `PostgreSQL`
29pub type UpdateBuilder<'a, Schema, State, Table, Marker = (), Row = ()> =
30    super::QueryBuilder<'a, Schema, State, Table, Marker, Row>;
31
32type ReturningMarker<Table, Columns> = drizzle_core::Scoped<
33    <Columns as drizzle_core::IntoSelectTarget>::Marker,
34    drizzle_core::Cons<Table, drizzle_core::Nil>,
35>;
36
37type ReturningRow<Table, Columns> =
38    <<Columns as drizzle_core::IntoSelectTarget>::Marker as drizzle_core::ResolveRow<Table>>::Row;
39
40type ReturningBuilder<'a, S, T, Columns> = UpdateBuilder<
41    'a,
42    S,
43    UpdateReturningSet,
44    T,
45    ReturningMarker<T, Columns>,
46    ReturningRow<T, Columns>,
47>;
48
49//------------------------------------------------------------------------------
50// Initial State Implementation
51//------------------------------------------------------------------------------
52
53impl<'a, Schema, Table> UpdateBuilder<'a, Schema, UpdateInitial, Table>
54where
55    Table: SQLTable<'a, PostgresSchemaType, PostgresValue<'a>>,
56{
57    /// Sets the values to update and transitions to the `SetClauseSet` state
58    #[inline]
59    pub fn set(
60        self,
61        values: Table::Update,
62    ) -> UpdateBuilder<'a, Schema, UpdateSetClauseSet, Table> {
63        let sql = crate::helpers::set::<Table, PostgresSchemaType, PostgresValue<'a>>(&values);
64        drop(values);
65        UpdateBuilder {
66            sql: self.sql.append(sql),
67            schema: PhantomData,
68            state: PhantomData,
69            table: PhantomData,
70            marker: PhantomData,
71            row: PhantomData,
72            grouped: PhantomData,
73        }
74    }
75}
76
77//------------------------------------------------------------------------------
78// Post-SET Implementation
79//------------------------------------------------------------------------------
80
81impl<'a, S, T> UpdateBuilder<'a, S, UpdateSetClauseSet, T> {
82    /// Adds a FROM clause and transitions to the `FromSet` state
83    #[inline]
84    pub fn from(
85        self,
86        source: impl ToSQL<'a, PostgresValue<'a>>,
87    ) -> UpdateBuilder<'a, S, UpdateFromSet, T> {
88        let from_sql = crate::helpers::from(source);
89        UpdateBuilder {
90            sql: self.sql.append(from_sql),
91            schema: PhantomData,
92            state: PhantomData,
93            table: PhantomData,
94            marker: PhantomData,
95            row: PhantomData,
96            grouped: PhantomData,
97        }
98    }
99
100    /// Adds a WHERE condition and transitions to the `WhereSet` state
101    #[inline]
102    pub fn r#where<E>(self, condition: E) -> UpdateBuilder<'a, S, UpdateWhereSet, T>
103    where
104        E: drizzle_core::expr::Expr<'a, PostgresValue<'a>>,
105        E::SQLType: drizzle_core::types::BooleanLike,
106    {
107        let where_sql = crate::helpers::r#where(condition);
108        UpdateBuilder {
109            sql: self.sql.append(where_sql),
110            schema: PhantomData,
111            state: PhantomData,
112            table: PhantomData,
113            marker: PhantomData,
114            row: PhantomData,
115            grouped: PhantomData,
116        }
117    }
118
119    /// Adds a RETURNING clause and transitions to the `ReturningSet` state
120    #[inline]
121    pub fn returning<Columns>(self, columns: Columns) -> ReturningBuilder<'a, S, T, Columns>
122    where
123        Columns: ToSQL<'a, PostgresValue<'a>> + drizzle_core::IntoSelectTarget,
124        Columns::Marker: drizzle_core::ResolveRow<T>,
125    {
126        let returning_sql = crate::helpers::returning(columns);
127        UpdateBuilder {
128            sql: self.sql.append(returning_sql),
129            schema: PhantomData,
130            state: PhantomData,
131            table: PhantomData,
132            marker: PhantomData,
133            row: PhantomData,
134            grouped: PhantomData,
135        }
136    }
137}
138
139//------------------------------------------------------------------------------
140// Post-FROM Implementation
141//------------------------------------------------------------------------------
142
143impl<'a, S, T> UpdateBuilder<'a, S, UpdateFromSet, T> {
144    /// Adds a WHERE condition after FROM
145    #[inline]
146    pub fn r#where<E>(self, condition: E) -> UpdateBuilder<'a, S, UpdateWhereSet, T>
147    where
148        E: drizzle_core::expr::Expr<'a, PostgresValue<'a>>,
149        E::SQLType: drizzle_core::types::BooleanLike,
150    {
151        let where_sql = crate::helpers::r#where(condition);
152        UpdateBuilder {
153            sql: self.sql.append(where_sql),
154            schema: PhantomData,
155            state: PhantomData,
156            table: PhantomData,
157            marker: PhantomData,
158            row: PhantomData,
159            grouped: PhantomData,
160        }
161    }
162
163    /// Adds a RETURNING clause after FROM
164    #[inline]
165    pub fn returning<Columns>(self, columns: Columns) -> ReturningBuilder<'a, S, T, Columns>
166    where
167        Columns: ToSQL<'a, PostgresValue<'a>> + drizzle_core::IntoSelectTarget,
168        Columns::Marker: drizzle_core::ResolveRow<T>,
169    {
170        let returning_sql = crate::helpers::returning(columns);
171        UpdateBuilder {
172            sql: self.sql.append(returning_sql),
173            schema: PhantomData,
174            state: PhantomData,
175            table: PhantomData,
176            marker: PhantomData,
177            row: PhantomData,
178            grouped: PhantomData,
179        }
180    }
181}
182
183//------------------------------------------------------------------------------
184// Post-WHERE Implementation
185//------------------------------------------------------------------------------
186
187impl<'a, S, T> UpdateBuilder<'a, S, UpdateWhereSet, T> {
188    /// Adds a RETURNING clause after WHERE
189    #[inline]
190    pub fn returning<Columns>(self, columns: Columns) -> ReturningBuilder<'a, S, T, Columns>
191    where
192        Columns: ToSQL<'a, PostgresValue<'a>> + drizzle_core::IntoSelectTarget,
193        Columns::Marker: drizzle_core::ResolveRow<T>,
194    {
195        let returning_sql = crate::helpers::returning(columns);
196        UpdateBuilder {
197            sql: self.sql.append(returning_sql),
198            schema: PhantomData,
199            state: PhantomData,
200            table: PhantomData,
201            marker: PhantomData,
202            row: PhantomData,
203            grouped: PhantomData,
204        }
205    }
206}
207
208#[cfg(test)]
209mod tests {
210    use super::*;
211    use drizzle_core::{SQL, ToSQL};
212
213    #[test]
214    fn test_update_builder_creation() {
215        let builder = UpdateBuilder::<(), UpdateInitial, ()> {
216            sql: SQL::raw("UPDATE test"),
217            schema: PhantomData,
218            state: PhantomData,
219            table: PhantomData,
220            marker: PhantomData,
221            row: PhantomData,
222            grouped: PhantomData,
223        };
224
225        assert_eq!(builder.to_sql().sql(), "UPDATE test");
226    }
227}