scufflecloud_core/dataloaders/
users.rs

1use 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}