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(crate) async fn enhance_relations_internal(
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(crate) async fn enhance_query_relations_internal(
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(crate) fn enhance_relation_aggregates_internal<'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(crate) fn enhance_object_group_bys_internal<'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_internal(query.entity.clone())
116 .with_trace_context(parent_trace_chain.to_vec())
117 .fetch_all_internal(&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(crate) fn enhance_child_queries_internal<'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_internal(query.entity.clone())
163 .with_trace_context(parent_trace_chain.to_vec())
164 .fetch_all_internal(&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_internal(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_internal(&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_internal(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_internal(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_internal(plan.target_entity.clone());
379 let query = self.query_for_plan(plan, parent_rows);
380 let mut child_rows = child_repo.fetch_all_internal(&query).await?;
381 for child in &mut child_rows {
382 child.remove(teaql_core::PARTITION_RANK_PROPERTY);
383 }
384 self.attach_relation_rows(parent_rows, plan, child_rows);
385
386 if !plan.children.is_empty() {
387 for parent in parent_rows.iter_mut() {
388 match parent.get_mut(&plan.relation_name) {
389 Some(Value::Object(child)) => {
390 child_repo
391 .enhance_child_record(child, &plan.children)
392 .await?;
393 }
394 Some(Value::List(values)) => {
395 for value in values.iter_mut() {
396 if let Value::Object(child) = value {
397 child_repo
398 .enhance_child_record(child, &plan.children)
399 .await?;
400 }
401 }
402 }
403 _ => {}
404 }
405 }
406 }
407 Ok(())
408 })
409 }
410
411 fn enhance_child_record<'b>(
412 &'b self,
413 child: &'b mut Record,
414 plans: &'b [RelationLoadPlan],
415 ) -> std::pin::Pin<
416 Box<dyn std::future::Future<Output = Result<(), DataServiceError<E::Error>>> + Send + 'b>,
417 > {
418 Box::pin(async move {
419 for plan in plans {
420 self.enhance_plan(slice::from_mut(child), plan).await?;
421 }
422 Ok(())
423 })
424 }
425
426 fn query_for_plan(&self, plan: &RelationLoadPlan, parent_rows: &[Record]) -> SelectQuery {
427 let ids = parent_rows
428 .iter()
429 .filter_map(|row| row.get(&plan.local_key).cloned())
430 .collect::<Vec<_>>();
431
432 let mut query = plan
433 .query
434 .clone()
435 .unwrap_or_else(|| SelectQuery::new(plan.target_entity.clone()));
436 query.entity = plan.target_entity.clone();
437 ensure_projection(&mut query, &plan.foreign_key);
438 for child in &plan.children {
439 ensure_projection(&mut query, &child.local_key);
440 }
441 if !ids.is_empty() {
442 query = query.and_filter(Expr::in_list(plan.foreign_key.clone(), ids));
443 }
444 if query.slice.is_some() {
445 query.partition_by = Some(plan.foreign_key.clone());
446 }
447 query
448 }
449
450 fn attach_relation_rows(
451 &self,
452 parent_rows: &mut [Record],
453 plan: &RelationLoadPlan,
454 child_rows: Vec<Record>,
455 ) {
456 let inverse_relation = self
457 .data_service
458 .metadata
459 .context
460 .entity(&plan.target_entity)
461 .and_then(|descriptor| {
462 descriptor.relations.iter().find(|relation| {
463 relation.target_entity == plan.parent_entity
464 && relation.local_key == plan.foreign_key
465 && relation.foreign_key == plan.local_key
466 })
467 })
468 .map(|relation| (relation.name.clone(), relation.many));
469
470 let mut buckets: BTreeMap<String, Vec<Record>> = BTreeMap::new();
471 for child in child_rows.clone() {
472 if let Some(key) = child.get(&plan.foreign_key) {
473 buckets
474 .entry(graph_identity_key(key))
475 .or_default()
476 .push(child);
477 }
478 }
479
480 for parent in parent_rows.iter_mut() {
481 let Some(local_value) = parent.get(&plan.local_key) else {
482 continue;
483 };
484 let bucket_key = graph_identity_key(local_value);
485 let related = buckets.get(&bucket_key).cloned().unwrap_or_default();
486 let related = match &inverse_relation {
487 Some((inverse_relation, inverse_many)) => {
488 let mut parent_object = parent.clone();
489 parent_object.remove(&plan.relation_name);
490 related
491 .into_iter()
492 .map(|mut child| {
493 match *inverse_many {
494 true => {
495 let entry = child
496 .entry(inverse_relation.clone())
497 .or_insert_with(|| Value::List(Vec::new()));
498 if let Value::List(list) = entry {
499 list.push(Value::object(parent_object.clone()));
500 }
501 }
502 false => {
503 child.insert(
504 inverse_relation.clone(),
505 Value::object(parent_object.clone()),
506 );
507 }
508 }
509 child
510 })
511 .collect::<Vec<_>>()
512 }
513 None => related,
514 };
515 match plan.many {
516 true => {
517 parent.insert(
518 plan.relation_name.clone(),
519 Value::List(related.into_iter().map(Value::object).collect()),
520 );
521 }
522 false => {
523 let value = related
524 .into_iter()
525 .next()
526 .map(Value::object)
527 .unwrap_or(Value::Null);
528 parent.insert(plan.relation_name.clone(), value);
529 }
530 }
531 }
532 }
533}