Skip to content

Commit c32cd87

Browse files
i-am-tomGil Mizrahi
i-am-tom
authored and
Gil Mizrahi
committed
It lives!
1 parent aac59d3 commit c32cd87

17 files changed

+557
-251
lines changed

‎crates/query-engine/sql/src/sql/helpers.rs‎

Lines changed: 43 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -401,15 +401,22 @@ pub fn select_mutation_rowset(
401401
(output_table_alias, output_column_alias): (TableAlias, ColumnAlias),
402402
(row_table_alias, row_column_alias): (TableAlias, ColumnAlias),
403403
aggregate_table_alias: TableAlias,
404-
row_select: Select,
405-
aggregate_select: Option<Select>,
404+
select: SelectSet,
406405
) -> Select {
407406
let row = vec![(
408407
output_column_alias,
409408
Expression::JsonBuildObject(BTreeMap::from([
410-
("type".to_string(), Box::new(Expression::Value(Value::String("procedure".to_string())))),
411-
("result".to_string(), Box::new(Expression::RowToJson(TableReference::AliasedTable(output_table_alias.clone()))))
412-
]))
409+
(
410+
"type".to_string(),
411+
Box::new(Expression::Value(Value::String("procedure".to_string()))),
412+
),
413+
(
414+
"result".to_string(),
415+
Box::new(Expression::RowToJson(TableReference::AliasedTable(
416+
output_table_alias.clone(),
417+
))),
418+
),
419+
])),
413420
)];
414421

415422
let mut final_select = simple_select(row);
@@ -418,25 +425,43 @@ pub fn select_mutation_rowset(
418425
select_rows_as_json_for_mutation(row_sel, row_column_alias, row_table_alias.clone())
419426
};
420427

421-
let mut select_star = star_select(From::Select {
422-
alias: row_table_alias.clone(),
423-
select: Box::new(wrap_row(row_select)),
424-
});
428+
match select {
429+
SelectSet::Rows(row_select) => {
430+
let select_star = star_select(From::Select {
431+
alias: row_table_alias.clone(),
432+
select: Box::new(wrap_row(row_select)),
433+
});
434+
435+
final_select.from = Some(From::Select {
436+
alias: output_table_alias,
437+
select: Box::new(select_star),
438+
});
439+
}
440+
441+
SelectSet::Aggregates(aggregate_select) => {
442+
final_select.from = Some(From::Select {
443+
alias: output_table_alias,
444+
select: Box::new(aggregate_select),
445+
});
446+
}
447+
448+
SelectSet::RowsAndAggregates(row_select, aggregate_select) => {
449+
let mut select_star = star_select(From::Select {
450+
alias: row_table_alias.clone(),
451+
select: Box::new(wrap_row(row_select)),
452+
});
425453

426-
match aggregate_select {
427-
None => {}
428-
Some(aggregate_select) => {
429454
select_star.joins = vec![Join::CrossJoin(CrossJoin {
430455
select: Box::new(aggregate_select),
431456
alias: aggregate_table_alias.clone(),
432457
})];
433-
}
434-
}
435458

436-
final_select.from = Some(From::Select {
437-
alias: output_table_alias,
438-
select: Box::new(select_star),
439-
});
459+
final_select.from = Some(From::Select {
460+
alias: output_table_alias,
461+
select: Box::new(select_star),
462+
});
463+
}
464+
};
440465

441466
final_select
442467
}

‎crates/query-engine/translation/src/translation/error.rs‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -28,6 +28,8 @@ pub enum Error {
2828
ColumnIsIdentityAlways(String),
2929
MissingColumnInInsert(String, String),
3030
NotImplementedYet(String),
31+
NoProcedureResultFieldsRequested,
32+
UnexpectedStructure(String),
3133
InternalError(String),
3234
}
3335

@@ -119,6 +121,11 @@ impl std::fmt::Display for Error {
119121
Error::NotImplementedYet(thing) => {
120122
write!(f, "Queries containing {} are not supported.", thing)
121123
}
124+
Error::NoProcedureResultFieldsRequested => write!(
125+
f,
126+
"Procedure requests must ask for 'affected_rows' or use the 'returning' clause."
127+
),
128+
Error::UnexpectedStructure(structure) => write!(f, "Unexpected {}.", structure),
122129
Error::InternalError(thing) => {
123130
write!(f, "Internal error: {}.", thing)
124131
}

‎crates/query-engine/translation/src/translation/mutation/translate.rs‎

Lines changed: 122 additions & 28 deletions
Original file line numberDiff line numberDiff line change
@@ -64,16 +64,49 @@ fn translate_mutation(
6464
alias: table_alias.clone(),
6565
};
6666

67-
let fields = match fields {
68-
Some(models::NestedField::Object(models::NestedObject{ fields })) => Ok(Some(fields)),
69-
Some(models::NestedField::Array(_)) => Err(Error::NotImplementedYet("nested array fields".to_string())),
70-
None => Ok(None)
67+
let (aggregates, fields) = match fields {
68+
Some(models::NestedField::Object(models::NestedObject { fields })) => {
69+
let affected_rows = fields
70+
.get("affected_rows")
71+
.map(|_| indexmap!("affected_rows".to_string() => models::Aggregate::StarCount{}));
72+
73+
let returning = match fields.get("returning") {
74+
Some(field) => match field {
75+
models::Field::Column { fields, .. } => match fields {
76+
Some(nested_fields) => match nested_fields {
77+
models::NestedField::Object(models::NestedObject { fields }) => {
78+
Some(fields.clone())
79+
}
80+
models::NestedField::Array(_) => {
81+
Err(Error::NotImplementedYet("nested array fields".to_string()))?
82+
}
83+
},
84+
None => None,
85+
},
86+
models::Field::Relationship { .. } => Err(Error::UnexpectedStructure(
87+
"relationship variant in procedure fields".to_string(),
88+
))?,
89+
},
90+
None => None,
91+
};
92+
93+
if affected_rows.is_none() && returning.is_none() {
94+
Err(Error::NoProcedureResultFieldsRequested)?
95+
}
96+
97+
Ok((affected_rows, returning))
98+
}
99+
100+
Some(models::NestedField::Array(_)) => {
101+
Err(Error::NotImplementedYet("nested array fields".to_string()))
102+
}
103+
None => Err(Error::NoProcedureResultFieldsRequested)?,
71104
}?;
72105

73106
// define the query selecting from the native query,
74107
// selecting the affected_rows as aggregate and the fields.
75108
let query = models::Query {
76-
aggregates: Some(indexmap!("affected_rows".to_string() => models::Aggregate::StarCount{})),
109+
aggregates,
77110
fields,
78111
limit: None,
79112
offset: None,
@@ -113,13 +146,21 @@ fn translate_mutation(
113146
};
114147

115148
// fields
116-
let (_, returning_select) = crate::translation::query::root::translate_rows_query(
117-
env,
118-
&mut state,
119-
&current_table,
120-
&from_clause,
121-
&query,
122-
)?;
149+
let returning_select = match &query.fields {
150+
Some(_) => {
151+
let (_, returning_select) = crate::translation::query::root::translate_rows_query(
152+
env,
153+
&mut state,
154+
&current_table,
155+
&from_clause,
156+
&query,
157+
)?;
158+
159+
Some(returning_select)
160+
}
161+
162+
None => None,
163+
};
123164

124165
// affected rows
125166
let aggregate_select = crate::translation::query::root::translate_aggregate_query(
@@ -141,8 +182,14 @@ fn translate_mutation(
141182
sql::helpers::make_column_alias("returning".to_string()),
142183
),
143184
state.make_table_alias("aggregates".to_string()),
144-
returning_select,
145-
aggregate_select,
185+
match (returning_select, aggregate_select) {
186+
(Some(returning_select), None) => sql::helpers::SelectSet::Rows(returning_select),
187+
(None, Some(aggregate_select)) => sql::helpers::SelectSet::Aggregates(aggregate_select),
188+
(Some(returning_select), Some(aggregate_select)) => {
189+
sql::helpers::SelectSet::RowsAndAggregates(returning_select, aggregate_select)
190+
}
191+
(None, None) => Err(Error::NoProcedureResultFieldsRequested)?,
192+
},
146193
);
147194

148195
let common_table_expression = sql::ast::CommonTableExpression {
@@ -200,16 +247,49 @@ fn translate_native_query(
200247
alias: table_alias.clone(),
201248
};
202249

203-
let fields = match fields {
204-
Some(models::NestedField::Object(models::NestedObject{ fields })) => Ok(Some(fields)),
205-
Some(models::NestedField::Array(_)) => Err(Error::NotImplementedYet("nested array fields".to_string())),
206-
None => Ok(None)
250+
let (aggregates, fields) = match fields {
251+
Some(models::NestedField::Object(models::NestedObject { fields })) => {
252+
let affected_rows = fields
253+
.get("affected_rows")
254+
.map(|_| indexmap!("affected_rows".to_string() => models::Aggregate::StarCount{}));
255+
256+
let returning = match fields.get("returning") {
257+
Some(field) => match field {
258+
models::Field::Column { fields, .. } => match fields {
259+
Some(nested_fields) => match nested_fields {
260+
models::NestedField::Object(models::NestedObject { fields }) => {
261+
Some(fields.clone())
262+
}
263+
models::NestedField::Array(_) => {
264+
Err(Error::NotImplementedYet("nested array fields".to_string()))?
265+
}
266+
},
267+
None => None,
268+
},
269+
models::Field::Relationship { .. } => Err(Error::UnexpectedStructure(
270+
"relationship variant in procedure fields".to_string(),
271+
))?,
272+
},
273+
None => None,
274+
};
275+
276+
if affected_rows.is_none() && returning.is_none() {
277+
Err(Error::NoProcedureResultFieldsRequested)?
278+
}
279+
280+
Ok((affected_rows, returning))
281+
}
282+
283+
Some(models::NestedField::Array(_)) => {
284+
Err(Error::NotImplementedYet("nested array fields".to_string()))
285+
}
286+
None => Err(Error::NoProcedureResultFieldsRequested)?,
207287
}?;
208288

209289
// define the query selecting from the native query,
210290
// selecting the affected_rows as aggregate and the fields.
211291
let query = models::Query {
212-
aggregates: Some(indexmap!("affected_rows".to_string() => models::Aggregate::StarCount{})),
292+
aggregates,
213293
fields,
214294
limit: None,
215295
offset: None,
@@ -223,13 +303,21 @@ fn translate_native_query(
223303
};
224304

225305
// fields
226-
let (_, returning_select) = crate::translation::query::root::translate_rows_query(
227-
env,
228-
&mut state,
229-
&current_table,
230-
&from_clause,
231-
&query,
232-
)?;
306+
let returning_select = match &query.fields {
307+
Some(_) => {
308+
let (_, returning_select) = crate::translation::query::root::translate_rows_query(
309+
env,
310+
&mut state,
311+
&current_table,
312+
&from_clause,
313+
&query,
314+
)?;
315+
316+
Some(returning_select)
317+
}
318+
319+
None => None,
320+
};
233321

234322
// affected rows
235323
let aggregate_select = crate::translation::query::root::translate_aggregate_query(
@@ -251,8 +339,14 @@ fn translate_native_query(
251339
sql::helpers::make_column_alias("returning".to_string()),
252340
),
253341
state.make_table_alias("aggregates".to_string()),
254-
returning_select,
255-
aggregate_select,
342+
match (returning_select, aggregate_select) {
343+
(Some(returning_select), None) => sql::helpers::SelectSet::Rows(returning_select),
344+
(None, Some(aggregate_select)) => sql::helpers::SelectSet::Aggregates(aggregate_select),
345+
(Some(returning_select), Some(aggregate_select)) => {
346+
sql::helpers::SelectSet::RowsAndAggregates(returning_select, aggregate_select)
347+
}
348+
(None, None) => Err(Error::NoProcedureResultFieldsRequested)?,
349+
},
256350
);
257351

258352
// add the procedure native query definition is a with clause.

‎crates/query-engine/translation/tests/goldenfiles/mutations/simple/request.json‎

Lines changed: 20 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -7,9 +7,26 @@
77
"track_id": 90
88
},
99
"fields": {
10-
"playlist_id": {
11-
"type": "column",
12-
"column": "PlaylistId"
10+
"type": "object",
11+
"fields": {
12+
"affected_rows": {
13+
"column": "affected_rows",
14+
"type": "column"
15+
},
16+
17+
"returning": {
18+
"type": "column",
19+
"column": "returning",
20+
"fields": {
21+
"type": "object",
22+
"fields": {
23+
"playlist_id": {
24+
"type": "column",
25+
"column": "PlaylistId"
26+
}
27+
}
28+
}
29+
}
1330
}
1431
}
1532
}

‎crates/query-engine/translation/tests/goldenfiles/mutations/v1_insert/request.json‎

Lines changed: 24 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -10,13 +10,30 @@
1010
}
1111
},
1212
"fields": {
13-
"artist_id": {
14-
"type": "column",
15-
"column": "id"
16-
},
17-
"name": {
18-
"type": "column",
19-
"column": "name"
13+
"type": "object",
14+
"fields": {
15+
"affected_rows": {
16+
"column": "affected_rows",
17+
"type": "column"
18+
},
19+
20+
"returning": {
21+
"type": "column",
22+
"column": "returning",
23+
"fields": {
24+
"type": "object",
25+
"fields": {
26+
"artist_id": {
27+
"type": "column",
28+
"column": "id"
29+
},
30+
"name": {
31+
"type": "column",
32+
"column": "name"
33+
}
34+
}
35+
}
36+
}
2037
}
2138
}
2239
}

0 commit comments

Comments
 (0)