Skip to main content

gluesql_core/executor/
update.rs

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