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 = if aggregate.single_result {
236 aggregate.alias.clone()
237 } else {
238 "count".to_owned()
239 };
240 query = query.aggregate(Aggregate::count(alias));
241 }
242 if !query
243 .group_by
244 .iter()
245 .any(|field| field == &plan.foreign_key)
246 {
247 query = query.group_by(plan.foreign_key.clone());
248 }
249 query = query.and_filter(Expr::in_list(plan.foreign_key.clone(), ids));
250
251 let mut chain = parent_trace_chain.to_vec();
252 chain.push(teaql_core::TraceNode {
253 entity_type: query.entity.clone(),
254 entity_id: None,
255 comment: aggregate.alias.clone(),
256 });
257
258 let mut aggregate_rows = child_repo
259 .with_trace_context(chain)
260 .fetch_all(&query)
261 .await?;
262 let foreign_key_column = self
263 .data_service
264 .metadata
265 .context
266 .entity(&plan.target_entity)
267 .and_then(|descriptor| {
268 descriptor
269 .properties
270 .iter()
271 .find(|property| property.name == plan.foreign_key)
272 .map(|property| property.column_name.clone())
273 });
274 if let Some(foreign_key_column) =
275 foreign_key_column.filter(|column| column != &plan.foreign_key)
276 {
277 for row in &mut aggregate_rows {
278 if !row.contains_key(&plan.foreign_key) {
279 if let Some(value) = row.remove(&foreign_key_column) {
280 row.insert(plan.foreign_key.clone(), value);
281 }
282 }
283 }
284 }
285 attach_relation_aggregate_rows(parent_rows, &plan, aggregate, aggregate_rows);
286 Ok(())
287 }
288
289 fn build_relation_plans(
290 &self,
291 entity: &str,
292 loads: &[String],
293 ) -> Result<Vec<RelationLoadPlan>, RuntimeError> {
294 let descriptor = self.data_service.metadata.context.require_entity(entity)?;
295 let mut grouped: BTreeMap<String, Vec<String>> = BTreeMap::new();
296 for load in loads {
297 if let Some((head, tail)) = load.split_once('.') {
298 grouped
299 .entry(head.to_owned())
300 .or_default()
301 .push(tail.to_owned());
302 } else {
303 grouped.entry(load.clone()).or_default();
304 }
305 }
306
307 grouped
308 .into_iter()
309 .map(|(name, child_loads)| {
310 let relation = descriptor.relation_by_name(&name).ok_or_else(|| {
311 RuntimeError::MissingRelation {
312 entity: entity.to_owned(),
313 relation: name.clone(),
314 }
315 })?;
316 let child_repo = self.scoped_data_service(relation.target_entity.clone());
317 let children =
318 child_repo.build_relation_plans(&relation.target_entity, &child_loads)?;
319 Ok(RelationLoadPlan {
320 parent_entity: entity.to_owned(),
321 relation_name: relation.name.clone(),
322 path: relation.name.clone(),
323 target_entity: relation.target_entity.clone(),
324 local_key: relation.local_key.clone(),
325 foreign_key: relation.foreign_key.clone(),
326 many: relation.many,
327 query: None,
328 children,
329 })
330 })
331 .collect()
332 }
333
334 fn build_relation_plans_from_loads(
335 &self,
336 entity: &str,
337 loads: &[RelationLoad],
338 ) -> Result<Vec<RelationLoadPlan>, RuntimeError> {
339 let descriptor = self.data_service.metadata.context.require_entity(entity)?;
340 loads
341 .iter()
342 .map(|load| {
343 let relation = descriptor.relation_by_name(&load.name).ok_or_else(|| {
344 RuntimeError::MissingRelation {
345 entity: entity.to_owned(),
346 relation: load.name.clone(),
347 }
348 })?;
349 let relation_query = load.query.as_deref().cloned();
350 let child_loads = relation_query
351 .as_ref()
352 .map(|query| query.relations.as_slice())
353 .unwrap_or_default();
354 let child_repo = self.scoped_data_service(relation.target_entity.clone());
355 let children = child_repo
356 .build_relation_plans_from_loads(&relation.target_entity, child_loads)?;
357 Ok(RelationLoadPlan {
358 parent_entity: entity.to_owned(),
359 relation_name: relation.name.clone(),
360 path: relation.name.clone(),
361 target_entity: relation.target_entity.clone(),
362 local_key: relation.local_key.clone(),
363 foreign_key: relation.foreign_key.clone(),
364 many: relation.many,
365 query: relation_query,
366 children,
367 })
368 })
369 .collect()
370 }
371 fn enhance_plan<'b>(
372 &'b self,
373 parent_rows: &'b mut [Record],
374 plan: &'b RelationLoadPlan,
375 ) -> std::pin::Pin<
376 Box<dyn std::future::Future<Output = Result<(), DataServiceError<E::Error>>> + Send + 'b>,
377 > {
378 Box::pin(async move {
379 let child_repo = self.scoped_data_service(plan.target_entity.clone());
380 let query = self.query_for_plan(plan, parent_rows);
381 let child_rows = child_repo.fetch_all(&query).await?;
382 self.attach_relation_rows(parent_rows, plan, child_rows);
383
384 if !plan.children.is_empty() {
385 for parent in parent_rows.iter_mut() {
386 match parent.get_mut(&plan.relation_name) {
387 Some(Value::Object(child)) => {
388 child_repo
389 .enhance_child_record(child, &plan.children)
390 .await?;
391 }
392 Some(Value::List(values)) => {
393 for value in values.iter_mut() {
394 if let Value::Object(child) = value {
395 child_repo
396 .enhance_child_record(child, &plan.children)
397 .await?;
398 }
399 }
400 }
401 _ => {}
402 }
403 }
404 }
405 Ok(())
406 })
407 }
408
409 fn enhance_child_record<'b>(
410 &'b self,
411 child: &'b mut Record,
412 plans: &'b [RelationLoadPlan],
413 ) -> std::pin::Pin<
414 Box<dyn std::future::Future<Output = Result<(), DataServiceError<E::Error>>> + Send + 'b>,
415 > {
416 Box::pin(async move {
417 for plan in plans {
418 self.enhance_plan(slice::from_mut(child), plan).await?;
419 }
420 Ok(())
421 })
422 }
423
424 fn query_for_plan(&self, plan: &RelationLoadPlan, parent_rows: &[Record]) -> SelectQuery {
425 let ids = parent_rows
426 .iter()
427 .filter_map(|row| row.get(&plan.local_key).cloned())
428 .collect::<Vec<_>>();
429
430 let mut query = plan
431 .query
432 .clone()
433 .unwrap_or_else(|| SelectQuery::new(plan.target_entity.clone()));
434 query.entity = plan.target_entity.clone();
435 ensure_projection(&mut query, &plan.foreign_key);
436 for child in &plan.children {
437 ensure_projection(&mut query, &child.local_key);
438 }
439 if !ids.is_empty() {
440 query = query.and_filter(Expr::in_list(plan.foreign_key.clone(), ids));
441 }
442 query
443 }
444
445 fn attach_relation_rows(
446 &self,
447 parent_rows: &mut [Record],
448 plan: &RelationLoadPlan,
449 child_rows: Vec<Record>,
450 ) {
451 let inverse_relation = self
452 .data_service
453 .metadata
454 .context
455 .entity(&plan.target_entity)
456 .and_then(|descriptor| {
457 descriptor.relations.iter().find(|relation| {
458 relation.target_entity == plan.parent_entity
459 && relation.local_key == plan.foreign_key
460 && relation.foreign_key == plan.local_key
461 })
462 })
463 .map(|relation| (relation.name.clone(), relation.many));
464
465 let mut buckets: BTreeMap<String, Vec<Record>> = BTreeMap::new();
466 for child in child_rows.clone() {
467 if let Some(key) = child.get(&plan.foreign_key) {
468 buckets
469 .entry(graph_identity_key(key))
470 .or_default()
471 .push(child);
472 }
473 }
474
475 for parent in parent_rows.iter_mut() {
476 let Some(local_value) = parent.get(&plan.local_key) else {
477 continue;
478 };
479 let bucket_key = graph_identity_key(local_value);
480 let related = buckets.get(&bucket_key).cloned().unwrap_or_default();
481 let related = if let Some((inverse_relation, inverse_many)) = &inverse_relation {
482 let mut parent_object = parent.clone();
483 parent_object.remove(&plan.relation_name);
484 related
485 .into_iter()
486 .map(|mut child| {
487 if *inverse_many {
488 let entry = child
489 .entry(inverse_relation.clone())
490 .or_insert_with(|| Value::List(Vec::new()));
491 if let Value::List(list) = entry {
492 list.push(Value::object(parent_object.clone()));
493 }
494 } else {
495 child.insert(
496 inverse_relation.clone(),
497 Value::object(parent_object.clone()),
498 );
499 }
500 child
501 })
502 .collect::<Vec<_>>()
503 } else {
504 related
505 };
506 if plan.many {
507 parent.insert(
508 plan.relation_name.clone(),
509 Value::List(related.into_iter().map(Value::object).collect()),
510 );
511 } else {
512 let value = related
513 .into_iter()
514 .next()
515 .map(Value::object)
516 .unwrap_or(Value::Null);
517 parent.insert(plan.relation_name.clone(), value);
518 }
519 }
520 }
521}