use crate::{
error::TitoError,
types::{TitoEngine, TitoEmbeddedRelationshipConfig, TitoModelTrait},
TitoModel,
};
use serde::{de::DeserializeOwned, Serialize};
use serde_json::Value;
use std::collections::HashMap;
impl<
E: TitoEngine,
T: Default
+ Clone
+ Serialize
+ DeserializeOwned
+ Unpin
+ std::marker::Send
+ Sync
+ TitoModelTrait,
> TitoModel<E, T>
{
pub fn stitch_relationship(
&self,
item: &mut Value,
rel_map: &HashMap<String, Value>,
config: &TitoEmbeddedRelationshipConfig,
) {
let source_parts: Vec<&str> = config.source_field_name.split('.').collect();
let dest_parts: Vec<&str> = config.destination_field_name.split('.').collect();
Self::_stitch_recursive_helper(item, &source_parts, &dest_parts, rel_map, config);
}
fn _stitch_recursive_helper(
root_json_node: &mut Value,
source_path: &[&str],
dest_path: &[&str],
rel_map: &HashMap<String, Value>,
config: &TitoEmbeddedRelationshipConfig,
) {
Self::_stitch_iterative(root_json_node, source_path, dest_path, rel_map, config);
}
fn _stitch_iterative(
mut current_json_node: &mut Value,
mut source_path_remaining: &[&str],
mut dest_path_remaining: &[&str],
rel_map: &HashMap<String, Value>,
config: &TitoEmbeddedRelationshipConfig,
) {
loop {
if source_path_remaining.len() == 1 && dest_path_remaining.len() == 1 {
let source_key = source_path_remaining[0];
let dest_key = dest_path_remaining[0];
if let Some(obj_to_modify) = current_json_node.as_object_mut() {
if let Some(id_val_at_source_key) = obj_to_modify.get(source_key) {
if let Some(id_str) = id_val_at_source_key.as_str() {
let rel_lookup_key = format!("table:{}:{}", config.model, id_str);
if let Some(related_data) = rel_map.get(&rel_lookup_key) {
obj_to_modify.insert(dest_key.to_string(), related_data.clone());
}
} else if let Some(ids_array) = id_val_at_source_key.as_array() {
let mut stitched_related_items_array = Vec::new();
for id_elem_in_array in ids_array {
if let Some(id_str_elem) = id_elem_in_array.as_str() {
let rel_lookup_key =
format!("table:{}:{}", config.model, id_str_elem);
if let Some(related_data) = rel_map.get(&rel_lookup_key) {
stitched_related_items_array.push(related_data.clone());
}
}
}
obj_to_modify.insert(
dest_key.to_string(),
Value::Array(stitched_related_items_array),
);
}
}
}
break;
}
if !source_path_remaining.is_empty()
&& !dest_path_remaining.is_empty()
&& source_path_remaining[0] == dest_path_remaining[0]
{
let common_key = source_path_remaining[0];
if let Some(next_json_node_candidate) = current_json_node.get_mut(common_key) {
match next_json_node_candidate {
Value::Object(_) => {
current_json_node = next_json_node_candidate;
source_path_remaining = &source_path_remaining[1..];
dest_path_remaining = &dest_path_remaining[1..];
continue;
}
Value::Array(arr) => {
let remaining_source = &source_path_remaining[1..];
let remaining_dest = &dest_path_remaining[1..];
for element_in_array in arr.iter_mut() {
if remaining_source.len() == 1 && remaining_dest.len() == 1 {
let source_key = remaining_source[0];
let dest_key = remaining_dest[0];
if let Some(element_obj) = element_in_array.as_object_mut() {
if let Some(id_val) = element_obj.get(source_key) {
if let Some(id_str) = id_val.as_str() {
let rel_lookup_key = format!("table:{}:{}", config.model, id_str);
if let Some(related_data) = rel_map.get(&rel_lookup_key) {
element_obj.insert(dest_key.to_string(), related_data.clone());
}
}
}
}
}
}
break;
}
_ => break,
}
} else {
break;
}
} else {
break;
}
}
}
pub fn get_relationship_data(
&self,
items: &Vec<(String, Value)>,
rels_config: &[TitoEmbeddedRelationshipConfig],
rels: &Vec<String>, ) -> Vec<(TitoEmbeddedRelationshipConfig, String)> {
let mut relationship_keys_to_fetch = Vec::new();
for (_, item_value) in items {
for config in rels_config {
if rels.contains(&config.destination_field_name) {
if let Some(found_values_at_source_path) =
self.get_nested_values(item_value, &config.source_field_name)
{
for value_or_array_of_values in found_values_at_source_path {
if let Some(id_str) = value_or_array_of_values.as_str() {
if id_str != "__null__" {
relationship_keys_to_fetch.push((
config.clone(),
format!("table:{}:{}", config.model, id_str),
));
}
} else if let Some(id_array) = value_or_array_of_values.as_array() {
for id_element in id_array {
if let Some(id_str) = id_element.as_str() {
if id_str != "__null__" {
relationship_keys_to_fetch.push((
config.clone(),
format!("table:{}:{}", config.model, id_str),
));
}
}
}
}
}
}
}
}
}
relationship_keys_to_fetch
}
pub async fn fetch_and_stitch_relationships(
&self,
items: Vec<(String, Value)>,
rels: Vec<String>,
tx: &E::Transaction,
) -> Result<Vec<(String, Value)>, TitoError> {
if rels.is_empty() {
return Ok(items);
}
let rels_config = self.get_embedded_relationships();
let relationship_data = self.get_relationship_data(&items, &rels_config, &rels);
let rel_keys: Vec<String> = relationship_data.into_iter().map(|item| item.1).collect();
let rel_items = self.batch_get(rel_keys.clone(), 10, 2, tx).await?;
let mut rel_map: HashMap<String, Value> = HashMap::new();
for kv in rel_items {
rel_map.insert(kv.0, kv.1);
}
let final_items = items
.into_iter()
.map(|mut item| {
for config in &rels_config {
if rels.contains(&config.destination_field_name) {
self.stitch_relationship(&mut item.1, &rel_map, config);
}
}
item
})
.collect();
Ok(final_items)
}
pub fn extract_relationship(&self, input: &str) -> Option<String> {
if let (Some(start), Some(end)) = (input.find("-").map(|idx| idx + 1), input.rfind("-")) {
if start < end {
return Some(input[start..end].trim().to_string());
}
}
None
}
}