async_graphql_dataloader/
loader.rs

1// src/loader.rs
2use crate::batcher::Batcher;
3use crate::cache::Cache;
4use crate::error::DataLoaderError;
5use async_trait::async_trait;
6use std::collections::HashMap;
7
8#[async_trait]
9pub trait BatchLoad: Send + Sync {
10    type Key: Send + Sync + Clone + Eq + std::hash::Hash + std::fmt::Debug + 'static;
11    type Value: Send + Sync + Clone + 'static;
12    type Error: Send + Sync + From<String> + std::fmt::Display + 'static;
13
14    async fn load(
15        &self,
16        keys: &[Self::Key],
17    ) -> HashMap<Self::Key, Result<Self::Value, Self::Error>>;
18}
19
20pub struct DataLoader<L: BatchLoad> {
21    batcher: std::sync::Arc<Batcher<L>>,
22    cache: std::sync::Arc<Cache<L::Key, Result<L::Value, DataLoaderError>>>,
23    max_batch_size: usize,
24    delay: std::time::Duration,
25}
26
27impl<L> DataLoader<L>
28where
29    L: BatchLoad + 'static,
30    L::Key: Clone + Eq + std::hash::Hash + std::fmt::Debug,
31    L::Value: Clone,
32    L::Error: From<String> + std::fmt::Display,
33{
34    pub fn new(loader: L) -> Self {
35        let loader_arc = std::sync::Arc::new(loader);
36        let batcher = std::sync::Arc::new(Batcher::new(std::sync::Arc::clone(&loader_arc)));
37
38        Self {
39            batcher,
40            cache: std::sync::Arc::new(Cache::new()),
41            max_batch_size: 100,
42            delay: std::time::Duration::from_millis(10),
43        }
44    }
45
46    pub fn with_max_batch_size(mut self, size: usize) -> Self {
47        self.max_batch_size = size;
48        self
49    }
50
51    pub fn with_delay(mut self, delay: std::time::Duration) -> Self {
52        self.delay = delay;
53        self
54    }
55
56    pub async fn load(&self, key: L::Key) -> Result<L::Value, DataLoaderError> {
57        // Verifica cache primeiro
58        if let Some(cached) = self.cache.get(&key) {
59            return cached;
60        }
61
62        // Agenda no batcher
63        let result = self.batcher.schedule(key.clone()).await;
64
65        // Cache o resultado
66        if let Ok(ref value) = result {
67            self.cache.set(key.clone(), Ok(value.clone()));
68        } else if let Err(ref error) = result {
69            self.cache.set(key.clone(), Err(error.clone()));
70        }
71
72        result
73    }
74
75    pub fn clear(&self) {
76        self.cache.clear();
77    }
78
79    pub fn prime(&self, key: L::Key, value: Result<L::Value, DataLoaderError>) {
80        self.cache.set(key, value);
81    }
82}
83
84impl<L> Clone for DataLoader<L>
85where
86    L: BatchLoad,
87{
88    fn clone(&self) -> Self {
89        Self {
90            batcher: std::sync::Arc::clone(&self.batcher),
91            cache: std::sync::Arc::clone(&self.cache),
92            max_batch_size: self.max_batch_size,
93            delay: self.delay,
94        }
95    }
96}
async_graphql_dataloader/loader.rs

async_graphql_dataloader/
loader.rs
