teaql_runtime/data_service/
relation.rs1use std::collections::BTreeMap;
2use std::slice;
3
4use teaql_core::{
5 Aggregate, Expr, ObjectGroupBy, Record, RelationAggregate, RelationLoad, SelectQuery, Value,
6};
7
8use crate::{DataServiceError, RuntimeError};
9
10use super::{EntityDataService, RelationLoadPlan, helpers::*};
11
12impl<'a, E> EntityDataService<'a, E>
13where
14 E: teaql_data_service::QueryExecutor
15 + teaql_data_service::MutationExecutor
16 + Send
17 + Sync
18 + 'static,
19{
20 pub fn relation_loads(&self) -> Vec<String> {
21 self.behavior()
22 .map(|behavior| behavior.relation_loads(self.data_service.metadata.context))
23 .unwrap_or_default()
24 }
25
26 pub fn relation_plans(&self) -> Result<Vec<RelationLoadPlan>, RuntimeError> {
27 self.build_relation_plans(&self.entity, &self.relation_loads())
28 }
29
30 pub fn relation_query(
31 &self,
32 relation_name: &str,
33 parent_rows: &[Record],
34 ) -> Result<SelectQuery, RuntimeError> {
35 let plan = self
36 .relation_plans()?
37 .into_iter()
38 .find(|plan| plan.relation_name == relation_name)
39 .ok_or_else(|| RuntimeError::MissingRelation {
40 entity: self.entity.clone(),
41 relation: relation_name.to_owned(),
42 })?;
43 Ok(self.query_for_plan(&plan, parent_rows))
44 }
45
46 pub async fn enhance_relations(
47 &self,
48 parent_rows: &mut [Record],
49 ) -> Result<(), DataServiceError<E::Error>> {
50 let plans = self.relation_plans().map_err(DataServiceError::Runtime)?;
51 for plan in plans {
52 self.enhance_plan(parent_rows, &plan).await?;
53 }
54 Ok(())
55 }
56
57 pub async fn enhance_query_relations(
58 &self,
59 parent_rows: &mut [Record],
60 query: &SelectQuery,
61 ) -> Result<(), DataServiceError<E::Error>> {
62 let plans = self
63 .build_relation_plans_from_loads(&query.entity, &query.relations)
64 .map_err(DataServiceError::Runtime)?;
65 for plan in plans {
66 self.enhance_plan(parent_rows, &plan).await?;
67 }
68 Ok(())
69 }
70
71 pub fn enhance_relation_aggregates<'b>(
72 &'b self,
73 parent_rows: &'b mut [Record],
74 relation_aggregates: &'b [RelationAggregate],
75 parent_cache_options: Option<teaql_core::AggregationCacheOptions>,
76 parent_trace_chain: &'b [teaql_core::TraceNode],
77 ) -> std::pin::Pin<
78 Box<dyn std::future::Future<Output = Result<(), DataServiceError<E::Error>>> + Send + 'b>,
79 > {
80 Box::pin(async move {
81 for aggregate in relation_aggregates {
82 self.enhance_relation_aggregate(
83 parent_rows,
84 aggregate,
85 parent_cache_options,
86 parent_trace_chain,
87 )
88 .await?;
89 }
90 Ok(())
91 })
92 }
93
94 pub fn enhance_object_group_bys<'b>(
95 &'b self,
96 rows: &'b mut [Record],
97 object_group_bys: &'b [ObjectGroupBy],
98 parent_trace_chain: &'b [teaql_core::TraceNode],
99 ) -> std::pin::Pin<
100 Box<dyn std::future::Future<Output = Result<(), DataServiceError<E::Error>>> + Send + 'b>,
101 > {
102 Box::pin(async move {
103 for group_by in object_group_bys {
104 let ids = rows
105 .iter()
106 .filter_map(|row| row.get(&group_by.storage_field).cloned())
107 .collect::<Vec<_>>();
108 if ids.is_empty() {
109 continue;
110 }
111 let mut query = group_by.query.clone();
112 ensure_projection(&mut query, "id");
113 query = query.and_filter(Expr::in_list("id", ids));
114 let object_rows = self
115 .scoped_data_service(query.entity.clone())
116 .with_trace_context(parent_trace_chain.to_vec())
117 .fetch_all(&query)
118 .await?
119 .into_iter()
120 .filter_map(|row| {
121 row.get("id")
122 .cloned()
123 .map(|id| (graph_identity_key(&id), row))
124 })
125 .collect::<BTreeMap<_, _>>();
126 for row in rows.iter_mut() {
127 if let Some(key) = row.get(&group_by.storage_field).map(graph_identity_key) {
128 let value = object_rows
129 .get(&key)
130 .cloned()
131 .map(Value::object)
132 .unwrap_or(Value::Null);
133 row.insert(group_by.property_name.clone(), value);
134 }
135 }
136 }
137 Ok(())
138 })
139 }
140
141 pub fn enhance_child_queries<'b>(
142 &'b self,
143 rows: &'b mut [Record],
144 child_queries: &'b [SelectQuery],
145 parent_trace_chain: &'b [teaql_core::TraceNode],
146 ) -> std::pin::Pin<
147 Box<dyn std::future::Future<Output = Result<(), DataServiceError<E::Error>>> + Send + 'b>,
148 > {
149 Box::pin(async move {
150 for child_query in child_queries {
151 let ids = rows
152 .iter()
153 .filter_map(|row| row.get("id").cloned())
154 .collect::<Vec<_>>();
155 if ids.is_empty() {
156 continue;
157 }
158 let mut query = child_query.clone();
159 ensure_projection(&mut query, "id");
160 query = query.and_filter(Expr::in_list("id", ids));
161 let child_rows = self
162 .scoped_data_service(query.entity.clone())
163 .with_trace_context(parent_trace_chain.to_vec())
164 .fetch_all(&query)
165 .await?
166 .into_iter()
167 .filter_map(|row| {
168 row.get("id")
169 .cloned()
170 .map(|id| (graph_identity_key(&id), row))
171 })
172 .collect::<BTreeMap<_, _>>();
173 for row in rows.iter_mut() {
174 if let Some(key) = row.get("id").map(graph_identity_key) {
175 if let Some(child) = child_rows.get(&key) {
176 row.extend(child.clone());
177 }
178 }
179 }
180 }
181 Ok(())
182 })
183 }
184
185 async fn enhance_relation_aggregate(
186 &self,
187 parent_rows: &mut [Record],
188 aggregate: &RelationAggregate,
189 parent_cache_options: Option<teaql_core::AggregationCacheOptions>,
190 parent_trace_chain: &[teaql_core::TraceNode],
191 ) -> Result<(), DataServiceError<E::Error>> {
192 let plan = self
193 .build_relation_plans_from_loads(
194 &self.entity,
195 &[RelationLoad::with_query(
196 aggregate.relation_name.clone(),
197 aggregate.query.clone(),
198 )],
199 )
200 .map_err(DataServiceError::Runtime)?
201 .into_iter()
202 .next()
203 .ok_or_else(|| {
204 DataServiceError::Runtime(RuntimeError::MissingRelation {
205 entity: self.entity.clone(),
206 relation: aggregate.relation_name.clone(),
207 })
208 })?;
209
210 let ids = parent_rows
211 .iter()
212 .filter_map(|row| row.get(&plan.local_key).cloned())
213 .collect::<Vec<_>>();
214 if ids.is_empty() {
215 attach_empty_relation_aggregate(parent_rows, &aggregate.alias, aggregate.single_result);
216 return Ok(());
217 }
218
219 let child_repo = self.scoped_data_service(plan.target_entity.clone());
220 let mut query = aggregate.query.clone();
221 query.entity = plan.target_entity.clone();
222 if query.aggregation_cache.is_none() {
223 if let Some(options) = parent_cache_options.filter(|options| options.propagate) {
224 query.aggregation_cache = Some(teaql_core::AggregationCacheOptions::enabled(
225 options.propagate_cache_expired_millis,
226 ));
227 }
228 }
229 query.projection.clear();
230 query.expr_projection.clear();
231 query.order_by.clear();
232 query.slice = None;
233 query.relations.clear();
234 if query.aggregates.is_empty() {
235 let alias = aggregate_alias(aggregate.single_result, &aggregate.alias);
236 query = query.aggregate(Aggregate::count(alias));
237 }
238 if !query
239 .group_by
240 .iter()
241 .any(|field| field == &plan.foreign_key)
242 {
243 query = query.group_by(plan.foreign_key.clone());
244 }
245 query = query.and_filter(Expr::in_list(plan.foreign_key.clone(), ids));
246
247 let mut chain = parent_trace_chain.to_vec();
248 chain.push(teaql_core::TraceNode {
249 entity_type: query.entity.clone(),
250 entity_id: None,
251 comment: aggregate.alias.clone(),
252 });
253
254 let mut aggregate_rows = child_repo
255 .with_trace_context(chain)
256 .fetch_all(&query)
257 .await?;
258 let foreign_key_column = self
259 .data_service
260 .metadata
261 .context
262 .entity(&plan.target_entity)
263 .and_then(|descriptor| {
264 descriptor
265 .properties
266 .iter()
267 .find(|property| property.name == plan.foreign_key)
268 .map(|property| property.column_name.clone())
269 });
270 if let Some(foreign_key_column) =
271 foreign_key_column.filter(|column| column != &plan.foreign_key)
272 {
273 for row in &mut aggregate_rows {
274 if !row.contains_key(&plan.foreign_key) {
275 if let Some(value) = row.remove(&foreign_key_column) {
276 row.insert(plan.foreign_key.clone(), value);
277 }
278 }
279 }
280 }
281 attach_relation_aggregate_rows(parent_rows, &plan, aggregate, aggregate_rows);
282 Ok(())
283 }
284
285 fn build_relation_plans(
286 &self,
287 entity: &str,
288 loads: &[String],
289 ) -> Result<Vec<RelationLoadPlan>, RuntimeError> {
290 let descriptor = self.data_service.metadata.context.require_entity(entity)?;
291 let mut grouped: BTreeMap<String, Vec<String>> = BTreeMap::new();
292 for load in loads {
293 match load.split_once('.') {
294 Some((head, tail)) => {
295 grouped
296 .entry(head.to_owned())
297 .or_default()
298 .push(tail.to_owned());
299 }
300 None => {
301 grouped.entry(load.clone()).or_default();
302 }
303 }
304 }
305
306 grouped
307 .into_iter()
308 .map(|(name, child_loads)| {
309 let relation = descriptor.relation_by_name(&name).ok_or_else(|| {
310 RuntimeError::MissingRelation {
311 entity: entity.to_owned(),
312 relation: name.clone(),
313 }
314 })?;
315 let child_repo = self.scoped_data_service(relation.target_entity.clone());
316 let children =
317 child_repo.build_relation_plans(&relation.target_entity, &child_loads)?;
318 Ok(RelationLoadPlan {
319 parent_entity: entity.to_owned(),
320 relation_name: relation.name.clone(),
321 path: relation.name.clone(),
322 target_entity: relation.target_entity.clone(),
323 local_key: relation.local_key.clone(),
324 foreign_key: relation.foreign_key.clone(),
325 many: relation.many,
326 query: None,
327 children,
328 })
329 })
330 .collect()
331 }
332
333 fn build_relation_plans_from_loads(
334 &self,
335 entity: &str,
336 loads: &[RelationLoad],
337 ) -> Result<Vec<RelationLoadPlan>, RuntimeError> {
338 let descriptor = self.data_service.metadata.context.require_entity(entity)?;
339 loads
340 .iter()
341 .map(|load| {
342 let relation = descriptor.relation_by_name(&load.name).ok_or_else(|| {
343 RuntimeError::MissingRelation {
344 entity: entity.to_owned(),
345 relation: load.name.clone(),
346 }
347 })?;
348 let relation_query = load.query.as_deref().cloned();
349 let child_loads = relation_query
350 .as_ref()
351 .map(|query| query.relations.as_slice())
352 .unwrap_or_default();
353 let child_repo = self.scoped_data_service(relation.target_entity.clone());
354 let children = child_repo
355 .build_relation_plans_from_loads(&relation.target_entity, child_loads)?;
356 Ok(RelationLoadPlan {
357 parent_entity: entity.to_owned(),
358 relation_name: relation.name.clone(),
359 path: relation.name.clone(),
360 target_entity: relation.target_entity.clone(),
361 local_key: relation.local_key.clone(),
362 foreign_key: relation.foreign_key.clone(),
363 many: relation.many,
364 query: relation_query,
365 children,
366 })
367 })
368 .collect()
369 }
370 fn enhance_plan<'b>(
371 &'b self,
372 parent_rows: &'b mut [Record],
373 plan: &'b RelationLoadPlan,
374 ) -> std::pin::Pin<
375 Box<dyn std::future::Future<Output = Result<(), DataServiceError<E::Error>>> + Send + 'b>,
376 > {
377 Box::pin(async move {
378 let child_repo = self.scoped_data_service(plan.target_entity.clone());
379 let query = self.query_for_plan(plan, parent_rows);
380 let child_rows = child_repo.fetch_all(&query).await?;
381 self.attach_relation_rows(parent_rows, plan, child_rows);
382
383 if !plan.children.is_empty() {
384 for parent in parent_rows.iter_mut() {
385 match parent.get_mut(&plan.relation_name) {
386 Some(Value::Object(child)) => {
387 child_repo
388 .enhance_child_record(child, &plan.children)
389 .await?;
390 }
391 Some(Value::List(values)) => {
392 for value in values.iter_mut() {
393 if let Value::Object(child) = value {
394 child_repo
395 .enhance_child_record(child, &plan.children)
396 .await?;
397 }
398 }
399 }
400 _ => {}
401 }
402 }
403 }
404 Ok(())
405 })
406 }
407
408 fn enhance_child_record<'b>(
409 &'b self,
410 child: &'b mut Record,
411 plans: &'b [RelationLoadPlan],
412 ) -> std::pin::Pin<
413 Box<dyn std::future::Future<Output = Result<(), DataServiceError<E::Error>>> + Send + 'b>,
414 > {
415 Box::pin(async move {
416 for plan in plans {
417 self.enhance_plan(slice::from_mut(child), plan).await?;
418 }
419 Ok(())
420 })
421 }
422
423 fn query_for_plan(&self, plan: &RelationLoadPlan, parent_rows: &[Record]) -> SelectQuery {
424 let ids = parent_rows
425 .iter()
426 .filter_map(|row| row.get(&plan.local_key).cloned())
427 .collect::<Vec<_>>();
428
429 let mut query = plan
430 .query
431 .clone()
432 .unwrap_or_else(|| SelectQuery::new(plan.target_entity.clone()));
433 query.entity = plan.target_entity.clone();
434 ensure_projection(&mut query, &plan.foreign_key);
435 for child in &plan.children {
436 ensure_projection(&mut query, &child.local_key);
437 }
438 if !ids.is_empty() {
439 query = query.and_filter(Expr::in_list(plan.foreign_key.clone(), ids));
440 }
441 query
442 }
443
444 fn attach_relation_rows(
445 &self,
446 parent_rows: &mut [Record],
447 plan: &RelationLoadPlan,
448 child_rows: Vec<Record>,
449 ) {
450 let inverse_relation = self
451 .data_service
452 .metadata
453 .context
454 .entity(&plan.target_entity)
455 .and_then(|descriptor| {
456 descriptor.relations.iter().find(|relation| {
457 relation.target_entity == plan.parent_entity
458 && relation.local_key == plan.foreign_key
459 && relation.foreign_key == plan.local_key
460 })
461 })
462 .map(|relation| (relation.name.clone(), relation.many));
463
464 let mut buckets: BTreeMap<String, Vec<Record>> = BTreeMap::new();
465 for child in child_rows.clone() {
466 if let Some(key) = child.get(&plan.foreign_key) {
467 buckets
468 .entry(graph_identity_key(key))
469 .or_default()
470 .push(child);
471 }
472 }
473
474 for parent in parent_rows.iter_mut() {
475 let Some(local_value) = parent.get(&plan.local_key) else {
476 continue;
477 };
478 let bucket_key = graph_identity_key(local_value);
479 let related = buckets.get(&bucket_key).cloned().unwrap_or_default();
480 let related = match &inverse_relation {
481 Some((inverse_relation, inverse_many)) => {
482 let mut parent_object = parent.clone();
483 parent_object.remove(&plan.relation_name);
484 related
485 .into_iter()
486 .map(|mut child| {
487 match *inverse_many {
488 true => {
489 let entry = child
490 .entry(inverse_relation.clone())
491 .or_insert_with(|| Value::List(Vec::new()));
492 if let Value::List(list) = entry {
493 list.push(Value::object(parent_object.clone()));
494 }
495 }
496 false => {
497 child.insert(
498 inverse_relation.clone(),
499 Value::object(parent_object.clone()),
500 );
501 }
502 }
503 child
504 })
505 .collect::<Vec<_>>()
506 }
507 None => related,
508 };
509 match plan.many {
510 true => {
511 parent.insert(
512 plan.relation_name.clone(),
513 Value::List(related.into_iter().map(Value::object).collect()),
514 );
515 }
516 false => {
517 let value = related
518 .into_iter()
519 .next()
520 .map(Value::object)
521 .unwrap_or(Value::Null);
522 parent.insert(plan.relation_name.clone(), value);
523 }
524 }
525 }
526 }
527}