scufflecloud_big_bin/dataloaders/
users.rs

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