1use {
2 super::{
3 context::RowContext,
4 evaluate::{evaluate, Evaluated},
5 },
6 crate::{
7 ast::{Assignment, ColumnDef, ColumnUniqueOption, ForeignKey},
8 data::{Key, Row, Value},
9 result::{Error, Result},
10 store::GStore,
11 },
12 futures::stream::{self, StreamExt, TryStreamExt},
13 serde::Serialize,
14 std::{borrow::Cow, fmt::Debug, rc::Rc},
15 thiserror::Error,
16 utils::HashMapExt,
17};
18
19#[derive(Error, Serialize, Debug, PartialEq, Eq)]
20pub enum UpdateError {
21 #[error("column not found {0}")]
22 ColumnNotFound(String),
23
24 #[error("update on primary key is not supported: {0}")]
25 UpdateOnPrimaryKeyNotSupported(String),
26
27 #[error("conflict on schema, row data does not fit to schema")]
28 ConflictOnSchema,
29
30 #[error("cannot find referenced value on {table_name}.{column_name} with value {referenced_value:?}")]
31 CannotFindReferencedValue {
32 table_name: String,
33 column_name: String,
34 referenced_value: String,
35 },
36}
37
38pub struct Update<'a, T: GStore> {
39 storage: &'a T,
40 table_name: &'a str,
41 fields: &'a [Assignment],
42 column_defs: Option<&'a [ColumnDef]>,
43}
44
45impl<'a, T: GStore> Update<'a, T> {
46 pub fn new(
47 storage: &'a T,
48 table_name: &'a str,
49 fields: &'a [Assignment],
50 column_defs: Option<&'a [ColumnDef]>,
51 ) -> Result<Self> {
52 if let Some(column_defs) = column_defs {
53 for assignment in fields.iter() {
54 let Assignment { id, .. } = assignment;
55
56 if column_defs.iter().all(|col_def| &col_def.name != id) {
57 return Err(UpdateError::ColumnNotFound(id.to_owned()).into());
58 } else if column_defs.iter().any(|ColumnDef { name, unique, .. }| {
59 name == id && matches!(unique, Some(ColumnUniqueOption { is_primary: true }))
60 }) {
61 return Err(UpdateError::UpdateOnPrimaryKeyNotSupported(id.to_owned()).into());
62 }
63 }
64 }
65
66 Ok(Self {
67 storage,
68 table_name,
69 fields,
70 column_defs,
71 })
72 }
73
74 pub async fn apply(&self, row: Row, foreign_keys: &[ForeignKey]) -> Result<Row> {
75 let context = RowContext::new(self.table_name, Cow::Borrowed(&row), None);
76 let context = Some(Rc::new(context));
77
78 let assignments = stream::iter(self.fields.iter())
79 .then(|assignment| {
80 let Assignment {
81 id,
82 value: value_expr,
83 } = assignment;
84 let context = context.as_ref().map(Rc::clone);
85
86 async move {
87 let evaluated = evaluate(self.storage, context, None, value_expr).await?;
88 let value = match self.column_defs {
89 Some(column_defs) => {
90 let ColumnDef {
91 data_type,
92 nullable,
93 ..
94 } = column_defs
95 .iter()
96 .find(|column_def| id == &column_def.name)
97 .ok_or(UpdateError::ConflictOnSchema)?;
98
99 let value = match evaluated {
100 Evaluated::Literal(v) => Value::try_from_literal(data_type, &v)?,
101 Evaluated::Value(v) => {
102 v.validate_type(data_type)?;
103 v
104 }
105 Evaluated::StrSlice {
106 source: s,
107 range: r,
108 } => Value::Str(s[r].to_owned()),
109 };
110
111 value.validate_null(*nullable)?;
112 value
113 }
114 None => evaluated.try_into()?,
115 };
116
117 Ok::<_, Error>((id.as_ref(), value))
118 }
119 })
120 .and_then(|(id, value)| async move {
121 if value == Value::Null {
122 return Ok((id, value));
123 }
124
125 for foreign_key in foreign_keys {
126 let ForeignKey {
127 referencing_column_name,
128 referenced_table_name,
129 referenced_column_name,
130 ..
131 } = foreign_key;
132
133 if referencing_column_name != id {
134 continue;
135 }
136
137 let no_referenced = self
138 .storage
139 .fetch_data(referenced_table_name, &Key::try_from(&value)?)
140 .await?
141 .is_none();
142
143 if no_referenced {
144 return Err(UpdateError::CannotFindReferencedValue {
145 table_name: referenced_table_name.to_owned(),
146 column_name: referenced_column_name.to_owned(),
147 referenced_value: String::from(value),
148 }
149 .into());
150 }
151 }
152
153 Ok((id, value))
154 })
155 .try_collect::<Vec<(&str, Value)>>()
156 .await?;
157
158 Ok(match row {
159 Row::Vec { columns, values } => {
160 let values = columns
161 .iter()
162 .zip(values)
163 .map(|(column, value)| {
164 assignments
165 .iter()
166 .find_map(|(id, new_value)| (column == id).then_some(new_value.clone()))
167 .unwrap_or(value)
168 })
169 .collect();
170
171 Row::Vec { columns, values }
172 }
173 Row::Map(values) => {
174 let assignments = assignments
175 .into_iter()
176 .map(|(id, value)| (id.to_owned(), value));
177
178 Row::Map(values.concat(assignments))
179 }
180 })
181 }
182}