use std::sync::Arc; use youmubot_db_sql::{models::osu as models, Pool}; use youmubot_prelude::*; use crate::{ models::{ApprovalStatus, Beatmap, Mode}, OsuClient, }; /// BeatmapMetaCache intercepts beatmap-by-id requests and caches them for later recalling. /// Does not cache non-Ranked beatmaps. #[derive(Clone)] pub struct BeatmapMetaCache { client: Arc, pool: Pool, } impl std::fmt::Debug for BeatmapMetaCache { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { write!(f, "") } } impl TypeMapKey for BeatmapMetaCache { type Value = BeatmapMetaCache; } impl BeatmapMetaCache { /// Create a new beatmap cache. pub fn new(client: Arc, pool: Pool) -> Self { BeatmapMetaCache { client, pool } } /// Clean the cache. pub async fn clear(&self) -> Result<()> { models::CachedBeatmap::clear_all(&self.pool).await?; Ok(()) } #[allow(clippy::wrong_self_convention)] fn to_cached_beatmap(beatmap: &Beatmap, mode: Option) -> models::CachedBeatmap { models::CachedBeatmap { beatmap_id: beatmap.beatmap_id as i64, mode: mode.unwrap_or(beatmap.mode) as u8, cached_at: chrono::Utc::now(), beatmap: bincode::serialize(&beatmap).unwrap(), } } async fn insert_if_possible(&self, id: u64, mode: Option) -> Result { let beatmap = self .client .beatmaps(crate::BeatmapRequestKind::Beatmap(id), |f| { if let Some(mode) = mode { f.mode(mode, true); } f }) .await .and_then(|v| { v.into_iter() .next() .ok_or_else(|| Error::msg("beatmap not found")) })?; if let ApprovalStatus::Ranked(_) = beatmap.approval { let mut c = Self::to_cached_beatmap(&beatmap, mode); c.store(&self.pool).await.pls_ok(); }; Ok(beatmap) } async fn get_beatmap_db(&self, id: u64, mode: Mode) -> Result> { Ok( models::CachedBeatmap::by_id(id as i64, mode as u8, &self.pool) .await? .map(|v| bincode::deserialize(&v.beatmap[..]).unwrap()), ) } /// Get the given beatmap pub async fn get_beatmap(&self, id: u64, mode: Mode) -> Result { match self.get_beatmap_db(id, mode).await? { Some(v) => Ok(v), None => self.insert_if_possible(id, Some(mode)).await, } } /// Get a beatmap without a mode... pub async fn get_beatmap_default(&self, id: u64) -> Result { for mode in [Mode::Std, Mode::Taiko, Mode::Catch, Mode::Mania].into_iter() { if let Ok(Some(bm)) = self.get_beatmap_db(id, mode).await { if bm.mode == mode { return Ok(bm); } } } self.insert_if_possible(id, None).await } /// Get a beatmapset from its ID. pub async fn get_beatmapset(&self, id: u64) -> Result> { let bms = models::CachedBeatmap::by_beatmapset(id as i64, &self.pool).await?; if !bms.is_empty() { return Ok(bms .into_iter() .map(|v| bincode::deserialize(&v.beatmap[..]).unwrap()) .collect()); } let mut beatmaps = self .client .beatmaps(crate::BeatmapRequestKind::Beatmapset(id), |f| f) .await?; if beatmaps.is_empty() { return Err(Error::msg("beatmapset not found")); } beatmaps.sort_by_key(|b| (b.mode as u8, (b.difficulty.stars * 1000.0) as u64)); // Cast so that Ord is maintained if let ApprovalStatus::Ranked(_) = &beatmaps[0].approval { // Save each beatmap. let mut t = self.pool.begin().await?; for b in &beatmaps { let mut b = Self::to_cached_beatmap(b, None); b.store(&mut *t).await?; // Save the beatmapset mapping. b.link_beatmapset(id as i64, &mut *t).await?; } t.commit().await?; } Ok(beatmaps) } }