1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
use crate::{
Result,
engine::exec::{Action, Exec, Output, VarId},
};
use toasty_core::{
driver::{ExecResponse, Rows, operation},
schema::db::{ColumnId, TableId},
stmt::{self, ValueStream},
};
#[derive(Debug, Clone)]
pub(crate) struct UpdateByKey {
/// If specified, use the input to generate the list of keys to update
pub input: VarId,
/// Where to store the result of the update
pub output: Output,
/// Which table to update
pub table: TableId,
/// Assignments
pub assignments: stmt::Assignments,
/// Only update keys that match the filter
pub filter: Option<stmt::Expr>,
/// Fail the update if the condition is not met
pub condition: Option<stmt::Expr>,
/// The columns to return for each updated row *after* the update. When
/// `None`, just return the count of updated rows.
pub returning: Option<Vec<ColumnId>>,
}
impl Exec<'_> {
pub(super) async fn action_update_by_key(&mut self, action: &UpdateByKey) -> Result<()> {
let keys = self
.vars
.load(action.input)
.await?
.values
.collect_as_value()
.await?
.into_list_unwrap();
// Shred a multi-key update into one single-key op per key so each key's
// filter is adjudicated independently — matching SQL's per-row
// semantics, and mirroring how delete fans out. These updates are not
// atomic.
let mut total_count = 0u64;
let mut rows = vec![];
for key in keys {
match self.exec_update_one(action, key).await? {
Rows::Count(n) => total_count += n,
other => rows.extend(other.into_value_stream().collect().await?),
}
}
// The output shape is a property of the action, not the results: with
// zero keys there is nothing to match on, yet a `returning` update must
// still yield an (empty) stream rather than a count.
let res = if action.returning.is_some() {
Rows::value_stream(ValueStream::from_vec(rows))
} else {
Rows::Count(total_count)
};
self.vars.store(
action.output.var,
action.output.num_uses,
ExecResponse::from_rows(res),
);
Ok(())
}
/// Execute a single-key `UpdateByKey` op for one resolved key.
async fn exec_update_one(&mut self, action: &UpdateByKey, key: stmt::Value) -> Result<Rows> {
let op = operation::UpdateByKey {
table: action.table,
keys: vec![key],
assignments: action.assignments.clone(),
filter: action.filter.clone(),
condition: action.condition.clone(),
returning: action.returning.clone(),
};
let res = self.connection.exec(&self.engine.schema, op.into()).await?;
debug_assert_eq!(!res.values.is_count(), action.returning.is_some());
Ok(res.values)
}
}
impl From<UpdateByKey> for Action {
fn from(src: UpdateByKey) -> Self {
Self::UpdateByKey(src)
}
}