use std::collections::HashMap;
use std::hash::Hash;
use std::sync::RwLock;
#[derive(Debug, Clone)]
pub struct GraphEdge {
pub parent_field: String,
pub relation: String,
pub sub_graph: Option<Box<EntityGraph>>,
}
#[derive(Debug, Clone, Default)]
pub struct EntityGraph {
edges: Vec<GraphEdge>,
}
impl EntityGraph {
pub fn new() -> Self {
Self { edges: Vec::new() }
}
pub fn add_edge(
&mut self,
parent_field: impl Into<String>,
relation: impl Into<String>,
) -> &mut Self {
self.edges.push(GraphEdge {
parent_field: parent_field.into(),
relation: relation.into(),
sub_graph: None,
});
self
}
pub fn add_edge_with_graph(
&mut self,
parent_field: impl Into<String>,
relation: impl Into<String>,
sub_graph: EntityGraph,
) -> &mut Self {
self.edges.push(GraphEdge {
parent_field: parent_field.into(),
relation: relation.into(),
sub_graph: Some(Box::new(sub_graph)),
});
self
}
pub fn edges(&self) -> &[GraphEdge] {
&self.edges
}
pub fn edge_count(&self) -> usize {
self.edges.len()
}
pub fn relations_of(&self, parent_field: &str) -> Vec<&GraphEdge> {
self.edges
.iter()
.filter(|e| e.parent_field == parent_field)
.collect()
}
pub fn all_relations(&self) -> Vec<String> {
let mut rels: Vec<String> = self.edges.iter().map(|e| e.relation.clone()).collect();
rels.sort();
rels.dedup();
rels
}
pub fn all_parent_fields(&self) -> Vec<String> {
let mut fields: Vec<String> = self.edges.iter().map(|e| e.parent_field.clone()).collect();
fields.sort();
fields.dedup();
fields
}
pub fn is_empty(&self) -> bool {
self.edges.is_empty()
}
pub fn all_relations_recursive(&self) -> Vec<String> {
let mut result = Vec::new();
for edge in &self.edges {
result.push(edge.relation.clone());
if let Some(sub) = &edge.sub_graph {
result.extend(sub.all_relations_recursive());
}
}
result.sort();
result.dedup();
result
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Default)]
pub enum BatchStrategy {
#[default]
In,
Join,
Subquery,
}
impl BatchStrategy {
pub fn name(&self) -> &'static str {
match self {
BatchStrategy::In => "in",
BatchStrategy::Join => "join",
BatchStrategy::Subquery => "subquery",
}
}
pub fn render_in_clause(column: &str, placeholders: usize) -> String {
if placeholders == 0 {
return format!("{} IN ()", column);
}
let marks: Vec<&str> = vec!["?"; placeholders];
format!("{} IN ({})", column, marks.join(", "))
}
}
#[derive(Debug, Clone, Copy)]
pub struct BatchSizeConfig {
pub size: usize,
pub strategy: BatchStrategy,
}
impl Default for BatchSizeConfig {
fn default() -> Self {
Self {
size: 100,
strategy: BatchStrategy::In,
}
}
}
impl BatchSizeConfig {
pub fn new(size: usize, strategy: BatchStrategy) -> Self {
Self { size, strategy }
}
pub fn with_size(size: usize) -> Self {
Self {
size,
strategy: BatchStrategy::In,
}
}
pub fn batch_count(&self, total: usize) -> usize {
if total == 0 {
0
} else {
total.div_ceil(self.size)
}
}
pub fn batch_range(&self, batch_index: usize, total: usize) -> std::ops::Range<usize> {
let start = batch_index * self.size;
let end = (start + self.size).min(total);
start..end
}
}
pub type BatchLoaderFn<K, V> = Box<dyn Fn(&[K]) -> HashMap<K, V> + Send + Sync>;
pub struct BatchLoader<K, V>
where
K: Hash + Eq + Clone + Send + Sync,
V: Clone + Send + Sync,
{
batch_size: usize,
loader: BatchLoaderFn<K, V>,
cache: RwLock<HashMap<K, V>>,
}
impl<K, V> BatchLoader<K, V>
where
K: Hash + Eq + Clone + Send + Sync,
V: Clone + Send + Sync,
{
pub fn new(batch_size: usize, loader: BatchLoaderFn<K, V>) -> Self {
Self {
batch_size,
loader,
cache: RwLock::new(HashMap::new()),
}
}
pub fn load_many(&self, keys: &[K]) -> HashMap<K, V> {
let mut result: HashMap<K, V> = HashMap::new();
let cached = self.cache.read().unwrap();
let mut to_load: Vec<K> = Vec::new();
for k in keys {
if let Some(v) = cached.get(k) {
result.insert(k.clone(), v.clone());
} else {
to_load.push(k.clone());
}
}
drop(cached);
if to_load.is_empty() {
return result;
}
let batch_size = self.batch_size.max(1);
let mut all_loaded: HashMap<K, V> = HashMap::new();
for chunk in to_load.chunks(batch_size) {
let loaded = (self.loader)(chunk);
all_loaded.extend(loaded);
}
let mut cache = self.cache.write().unwrap();
for (k, v) in &all_loaded {
cache.insert(k.clone(), v.clone());
}
drop(cache);
result.extend(all_loaded);
result
}
pub fn load_one(&self, key: &K) -> Option<V> {
let result = self.load_many(std::slice::from_ref(key));
result.get(key).cloned()
}
pub fn clear_cache(&self) {
self.cache.write().unwrap().clear();
}
pub fn cache_size(&self) -> usize {
self.cache.read().unwrap().len()
}
pub fn batch_size(&self) -> usize {
self.batch_size
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_new_graph_is_empty() {
let g = EntityGraph::new();
assert!(g.is_empty());
assert_eq!(g.edge_count(), 0);
}
#[test]
fn test_add_edge() {
let mut g = EntityGraph::new();
g.add_edge("user", "posts");
assert_eq!(g.edge_count(), 1);
assert!(!g.is_empty());
}
#[test]
fn test_add_multiple_edges() {
let mut g = EntityGraph::new();
g.add_edge("user", "posts")
.add_edge("user", "profile")
.add_edge("user", "comments");
assert_eq!(g.edge_count(), 3);
}
#[test]
fn test_add_edge_with_sub_graph() {
let mut sub = EntityGraph::new();
sub.add_edge("comments", "author");
let mut g = EntityGraph::new();
g.add_edge_with_graph("user", "posts", sub);
assert_eq!(g.edge_count(), 1);
assert!(g.edges()[0].sub_graph.is_some());
assert_eq!(g.edges()[0].sub_graph.as_ref().unwrap().edge_count(), 1);
}
#[test]
fn test_relations_of() {
let mut g = EntityGraph::new();
g.add_edge("user", "posts")
.add_edge("user", "profile")
.add_edge("post", "comments");
let user_relations = g.relations_of("user");
assert_eq!(user_relations.len(), 2);
assert_eq!(user_relations[0].relation, "posts");
assert_eq!(user_relations[1].relation, "profile");
let post_relations = g.relations_of("post");
assert_eq!(post_relations.len(), 1);
let none = g.relations_of("nonexistent");
assert!(none.is_empty());
}
#[test]
fn test_all_relations() {
let mut g = EntityGraph::new();
g.add_edge("user", "posts")
.add_edge("user", "profile")
.add_edge("post", "comments");
let rels = g.all_relations();
assert_eq!(rels, vec!["comments", "posts", "profile"]);
}
#[test]
fn test_all_parent_fields() {
let mut g = EntityGraph::new();
g.add_edge("user", "posts")
.add_edge("user", "profile")
.add_edge("post", "comments");
let fields = g.all_parent_fields();
assert_eq!(fields, vec!["post", "user"]);
}
#[test]
fn test_all_relations_recursive() {
let mut sub = EntityGraph::new();
sub.add_edge("comments", "author")
.add_edge("comments", "likes");
let mut g = EntityGraph::new();
g.add_edge("user", "posts")
.add_edge_with_graph("user", "comments", sub);
let all = g.all_relations_recursive();
assert!(all.contains(&"posts".to_string()));
assert!(all.contains(&"comments".to_string()));
assert!(all.contains(&"author".to_string()));
assert!(all.contains(&"likes".to_string()));
assert_eq!(all.len(), 4);
}
#[test]
fn test_default_graph_is_empty() {
let g = EntityGraph::default();
assert!(g.is_empty());
}
#[test]
fn test_strategy_name() {
assert_eq!(BatchStrategy::In.name(), "in");
assert_eq!(BatchStrategy::Join.name(), "join");
assert_eq!(BatchStrategy::Subquery.name(), "subquery");
}
#[test]
fn test_strategy_default_is_in() {
assert_eq!(BatchStrategy::default(), BatchStrategy::In);
}
#[test]
fn test_render_in_clause_empty() {
let sql = BatchStrategy::render_in_clause("id", 0);
assert_eq!(sql, "id IN ()");
}
#[test]
fn test_render_in_clause_single() {
let sql = BatchStrategy::render_in_clause("id", 1);
assert_eq!(sql, "id IN (?)");
}
#[test]
fn test_render_in_clause_multiple() {
let sql = BatchStrategy::render_in_clause("user_id", 3);
assert_eq!(sql, "user_id IN (?, ?, ?)");
}
#[test]
fn test_default_config() {
let config = BatchSizeConfig::default();
assert_eq!(config.size, 100);
assert_eq!(config.strategy, BatchStrategy::In);
}
#[test]
fn test_with_size() {
let config = BatchSizeConfig::with_size(50);
assert_eq!(config.size, 50);
assert_eq!(config.strategy, BatchStrategy::In);
}
#[test]
fn test_new_with_strategy() {
let config = BatchSizeConfig::new(200, BatchStrategy::Join);
assert_eq!(config.size, 200);
assert_eq!(config.strategy, BatchStrategy::Join);
}
#[test]
fn test_batch_count_zero() {
let config = BatchSizeConfig::with_size(100);
assert_eq!(config.batch_count(0), 0);
}
#[test]
fn test_batch_count_exact_multiple() {
let config = BatchSizeConfig::with_size(100);
assert_eq!(config.batch_count(100), 1);
assert_eq!(config.batch_count(200), 2);
assert_eq!(config.batch_count(500), 5);
}
#[test]
fn test_batch_count_with_remainder() {
let config = BatchSizeConfig::with_size(100);
assert_eq!(config.batch_count(1), 1);
assert_eq!(config.batch_count(99), 1);
assert_eq!(config.batch_count(101), 2);
assert_eq!(config.batch_count(150), 2);
assert_eq!(config.batch_count(201), 3);
}
#[test]
fn test_batch_range() {
let config = BatchSizeConfig::with_size(100);
assert_eq!(config.batch_range(0, 250), 0..100);
assert_eq!(config.batch_range(1, 250), 100..200);
assert_eq!(config.batch_range(2, 250), 200..250);
}
#[test]
fn test_batch_range_exact() {
let config = BatchSizeConfig::with_size(100);
assert_eq!(config.batch_range(0, 100), 0..100);
assert_eq!(config.batch_range(1, 100), 100..100); }
#[test]
fn test_batch_range_small_batch() {
let config = BatchSizeConfig::with_size(10);
assert_eq!(config.batch_range(0, 25), 0..10);
assert_eq!(config.batch_range(1, 25), 10..20);
assert_eq!(config.batch_range(2, 25), 20..25);
}
fn make_loader() -> BatchLoader<i64, String> {
let loader = Box::new(|ids: &[i64]| -> HashMap<i64, String> {
ids.iter().map(|id| (*id, format!("user_{}", id))).collect()
});
BatchLoader::new(2, loader)
}
#[test]
fn test_batch_loader_load_many_single_batch() {
let loader = make_loader();
let result = loader.load_many(&[1, 2]);
assert_eq!(result.len(), 2);
assert_eq!(result.get(&1), Some(&"user_1".to_string()));
assert_eq!(result.get(&2), Some(&"user_2".to_string()));
}
#[test]
fn test_batch_loader_load_many_multiple_batches() {
let loader = make_loader();
let result = loader.load_many(&[1, 2, 3, 4, 5]);
assert_eq!(result.len(), 5);
for id in 1..=5 {
assert_eq!(
result.get(&id),
Some(&format!("user_{}", id)),
"missing user {}",
id
);
}
}
#[test]
fn test_batch_loader_load_one() {
let loader = make_loader();
let result = loader.load_one(&42);
assert_eq!(result, Some("user_42".to_string()));
}
#[test]
fn test_batch_loader_load_one_missing() {
let loader: BatchLoader<i64, String> =
BatchLoader::new(10, Box::new(|_ids: &[i64]| HashMap::new()));
let result = loader.load_one(&100);
assert_eq!(result, None);
}
#[test]
fn test_batch_loader_caches_results() {
let call_count = std::sync::Arc::new(std::sync::Mutex::new(0u32));
let call_count_clone = call_count.clone();
let loader = Box::new(move |ids: &[i64]| -> HashMap<i64, String> {
*call_count_clone.lock().unwrap() += 1;
ids.iter().map(|id| (*id, format!("user_{}", id))).collect()
});
let batch_loader = BatchLoader::new(100, loader);
batch_loader.load_many(&[1, 2, 3]);
assert_eq!(*call_count.lock().unwrap(), 1);
batch_loader.load_many(&[1, 2, 3]);
assert_eq!(*call_count.lock().unwrap(), 1);
batch_loader.load_many(&[4, 5]);
assert_eq!(*call_count.lock().unwrap(), 2);
}
#[test]
fn test_batch_loader_partial_cache_hit() {
let call_count = std::sync::Arc::new(std::sync::Mutex::new(0u32));
let call_count_clone = call_count.clone();
let loader = Box::new(move |ids: &[i64]| -> HashMap<i64, String> {
*call_count_clone.lock().unwrap() += 1;
ids.iter().map(|id| (*id, format!("user_{}", id))).collect()
});
let batch_loader = BatchLoader::new(100, loader);
batch_loader.load_many(&[1, 2, 3]);
assert_eq!(*call_count.lock().unwrap(), 1);
let result = batch_loader.load_many(&[1, 2, 3, 4, 5]);
assert_eq!(result.len(), 5);
assert_eq!(*call_count.lock().unwrap(), 2);
assert_eq!(batch_loader.cache_size(), 5);
}
#[test]
fn test_batch_loader_clear_cache() {
let loader = make_loader();
loader.load_many(&[1, 2]);
assert_eq!(loader.cache_size(), 2);
loader.clear_cache();
assert_eq!(loader.cache_size(), 0);
}
#[test]
fn test_batch_loader_empty_input() {
let loader = make_loader();
let result = loader.load_many(&[]);
assert!(result.is_empty());
}
#[test]
fn test_batch_loader_batch_size_attribute() {
let loader = make_loader();
assert_eq!(loader.batch_size(), 2);
}
#[test]
fn test_batch_loader_with_size_1() {
let loader = BatchLoader::new(
1,
Box::new(|ids: &[i64]| ids.iter().map(|id| (*id, *id * 10)).collect()),
);
let result = loader.load_many(&[1, 2, 3]);
assert_eq!(result.len(), 3);
assert_eq!(result.get(&1), Some(&10));
assert_eq!(result.get(&2), Some(&20));
assert_eq!(result.get(&3), Some(&30));
}
#[test]
fn test_workflow_graph_and_batch_loader() {
let mut graph = EntityGraph::new();
graph.add_edge_with_graph("user", "posts", {
let mut sub = EntityGraph::new();
sub.add_edge("posts", "comments");
sub
});
assert_eq!(graph.all_relations_recursive().len(), 2);
let user_loader = BatchLoader::new(
50,
Box::new(|ids: &[i64]| ids.iter().map(|id| (*id, format!("User#{}", id))).collect()),
);
let user_ids: Vec<i64> = (1..=123).collect();
let users = user_loader.load_many(&user_ids);
assert_eq!(users.len(), 123);
assert_eq!(user_loader.cache_size(), 123);
}
#[test]
fn test_n_plus_1_problem_solved() {
let query_count = std::sync::Arc::new(std::sync::Mutex::new(0u32));
let query_count_clone = query_count.clone();
let post_loader = BatchLoader::new(
100,
Box::new(move |user_ids: &[i64]| {
*query_count_clone.lock().unwrap() += 1;
user_ids
.iter()
.map(|uid| (*uid, format!("posts_for_user_{}", uid)))
.collect()
}),
);
let user_ids: Vec<i64> = (1..=250).collect();
let _posts = post_loader.load_many(&user_ids);
assert_eq!(*query_count.lock().unwrap(), 3);
}
}