gluesql_core/executor/
update.rs

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}