1use {
2 super::{
3 select::select,
4 validate::{validate_unique, ColumnValidation},
5 },
6 crate::{
7 ast::{ColumnDef, ColumnUniqueOption, Expr, ForeignKey, Query, SetExpr, Values},
8 data::{Key, Row, Schema, Value},
9 executor::{evaluate::evaluate_stateless, limit::Limit},
10 result::Result,
11 store::{DataRow, GStore, GStoreMut},
12 },
13 futures::stream::{self, StreamExt, TryStreamExt},
14 serde::Serialize,
15 std::{fmt::Debug, rc::Rc},
16 thiserror::Error as ThisError,
17};
18
19#[derive(ThisError, Serialize, Debug, PartialEq, Eq)]
20pub enum InsertError {
21 #[error("table not found: {0}")]
22 TableNotFound(String),
23
24 #[error("lack of required column: {0}")]
25 LackOfRequiredColumn(String),
26
27 #[error("wrong column name: {0}")]
28 WrongColumnName(String),
29
30 #[error("column and values not matched")]
31 ColumnAndValuesNotMatched,
32
33 #[error("literals have more values than target columns")]
34 TooManyValues,
35
36 #[error("only single value accepted for schemaless row insert")]
37 OnlySingleValueAcceptedForSchemalessRow,
38
39 #[error("map type required: {0}")]
40 MapTypeValueRequired(String),
41
42 #[error("cannot find referenced value on {table_name}.{column_name} with value {referenced_value:?}")]
43 CannotFindReferencedValue {
44 table_name: String,
45 column_name: String,
46 referenced_value: String,
47 },
48
49 #[error("unreachable referencing column name: {0}")]
50 ConflictReferencingColumnName(String),
51}
52
53enum RowsData {
54 Append(Vec<DataRow>),
55 Insert(Vec<(Key, DataRow)>),
56}
57
58pub async fn insert<T: GStore + GStoreMut>(
59 storage: &mut T,
60 table_name: &str,
61 columns: &[String],
62 source: &Query,
63) -> Result<usize> {
64 let Schema {
65 column_defs,
66 foreign_keys,
67 ..
68 } = storage
69 .fetch_schema(table_name)
70 .await?
71 .ok_or_else(|| InsertError::TableNotFound(table_name.to_owned()))?;
72
73 let rows = match column_defs {
74 Some(column_defs) => {
75 fetch_vec_rows(
76 storage,
77 table_name,
78 column_defs,
79 columns,
80 source,
81 foreign_keys,
82 )
83 .await
84 }
85 None => fetch_map_rows(storage, source).await.map(RowsData::Append),
86 }?;
87
88 match rows {
89 RowsData::Append(rows) => {
90 let num_rows = rows.len();
91
92 storage
93 .append_data(table_name, rows)
94 .await
95 .map(|_| num_rows)
96 }
97 RowsData::Insert(rows) => {
98 let num_rows = rows.len();
99
100 storage
101 .insert_data(table_name, rows)
102 .await
103 .map(|_| num_rows)
104 }
105 }
106}
107
108async fn fetch_vec_rows<T: GStore>(
109 storage: &T,
110 table_name: &str,
111 column_defs: Vec<ColumnDef>,
112 columns: &[String],
113 source: &Query,
114 foreign_keys: Vec<ForeignKey>,
115) -> Result<RowsData> {
116 let labels = Rc::from(
117 column_defs
118 .iter()
119 .map(|column_def| column_def.name.to_owned())
120 .collect::<Vec<_>>(),
121 );
122 let column_defs = Rc::from(column_defs);
123 let column_validation = ColumnValidation::All(&column_defs);
124
125 #[derive(futures_enum::Stream)]
126 enum Rows<I1, I2> {
127 Values(I1),
128 Select(I2),
129 }
130
131 let rows = match &source.body {
132 SetExpr::Values(Values(values_list)) => {
133 let limit = Limit::new(source.limit.as_ref(), source.offset.as_ref()).await?;
134 let rows = stream::iter(values_list).then(|values| {
135 let column_defs = Rc::clone(&column_defs);
136 let labels = Rc::clone(&labels);
137
138 async move {
139 Ok(Row::Vec {
140 columns: labels,
141 values: fill_values(&column_defs, columns, values).await?,
142 })
143 }
144 });
145 let rows = limit.apply(rows);
146 let rows = rows.map(|row| row?.try_into_vec());
147
148 Rows::Values(rows)
149 }
150 SetExpr::Select(_) => {
151 let rows = select(storage, source, None).await?.map(|row| {
152 let values = row?.try_into_vec()?;
153
154 column_defs
155 .iter()
156 .zip(values.iter())
157 .try_for_each(|(column_def, value)| {
158 let ColumnDef {
159 data_type,
160 nullable,
161 ..
162 } = column_def;
163
164 value.validate_type(data_type)?;
165 value.validate_null(*nullable)
166 })?;
167
168 Ok(values)
169 });
170
171 Rows::Select(rows)
172 }
173 }
174 .try_collect::<Vec<Vec<Value>>>()
175 .await?;
176
177 validate_unique(
178 storage,
179 table_name,
180 column_validation,
181 rows.iter().map(|values| values.as_slice()),
182 )
183 .await?;
184
185 validate_foreign_key(storage, &column_defs, foreign_keys, &rows).await?;
186
187 let primary_key = column_defs.iter().position(|ColumnDef { unique, .. }| {
188 unique == &Some(ColumnUniqueOption { is_primary: true })
189 });
190
191 match primary_key {
192 Some(i) => rows
193 .into_iter()
194 .filter_map(|values| {
195 values
196 .get(i)
197 .map(Key::try_from)
198 .map(|result| result.map(|key| (key, values.into())))
199 })
200 .collect::<Result<Vec<_>>>()
201 .map(RowsData::Insert),
202 None => Ok(RowsData::Append(rows.into_iter().map(Into::into).collect())),
203 }
204}
205
206async fn validate_foreign_key<T: GStore>(
207 storage: &T,
208 column_defs: &Rc<[ColumnDef]>,
209 foreign_keys: Vec<ForeignKey>,
210 rows: &[Vec<Value>],
211) -> Result<()> {
212 for foreign_key in foreign_keys {
213 let ForeignKey {
214 referencing_column_name,
215 referenced_table_name,
216 referenced_column_name,
217 ..
218 } = &foreign_key;
219
220 let target_index = column_defs
221 .iter()
222 .enumerate()
223 .find(|(_, c)| &c.name == referencing_column_name)
224 .ok_or_else(|| {
225 InsertError::ConflictReferencingColumnName(referencing_column_name.to_owned())
226 })?;
227
228 for row in rows.iter() {
229 let value =
230 row.get(target_index.0)
231 .ok_or(InsertError::ConflictReferencingColumnName(
232 referencing_column_name.to_owned(),
233 ))?;
234
235 if value == &Value::Null {
236 continue;
237 }
238
239 let no_referenced = storage
240 .fetch_data(referenced_table_name, &Key::try_from(value)?)
241 .await?
242 .is_none();
243
244 if no_referenced {
245 return Err(InsertError::CannotFindReferencedValue {
246 table_name: referenced_table_name.to_owned(),
247 column_name: referenced_column_name.to_owned(),
248 referenced_value: String::from(value),
249 }
250 .into());
251 }
252 }
253 }
254
255 Ok(())
256}
257
258async fn fetch_map_rows<T: GStore>(storage: &T, source: &Query) -> Result<Vec<DataRow>> {
259 #[derive(futures_enum::Stream)]
260 enum Rows<I1, I2> {
261 Values(I1),
262 Select(I2),
263 }
264
265 let rows = match &source.body {
266 SetExpr::Values(Values(values_list)) => {
267 let limit = Limit::new(source.limit.as_ref(), source.offset.as_ref()).await?;
268 let rows = stream::iter(values_list).then(|values| async move {
269 if values.len() > 1 {
270 return Err(InsertError::OnlySingleValueAcceptedForSchemalessRow.into());
271 }
272
273 evaluate_stateless(None, &values[0])
274 .await?
275 .try_into()
276 .map(Row::Map)
277 });
278 let rows = limit.apply(rows);
279 let rows = rows.map_ok(Into::into);
280
281 Rows::Values(rows)
282 }
283 SetExpr::Select(_) => {
284 let rows = select(storage, source, None).await?.map(|row| {
285 let row = row?;
286
287 if let Row::Vec { values, .. } = &row {
288 if values.len() > 1 {
289 return Err(InsertError::OnlySingleValueAcceptedForSchemalessRow.into());
290 } else if !matches!(&values[0], Value::Map(_)) {
291 return Err(InsertError::MapTypeValueRequired((&values[0]).into()).into());
292 }
293 }
294
295 Ok(row.into())
296 });
297
298 Rows::Select(rows)
299 }
300 }
301 .try_collect::<Vec<DataRow>>()
302 .await?;
303
304 Ok(rows)
305}
306
307async fn fill_values(
308 column_defs: &[ColumnDef],
309 columns: &[String],
310 values: &[Expr],
311) -> Result<Vec<Value>> {
312 if !columns.is_empty() && values.len() != columns.len() {
313 return Err(InsertError::ColumnAndValuesNotMatched.into());
314 } else if values.len() > column_defs.len() {
315 return Err(InsertError::TooManyValues.into());
316 }
317
318 if let Some(wrong_column_name) = columns.iter().find(|column_name| {
319 !column_defs
320 .iter()
321 .any(|column_def| &&column_def.name == column_name)
322 }) {
323 return Err(InsertError::WrongColumnName(wrong_column_name.to_owned()).into());
324 }
325
326 #[derive(iter_enum::Iterator)]
327 enum Columns<I1, I2> {
328 All(I1),
329 Specified(I2),
330 }
331
332 let columns = if columns.is_empty() {
333 Columns::All(column_defs.iter().map(|ColumnDef { name, .. }| name))
334 } else {
335 Columns::Specified(columns.iter())
336 };
337
338 let column_name_value_list = columns.zip(values.iter()).collect::<Vec<(_, _)>>();
339
340 let values = stream::iter(column_defs)
341 .then(|column_def| {
342 let column_name_value_list = &column_name_value_list;
343
344 async move {
345 let ColumnDef {
346 name: def_name,
347 data_type,
348 nullable,
349 ..
350 } = column_def;
351
352 let value = column_name_value_list
353 .iter()
354 .find(|(name, _)| name == &def_name)
355 .map(|(_, value)| value);
356
357 match (value, &column_def.default, nullable) {
358 (Some(&expr), _, _) | (None, Some(expr), _) => evaluate_stateless(None, expr)
359 .await?
360 .try_into_value(data_type, *nullable),
361 (None, None, true) => Ok(Value::Null),
362 (None, None, false) => {
363 Err(InsertError::LackOfRequiredColumn(def_name.to_owned()).into())
364 }
365 }
366 }
367 })
368 .try_collect::<Vec<Value>>()
369 .await?;
370
371 Ok(values)
372}