1use std::marker::PhantomData;
2
3use crate::ast::{QueryNode, TableNode, UpdateAssignment, UpdateNode};
4use crate::expression::{Column, Expression, Selection};
5use crate::function::TypedExpression;
6use crate::schema::Table;
7use crate::value::IntoSqlValue;
8
9use super::{Cte, Query};
10
11#[derive(Clone, Debug)]
12pub struct UpdateQuery<T: Table, O = ()> {
13 node: UpdateNode,
14 marker: PhantomData<fn() -> (T, O)>,
15}
16
17pub fn update_table<T: Table>() -> UpdateQuery<T> {
18 UpdateQuery {
19 node: UpdateNode {
20 ctes: Vec::new(),
21 table: TableNode {
22 name: T::NAME,
23 alias: None,
24 },
25 assignments: Vec::new(),
26 from: Vec::new(),
27 filter: None,
28 returning: Vec::new(),
29 },
30 marker: PhantomData,
31 }
32}
33
34impl<T: Table, O> UpdateQuery<T, O> {
35 pub fn with<C: Table>(mut self, cte: Cte<C>) -> Self {
36 self.node.ctes.push(cte.node);
37 self
38 }
39
40 pub fn set<V>(mut self, column: Column<T, V>, value: impl IntoSqlValue<V>) -> Self {
41 self.node.assignments.push(UpdateAssignment {
42 table: column.table_name(),
43 column: column.name(),
44 value: Expression::Value(value.into_sql_value()),
45 });
46 self
47 }
48
49 pub fn set_expression<V>(mut self, column: Column<T, V>, value: TypedExpression<V>) -> Self {
50 self.node.assignments.push(UpdateAssignment {
51 table: column.table_name(),
52 column: column.name(),
53 value: value.expression(),
54 });
55 self
56 }
57
58 pub fn from<F: Table>(mut self) -> Self {
59 self.node.from.push(TableNode {
60 name: F::NAME,
61 alias: None,
62 });
63 self
64 }
65
66 pub fn filter(mut self, expression: Expression) -> Self {
67 self.node.filter = Some(match self.node.filter.take() {
68 Some(existing) => existing.and(expression),
69 None => expression,
70 });
71 self
72 }
73
74 pub fn returning<S: Selection>(mut self, selection: S) -> UpdateQuery<T, S::Output> {
75 self.node.returning.extend(selection.expressions());
76 UpdateQuery {
77 node: self.node,
78 marker: PhantomData,
79 }
80 }
81}
82
83impl<T: Table, O> Query for UpdateQuery<T, O> {
84 type Output = O;
85
86 fn compile(self, dialect: &impl crate::Dialect) -> crate::Result<crate::CompiledQuery> {
87 crate::compiler::compile(QueryNode::Update(self.node), dialect)
88 }
89}