scufflecloud_core/dataloaders/
users.rs1use std::collections::{HashMap, HashSet};
2use std::time::Duration;
3
4use diesel::{ExpressionMethods, QueryDsl, SelectableHelper};
5use diesel_async::pooled_connection::bb8::PooledConnection;
6use diesel_async::{AsyncPgConnection, RunQueryDsl};
7use scuffle_batching::{DataLoader, DataLoaderFetcher};
8use tokio::sync::Mutex;
9
10use crate::models::{User, UserId};
11use crate::schema::users;
12
13pub struct UserLoader(Mutex<PooledConnection<'static, AsyncPgConnection>>);
14
15impl DataLoaderFetcher for UserLoader {
16 type Key = UserId;
17 type Value = User;
18
19 async fn load(&self, keys: HashSet<Self::Key>) -> Option<HashMap<Self::Key, Self::Value>> {
20 let users = users::dsl::users
21 .filter(users::dsl::id.eq_any(keys))
22 .select(User::as_select())
23 .load::<User>(&mut self.0.lock().await)
24 .await
25 .map_err(|e| tracing::error!(err = %e, "failed to load users"))
26 .ok()?;
27
28 Some(users.into_iter().map(|u| (u.id, u)).collect())
29 }
30}
31
32impl UserLoader {
33 pub fn new(conn: PooledConnection<'static, AsyncPgConnection>) -> DataLoader<Self> {
34 DataLoader::new(Self(Mutex::new(conn)), 1000, 500, Duration::from_millis(5))
35 }
36}