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
}
pub fn detect_cycles(&self) -> Result<(), Vec<String>> {
let mut adj: std::collections::HashMap<String, Vec<String>> =
std::collections::HashMap::new();
self.collect_edges_recursive(&mut adj);
for neighbors in adj.values_mut() {
neighbors.sort();
}
let mut visited = std::collections::HashSet::new();
let mut visiting = std::collections::HashSet::new();
let mut path = Vec::new();
let mut sorted_nodes: Vec<String> = adj.keys().cloned().collect();
sorted_nodes.sort();
for node in &sorted_nodes {
if !visited.contains(node) {
dfs_cycle_detect(node, &adj, &mut visited, &mut visiting, &mut path)?;
}
}
Ok(())
}
fn collect_edges_recursive(&self, adj: &mut std::collections::HashMap<String, Vec<String>>) {
for edge in &self.edges {
adj.entry(edge.parent_field.clone())
.or_default()
.push(edge.relation.clone());
if let Some(sub) = &edge.sub_graph {
sub.collect_edges_recursive(adj);
}
}
}
pub fn detect_duplicate_edges(&self) -> Result<(), Vec<(String, String)>> {
let mut seen = std::collections::HashSet::new();
let mut duplicates = Vec::new();
for edge in &self.edges {
let key = (edge.parent_field.clone(), edge.relation.clone());
if !seen.insert(key.clone()) {
duplicates.push((edge.parent_field.clone(), edge.relation.clone()));
}
}
if duplicates.is_empty() {
Ok(())
} else {
Err(duplicates)
}
}
pub fn validate(&self) -> Result<(), String> {
if let Err(cycle) = self.detect_cycles() {
return Err(format!(
"EntityGraph 循环引用检测失败:{}",
cycle.join(" → ")
));
}
if let Err(duplicates) = self.detect_duplicate_edges() {
let dup_str: Vec<String> = duplicates
.iter()
.map(|(p, r)| format!("({}->{})", p, r))
.collect();
return Err(format!(
"EntityGraph 重复边检测失败:{}",
dup_str.join(", ")
));
}
Ok(())
}
}
fn dfs_cycle_detect(
node: &str,
adj: &std::collections::HashMap<String, Vec<String>>,
visited: &mut std::collections::HashSet<String>,
visiting: &mut std::collections::HashSet<String>,
path: &mut Vec<String>,
) -> Result<(), Vec<String>> {
if visiting.contains(node) {
let cycle_start = path.iter().position(|n| n == node).unwrap_or(0);
let mut cycle = path[cycle_start..].to_vec();
cycle.push(node.to_string());
return Err(cycle);
}
if visited.contains(node) {
return Ok(());
}
visiting.insert(node.to_string());
path.push(node.to_string());
if let Some(neighbors) = adj.get(node) {
for neighbor in neighbors {
dfs_cycle_detect(neighbor, adj, visited, visiting, path)?;
}
}
visiting.remove(node);
visited.insert(node.to_string());
path.pop();
Ok(())
}
#[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 mut to_load: Vec<K> = Vec::new();
if let Ok(cached) = self.cache.read() {
for k in keys {
if let Some(v) = cached.get(k) {
result.insert(k.clone(), v.clone());
} else {
to_load.push(k.clone());
}
}
} else {
to_load.extend(keys.iter().cloned());
}
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);
}
if let Ok(mut cache) = self.cache.write() {
for (k, v) in &all_loaded {
cache.insert(k.clone(), v.clone());
}
}
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) {
if let Ok(mut cache) = self.cache.write() {
cache.clear();
}
}
pub fn cache_size(&self) -> usize {
match self.cache.read() {
Ok(g) => g.len(),
Err(_) => 0,
}
}
pub fn batch_size(&self) -> usize {
self.batch_size
}
}
pub struct N1QueryDetector {
config: N1DetectionConfig,
counts: RwLock<HashMap<String, u64>>,
batch_counts: RwLock<HashMap<String, u64>>,
window_active: RwLock<bool>,
alerts: RwLock<Vec<N1Alert>>,
}
#[derive(Debug, Clone)]
pub struct N1DetectionConfig {
pub threshold: u64,
pub enabled: bool,
}
impl Default for N1DetectionConfig {
fn default() -> Self {
Self {
threshold: 5,
enabled: true,
}
}
}
impl N1DetectionConfig {
pub fn new() -> Self {
Self::default()
}
pub fn with_threshold(mut self, threshold: u64) -> Self {
self.threshold = threshold.max(1);
self
}
pub fn with_enabled(mut self, enabled: bool) -> Self {
self.enabled = enabled;
self
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct N1Alert {
pub relation: String,
pub query_count: u64,
pub batch_count: u64,
pub threshold: u64,
}
impl N1Alert {
pub fn no_batch_used(&self) -> bool {
self.batch_count == 0
}
pub fn suggested_batch_size(&self) -> usize {
let n = self.query_count as usize;
if n <= 50 {
50
} else if n <= 100 {
100
} else if n <= 500 {
500
} else {
1000
}
}
}
impl N1QueryDetector {
pub fn new(config: N1DetectionConfig) -> Self {
Self {
config,
counts: RwLock::new(HashMap::new()),
batch_counts: RwLock::new(HashMap::new()),
window_active: RwLock::new(false),
alerts: RwLock::new(Vec::new()),
}
}
pub fn with_defaults() -> Self {
Self::new(N1DetectionConfig::default())
}
pub fn is_enabled(&self) -> bool {
self.config.enabled
}
pub fn threshold(&self) -> u64 {
self.config.threshold
}
pub fn start_window(&self) {
if !self.config.enabled {
return;
}
if let Ok(mut counts) = self.counts.write() {
*counts = HashMap::new();
}
if let Ok(mut batch_counts) = self.batch_counts.write() {
*batch_counts = HashMap::new();
}
if let Ok(mut alerts) = self.alerts.write() {
*alerts = Vec::new();
}
if let Ok(mut window_active) = self.window_active.write() {
*window_active = true;
}
}
pub fn end_window(&self) -> Vec<N1Alert> {
if !self.config.enabled {
return Vec::new();
}
if let Ok(mut window_active) = self.window_active.write() {
*window_active = false;
}
let new_alerts: Vec<N1Alert> = match (self.counts.read(), self.batch_counts.read()) {
(Ok(counts), Ok(batch_counts)) => {
let mut alerts: Vec<N1Alert> = counts
.iter()
.filter_map(|(rel, &cnt)| {
if cnt >= self.config.threshold {
Some(N1Alert {
relation: rel.clone(),
query_count: cnt,
batch_count: *batch_counts.get(rel).unwrap_or(&0),
threshold: self.config.threshold,
})
} else {
None
}
})
.collect();
alerts.sort_by(|a, b| a.relation.cmp(&b.relation));
alerts
}
_ => Vec::new(),
};
if let Ok(mut alerts) = self.alerts.write() {
*alerts = new_alerts.clone();
}
new_alerts
}
pub fn record_single_load(&self, relation: &str) {
if !self.config.enabled {
return;
}
{
let active = self.window_active.read().map(|g| *g).unwrap_or(false);
if !active {
return;
}
}
if let Ok(mut counts) = self.counts.write() {
*counts.entry(relation.to_string()).or_insert(0) += 1;
}
}
pub fn record_batch_load(&self, relation: &str, _keys_count: usize) {
if !self.config.enabled {
return;
}
{
let active = self.window_active.read().map(|g| *g).unwrap_or(false);
if !active {
return;
}
}
if let Ok(mut batch_counts) = self.batch_counts.write() {
*batch_counts.entry(relation.to_string()).or_insert(0) += 1;
}
}
pub fn alerts(&self) -> Vec<N1Alert> {
self.alerts.read().map(|g| g.clone()).unwrap_or_default()
}
pub fn current_count(&self, relation: &str) -> u64 {
self.counts
.read()
.map(|g| g.get(relation).copied().unwrap_or(0))
.unwrap_or(0)
}
pub fn current_batch_count(&self, relation: &str) -> u64 {
self.batch_counts
.read()
.map(|g| g.get(relation).copied().unwrap_or(0))
.unwrap_or(0)
}
pub fn is_window_active(&self) -> bool {
self.window_active.read().map(|g| *g).unwrap_or(false)
}
pub fn has_n_plus_one(&self) -> bool {
!self.alerts().is_empty()
}
}
impl Default for N1QueryDetector {
fn default() -> Self {
Self::with_defaults()
}
}
#[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_n1_config_default() {
let cfg = N1DetectionConfig::default();
assert_eq!(cfg.threshold, 5);
assert!(cfg.enabled);
}
#[test]
fn test_n1_config_builder() {
let cfg = N1DetectionConfig::new()
.with_threshold(10)
.with_enabled(false);
assert_eq!(cfg.threshold, 10);
assert!(!cfg.enabled);
let cfg2 = N1DetectionConfig::new().with_threshold(0);
assert_eq!(cfg2.threshold, 1);
}
#[test]
fn test_n1_detector_default() {
let det = N1QueryDetector::default();
assert!(det.is_enabled());
assert_eq!(det.threshold(), 5);
assert!(!det.is_window_active());
assert!(!det.has_n_plus_one());
}
#[test]
fn test_n1_detector_disabled_is_noop() {
let det = N1QueryDetector::new(N1DetectionConfig::new().with_enabled(false));
det.start_window();
for _ in 0..100 {
det.record_single_load("posts");
}
assert_eq!(det.current_count("posts"), 0);
let alerts = det.end_window();
assert!(alerts.is_empty());
}
#[test]
fn test_n1_detector_records_outside_window_ignored() {
let det = N1QueryDetector::with_defaults();
det.record_single_load("posts");
assert_eq!(det.current_count("posts"), 0);
}
#[test]
fn test_n1_detector_below_threshold_no_alert() {
let det = N1QueryDetector::with_defaults(); det.start_window();
for _ in 0..4 {
det.record_single_load("posts");
}
assert_eq!(det.current_count("posts"), 4);
let alerts = det.end_window();
assert!(alerts.is_empty(), "below threshold should not alert");
assert!(!det.has_n_plus_one());
}
#[test]
fn test_n1_detector_at_threshold_triggers_alert() {
let det = N1QueryDetector::with_defaults(); det.start_window();
for _ in 0..5 {
det.record_single_load("posts");
}
let alerts = det.end_window();
assert_eq!(alerts.len(), 1);
assert_eq!(alerts[0].relation, "posts");
assert_eq!(alerts[0].query_count, 5);
assert_eq!(alerts[0].threshold, 5);
assert_eq!(alerts[0].batch_count, 0);
assert!(alerts[0].no_batch_used());
assert!(det.has_n_plus_one());
}
#[test]
fn test_n1_detector_above_threshold_triggers_alert() {
let det = N1QueryDetector::with_defaults();
det.start_window();
for _ in 0..10 {
det.record_single_load("posts");
}
let alerts = det.end_window();
assert_eq!(alerts.len(), 1);
assert_eq!(alerts[0].query_count, 10);
assert!(alerts[0].suggested_batch_size() >= 50);
}
#[test]
fn test_n1_detector_multiple_relations() {
let det = N1QueryDetector::with_defaults();
det.start_window();
for _ in 0..6 {
det.record_single_load("posts");
}
for _ in 0..3 {
det.record_single_load("comments"); }
for _ in 0..8 {
det.record_single_load("tags");
}
let alerts = det.end_window();
assert_eq!(alerts.len(), 2);
assert_eq!(alerts[0].relation, "posts");
assert_eq!(alerts[0].query_count, 6);
assert_eq!(alerts[1].relation, "tags");
assert_eq!(alerts[1].query_count, 8);
}
#[test]
fn test_n1_detector_batch_load_recorded_separately() {
let det = N1QueryDetector::with_defaults();
det.start_window();
for _ in 0..6 {
det.record_single_load("posts");
}
det.record_batch_load("posts", 100);
det.record_batch_load("posts", 50);
let alerts = det.end_window();
assert_eq!(alerts.len(), 1);
assert_eq!(alerts[0].query_count, 6);
assert_eq!(alerts[0].batch_count, 2);
assert!(!alerts[0].no_batch_used());
}
#[test]
fn test_n1_detector_batch_only_does_not_trigger() {
let det = N1QueryDetector::with_defaults();
det.start_window();
for _ in 0..100 {
det.record_batch_load("posts", 50);
}
assert_eq!(det.current_batch_count("posts"), 100);
assert_eq!(det.current_count("posts"), 0);
let alerts = det.end_window();
assert!(alerts.is_empty());
}
#[test]
fn test_n1_detector_start_window_resets() {
let det = N1QueryDetector::with_defaults();
det.start_window();
for _ in 0..10 {
det.record_single_load("posts");
}
let _ = det.end_window();
assert_eq!(det.alerts().len(), 1);
det.start_window();
assert_eq!(det.alerts().len(), 0);
assert_eq!(det.current_count("posts"), 0);
assert!(det.is_window_active());
}
#[test]
fn test_n1_detector_end_window_deactivates() {
let det = N1QueryDetector::with_defaults();
det.start_window();
assert!(det.is_window_active());
det.end_window();
assert!(!det.is_window_active());
det.record_single_load("posts");
assert_eq!(det.current_count("posts"), 0);
}
#[test]
fn test_n1_detector_custom_threshold() {
let det = N1QueryDetector::new(N1DetectionConfig::new().with_threshold(100));
det.start_window();
for _ in 0..50 {
det.record_single_load("posts");
}
let alerts = det.end_window();
assert!(alerts.is_empty(), "below custom threshold should not alert");
det.start_window();
for _ in 0..100 {
det.record_single_load("posts");
}
let alerts = det.end_window();
assert_eq!(alerts.len(), 1);
assert_eq!(alerts[0].threshold, 100);
assert_eq!(alerts[0].query_count, 100);
}
#[test]
fn test_n1_alert_suggested_batch_size() {
let mk = |cnt: u64| N1Alert {
relation: "x".into(),
query_count: cnt,
batch_count: 0,
threshold: 5,
};
assert_eq!(mk(5).suggested_batch_size(), 50);
assert_eq!(mk(50).suggested_batch_size(), 50);
assert_eq!(mk(51).suggested_batch_size(), 100);
assert_eq!(mk(100).suggested_batch_size(), 100);
assert_eq!(mk(101).suggested_batch_size(), 500);
assert_eq!(mk(500).suggested_batch_size(), 500);
assert_eq!(mk(501).suggested_batch_size(), 1000);
assert_eq!(mk(10000).suggested_batch_size(), 1000);
}
#[test]
fn test_n1_detector_real_n_plus_one_scenario() {
let det = N1QueryDetector::with_defaults();
det.start_window();
let user_ids: Vec<i64> = (1..=20).collect();
for _uid in &user_ids {
det.record_single_load("posts");
}
let alerts = det.end_window();
assert_eq!(alerts.len(), 1);
assert_eq!(alerts[0].query_count, 20);
assert!(alerts[0].no_batch_used());
det.start_window();
det.record_batch_load("posts", 20); let alerts2 = det.end_window();
assert!(alerts2.is_empty(), "batch loading should not trigger N+1");
}
#[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);
}
}