1use {
2 super::{context::RowContext, evaluate::evaluate_stateless, filter::check_expr},
3 crate::{
4 ast::{
5 ColumnDef, ColumnUniqueOption, Dictionary, Expr, IndexItem, Join, Query, Select,
6 SelectItem, SetExpr, TableAlias, TableFactor, TableWithJoins, ToSql, ToSqlUnquoted,
7 Values,
8 },
9 data::{get_alias, get_index, Key, Row, Value},
10 executor::{evaluate::evaluate, select::select},
11 result::Result,
12 store::{DataRow, GStore},
13 },
14 async_recursion::async_recursion,
15 futures::{
16 future,
17 stream::{self, Stream, StreamExt, TryStreamExt},
18 },
19 serde::Serialize,
20 std::{borrow::Cow, collections::HashMap, fmt::Debug, iter, rc::Rc},
21 thiserror::Error as ThisError,
22};
23
24#[derive(ThisError, Serialize, Debug, PartialEq, Eq)]
25pub enum FetchError {
26 #[error("table not found: {0}")]
27 TableNotFound(String),
28
29 #[error("table alias not found: {0}")]
30 TableAliasNotFound(String),
31
32 #[error("SERIES has wrong size: {0}")]
33 SeriesSizeWrong(i64),
34
35 #[error("table '{0}' has {1} columns available but {2} column aliases specified")]
36 TooManyColumnAliases(String, usize, usize),
37}
38
39pub async fn fetch<'a, T: GStore>(
40 storage: &'a T,
41 table_name: &'a str,
42 columns: Option<Rc<[String]>>,
43 where_clause: Option<&'a Expr>,
44) -> Result<impl Stream<Item = Result<(Key, Row)>> + 'a> {
45 let columns = columns.unwrap_or_else(|| Rc::from([]));
46 let rows = storage
47 .scan_data(table_name)
48 .await?
49 .try_filter_map(move |(key, data_row)| {
50 let row = match data_row {
51 DataRow::Vec(values) => Row::Vec {
52 columns: Rc::clone(&columns),
53 values,
54 },
55 DataRow::Map(values) => Row::Map(values),
56 };
57
58 async move {
59 let expr = match where_clause {
60 None => {
61 return Ok(Some((key, row)));
62 }
63 Some(expr) => expr,
64 };
65
66 let context = RowContext::new(table_name, Cow::Borrowed(&row), None);
67
68 check_expr(storage, Some(Rc::new(context)), None, expr)
69 .await
70 .map(|pass| pass.then_some((key, row)))
71 }
72 });
73
74 Ok(rows)
75}
76
77#[derive(futures_enum::Stream)]
78pub enum Rows<I1, I2, I3, I4> {
79 Derived(I1),
80 Table(I2),
81 Series(I3),
82 Dictionary(I4),
83}
84
85pub async fn fetch_relation_rows<'a, T: GStore>(
86 storage: &'a T,
87 table_factor: &'a TableFactor,
88 filter_context: &Option<Rc<RowContext<'a>>>,
89) -> Result<impl Stream<Item = Result<Row>> + 'a> {
90 let columns = Rc::from(
91 fetch_relation_columns(storage, table_factor)
92 .await?
93 .unwrap_or_default(),
94 );
95
96 match table_factor {
97 TableFactor::Derived { subquery, .. } => {
98 let filter_context = filter_context.as_ref().map(Rc::clone);
99 let rows =
100 select(storage, subquery, filter_context)
101 .await?
102 .map_ok(move |row| match row {
103 Row::Vec { values, .. } => Row::Vec {
104 columns: Rc::clone(&columns),
105 values,
106 },
107 Row::Map(values) => Row::Map(values),
108 });
109
110 Ok(Rows::Derived(rows))
111 }
112 TableFactor::Table { name, .. } => {
113 let rows = {
114 #[derive(futures_enum::Stream)]
115 enum Rows<I1, I2, I3, I4> {
116 Indexed(I1),
117 PrimaryKey(I2),
118 PrimaryKeyEmpty(I3),
119 FullScan(I4),
120 }
121
122 match get_index(table_factor) {
123 Some(IndexItem::NonClustered {
124 name: index_name,
125 asc,
126 cmp_expr,
127 }) => {
128 let cmp_value = match cmp_expr {
129 Some((op, expr)) => {
130 let evaluated = evaluate(storage, None, None, expr).await?;
131
132 Some((op, evaluated.try_into()?))
133 }
134 None => None,
135 };
136
137 let rows = storage
138 .scan_indexed_data(name, index_name, *asc, cmp_value)
139 .await?
140 .map_ok(move |(_, data_row)| match data_row {
141 DataRow::Vec(values) => Row::Vec {
142 columns: Rc::clone(&columns),
143 values,
144 },
145 DataRow::Map(values) => Row::Map(values),
146 });
147
148 Rows::Indexed(rows)
149 }
150 Some(IndexItem::PrimaryKey(expr)) => {
151 let filter_context = filter_context.as_ref().map(Rc::clone);
152 let key = evaluate(storage, filter_context, None, expr)
153 .await
154 .and_then(Value::try_from)
155 .and_then(Key::try_from)?;
156
157 match storage.fetch_data(name, &key).await? {
158 Some(data_row) => {
159 let row = match data_row {
160 DataRow::Vec(values) => Row::Vec {
161 columns: Rc::clone(&columns),
162 values,
163 },
164 DataRow::Map(values) => Row::Map(values),
165 };
166
167 Rows::PrimaryKey(stream::once(future::ready(Ok(row))))
168 }
169 None => Rows::PrimaryKeyEmpty(stream::empty()),
170 }
171 }
172 _ => {
173 let rows = storage.scan_data(name).await?.map_ok(move |(_, data_row)| {
174 match data_row {
175 DataRow::Vec(values) => Row::Vec {
176 columns: Rc::clone(&columns),
177 values,
178 },
179 DataRow::Map(values) => Row::Map(values),
180 }
181 });
182
183 Rows::FullScan(rows)
184 }
185 }
186 };
187
188 Ok(Rows::Table(rows))
189 }
190 TableFactor::Series { size, .. } => {
191 let value: Value = evaluate_stateless(None, size).await?.try_into()?;
192 let size: i64 = value.try_into()?;
193 let size = match size {
194 n if n >= 0 => size,
195 n => return Err(FetchError::SeriesSizeWrong(n).into()),
196 };
197
198 let columns = Rc::from(vec!["N".to_owned()]);
199 let rows = (1..=size).map(move |v| {
200 Ok(Row::Vec {
201 columns: Rc::clone(&columns),
202 values: vec![Value::I64(v)],
203 })
204 });
205
206 Ok(Rows::Series(stream::iter(rows)))
207 }
208 TableFactor::Dictionary { dict, .. } => {
209 let rows = {
210 #[derive(futures_enum::Stream)]
211 enum Rows<I1, I2, I3, I4> {
212 Tables(I1),
213 TableColumns(I2),
214 Indexes(I3),
215 Objects(I4),
216 }
217
218 match dict {
219 Dictionary::GlueObjects => {
220 let schemas = storage.fetch_all_schemas().await?;
221 let table_metas = storage
222 .scan_table_meta()
223 .await?
224 .collect::<Result<HashMap<_, _>>>()?;
225 let rows = schemas.into_iter().flat_map(move |schema| {
226 let meta = table_metas
227 .iter()
228 .find_map(|(table_name, hash_map)| {
229 (table_name == &schema.table_name).then(|| hash_map.clone())
230 })
231 .unwrap_or_default();
232
233 let table_rows = HashMap::from([
234 ("OBJECT_NAME".to_owned(), Value::Str(schema.table_name)),
235 ("OBJECT_TYPE".to_owned(), Value::Str("TABLE".to_owned())),
236 ])
237 .into_iter()
238 .chain(meta)
239 .collect::<HashMap<_, _>>();
240
241 let index_rows = schema.indexes.into_iter().map(|index| {
242 HashMap::from([
243 ("OBJECT_NAME".to_owned(), Value::Str(index.name)),
244 ("OBJECT_TYPE".to_owned(), Value::Str("INDEX".to_owned())),
245 ])
246 });
247
248 iter::once(table_rows)
249 .chain(index_rows)
250 .map(|hash_map| Ok(Row::Map(hash_map)))
251 });
252
253 Rows::Objects(stream::iter(rows))
254 }
255 Dictionary::GlueTables => {
256 let schemas = storage.fetch_all_schemas().await?;
257 let rows = schemas.into_iter().map(move |schema| {
258 Ok(Row::Vec {
259 columns: Rc::clone(&columns),
260 values: vec![
261 Value::Str(schema.table_name),
262 schema.comment.map(Value::Str).unwrap_or(Value::Null),
263 ],
264 })
265 });
266
267 Rows::Tables(stream::iter(rows))
268 }
269 Dictionary::GlueTableColumns => {
270 let schemas = storage.fetch_all_schemas().await?;
271 let rows = schemas.into_iter().flat_map(move |schema| {
272 let columns = Rc::clone(&columns);
273 let table_name = schema.table_name;
274
275 schema
276 .column_defs
277 .unwrap_or_default()
278 .into_iter()
279 .enumerate()
280 .map(move |(index, column_def)| {
281 let values = vec![
282 Value::Str(table_name.clone()),
283 Value::Str(column_def.name),
284 Value::I64(index as i64 + 1),
285 Value::Bool(column_def.nullable),
286 column_def
287 .unique
288 .map(|unique| Value::Str(unique.to_sql()))
289 .unwrap_or(Value::Null),
290 column_def
291 .default
292 .map(|expr| Value::Str(expr.to_sql()))
293 .unwrap_or(Value::Null),
294 column_def.comment.map(Value::Str).unwrap_or(Value::Null),
295 ];
296
297 Ok(Row::Vec {
298 columns: Rc::clone(&columns),
299 values,
300 })
301 })
302 });
303
304 Rows::TableColumns(stream::iter(rows))
305 }
306 Dictionary::GlueIndexes => {
307 let schemas = storage.fetch_all_schemas().await?;
308 let rows = schemas.into_iter().flat_map(move |schema| {
309 let column_defs = schema.column_defs.unwrap_or_default();
310 let primary_column = column_defs.iter().find_map(|column_def| {
311 let ColumnDef { name, unique, .. } = column_def;
312
313 (unique == &Some(ColumnUniqueOption { is_primary: true }))
314 .then_some(name)
315 });
316
317 let clustered = match primary_column {
318 Some(column_name) => {
319 let values = vec![
320 Value::Str(schema.table_name.clone()),
321 Value::Str("PRIMARY".to_owned()),
322 Value::Str("BOTH".to_owned()),
323 Value::Str(column_name.to_owned()),
324 Value::Bool(true),
325 ];
326
327 let row = Row::Vec {
328 columns: Rc::clone(&columns),
329 values,
330 };
331
332 vec![Ok(row)]
333 }
334 None => Vec::new(),
335 };
336
337 let columns = Rc::clone(&columns);
338 let non_clustered = schema.indexes.into_iter().map(move |index| {
339 let values = vec![
340 Value::Str(schema.table_name.clone()),
341 Value::Str(index.name),
342 Value::Str(index.order.to_string()),
343 Value::Str(index.expr.to_sql_unquoted()),
344 Value::Bool(false),
345 ];
346
347 Ok(Row::Vec {
348 columns: Rc::clone(&columns),
349 values,
350 })
351 });
352
353 clustered.into_iter().chain(non_clustered)
354 });
355
356 Rows::Indexes(stream::iter(rows))
357 }
358 }
359 };
360
361 Ok(Rows::Dictionary(rows))
362 }
363 }
364}
365
366pub async fn fetch_columns<T: GStore>(
367 storage: &T,
368 table_name: &str,
369) -> Result<Option<Vec<String>>> {
370 let columns = storage
371 .fetch_schema(table_name)
372 .await?
373 .ok_or_else(|| FetchError::TableNotFound(table_name.to_owned()))?
374 .column_defs
375 .map(|column_defs| {
376 column_defs
377 .into_iter()
378 .map(|column_def| column_def.name)
379 .collect()
380 });
381
382 Ok(columns)
383}
384
385#[async_recursion(?Send)]
386pub async fn fetch_relation_columns<T>(
387 storage: &T,
388 table_factor: &TableFactor,
389) -> Result<Option<Vec<String>>>
390where
391 T: GStore,
392{
393 match table_factor {
394 TableFactor::Table { name, alias, .. } => {
395 let columns = fetch_columns(storage, name).await?;
396 match (columns, alias) {
397 (columns, None) => Ok(columns),
398 (None, Some(_)) => Ok(None),
399 (Some(columns), Some(alias)) if alias.columns.len() > columns.len() => {
400 Err(FetchError::TooManyColumnAliases(
401 name.to_string(),
402 columns.len(),
403 alias.columns.len(),
404 )
405 .into())
406 }
407 (Some(columns), Some(alias)) => Ok(Some(
408 alias
409 .columns
410 .iter()
411 .cloned()
412 .chain(columns[alias.columns.len()..columns.len()].to_vec())
413 .collect(),
414 )),
415 }
416 }
417 TableFactor::Series { .. } => Ok(Some(vec!["N".to_owned()])),
418 TableFactor::Dictionary { dict, .. } => Ok(Some(match dict {
419 Dictionary::GlueObjects => vec![
420 "OBJECT_NAME".to_owned(),
421 "OBJECT_TYPE".to_owned(),
422 "CREATED".to_owned(),
423 ],
424 Dictionary::GlueTables => vec!["TABLE_NAME".to_owned(), "COMMENT".to_owned()],
425 Dictionary::GlueTableColumns => vec![
426 "TABLE_NAME".to_owned(),
427 "COLUMN_NAME".to_owned(),
428 "COLUMN_ID".to_owned(),
429 "NULLABLE".to_owned(),
430 "KEY".to_owned(),
431 "DEFAULT".to_owned(),
432 "COMMENT".to_owned(),
433 ],
434 Dictionary::GlueIndexes => vec![
435 "TABLE_NAME".to_owned(),
436 "INDEX_NAME".to_owned(),
437 "ORDER".to_owned(),
438 "EXPRESSION".to_owned(),
439 "UNIQUENESS".to_owned(),
440 ],
441 })),
442 TableFactor::Derived {
443 subquery: Query { body, .. },
444 alias:
445 TableAlias {
446 columns: alias_columns,
447 name,
448 },
449 } => match body {
450 SetExpr::Select(statement) => {
451 let Select {
452 from:
453 TableWithJoins {
454 relation, joins, ..
455 },
456 projection,
457 ..
458 } = statement.as_ref();
459
460 let labels = fetch_labels(storage, relation, joins, projection).await?;
461 match labels {
462 None => Ok(None),
463 Some(labels) if alias_columns.is_empty() => Ok(Some(labels)),
464 Some(labels) if alias_columns.len() > labels.len() => {
465 Err(FetchError::TooManyColumnAliases(
466 name.to_string(),
467 labels.len(),
468 alias_columns.len(),
469 )
470 .into())
471 }
472 Some(labels) => Ok(Some(
473 alias_columns
474 .iter()
475 .cloned()
476 .chain(labels[alias_columns.len()..labels.len()].to_vec())
477 .collect(),
478 )),
479 }
480 }
481 SetExpr::Values(Values(values_list)) => {
482 let total_len = values_list[0].len();
483 let alias_len = alias_columns.len();
484 if alias_len > total_len {
485 return Err(FetchError::TooManyColumnAliases(
486 name.into(),
487 total_len,
488 alias_len,
489 )
490 .into());
491 }
492 let labels = (alias_len + 1..=total_len).map(|i| format!("column{}", i));
493 let labels = alias_columns
494 .iter()
495 .cloned()
496 .chain(labels)
497 .collect::<Vec<_>>();
498
499 Ok(Some(labels))
500 }
501 },
502 }
503}
504
505async fn fetch_join_columns<'a, T: GStore>(
506 storage: &T,
507 joins: &'a [Join],
508) -> Result<Option<Vec<(&'a String, Vec<String>)>>> {
509 let columns = stream::iter(joins)
510 .filter_map(|join| async {
511 let relation = &join.relation;
512 let alias = get_alias(relation);
513
514 fetch_relation_columns(storage, relation)
515 .await
516 .map(|columns| Some((alias, columns?)))
517 .transpose()
518 })
519 .try_collect::<Vec<_>>()
520 .await?;
521
522 Ok((columns.len() == joins.len()).then_some(columns))
523}
524
525pub async fn fetch_labels<T: GStore>(
526 storage: &T,
527 relation: &TableFactor,
528 joins: &[Join],
529 projection: &[SelectItem],
530) -> Result<Option<Vec<String>>> {
531 let table_alias = get_alias(relation);
532 let columns = fetch_relation_columns(storage, relation).await?;
533 let join_columns = fetch_join_columns(storage, joins).await?;
534
535 if (columns.is_none() || join_columns.is_none())
536 && projection.iter().any(|item| {
537 matches!(
538 item,
539 SelectItem::Wildcard | SelectItem::QualifiedWildcard(_)
540 )
541 })
542 {
543 return Ok(None);
544 }
545
546 let columns = columns.unwrap_or_default();
547 let join_columns = join_columns.unwrap_or_default();
548
549 projection
550 .iter()
551 .flat_map(|item| match item {
552 SelectItem::Wildcard => {
553 let columns = columns.iter().cloned();
554 let join_columns = join_columns.iter().flat_map(|(_, columns)| columns.clone());
555
556 columns.chain(join_columns).map(Ok).collect()
557 }
558 SelectItem::QualifiedWildcard(target_table_alias) => {
559 if table_alias == target_table_alias {
560 return columns.iter().cloned().map(Ok).collect();
561 }
562
563 let labels = join_columns
564 .iter()
565 .find(|(table_alias, _)| table_alias == &target_table_alias)
566 .map(|(_, columns)| columns.clone());
567
568 match labels {
569 Some(columns) => columns.into_iter().map(Ok).collect(),
570 None => {
571 vec![Err(FetchError::TableAliasNotFound(
572 target_table_alias.to_owned(),
573 )
574 .into())]
575 }
576 }
577 }
578 SelectItem::Expr { label, .. } => vec![Ok(label.to_owned())],
579 })
580 .collect::<Result<_>>()
581 .map(Some)
582}