refactor(beatmap,beatmapset): use to ensure beatmap exists

app/signalr/hub/multiplayer.py

@@ -11,6 +11,7 @@ from app.database.multiplayer_event import MultiplayerEvent
 from app.database.playlists import Playlist
 from app.database.relationship import Relationship, RelationshipType
 from app.dependencies.database import engine, get_redis
+from app.dependencies.fetcher import get_fetcher
 from app.exception import InvokeException
 from app.log import logger
 from app.models.mods import APIMod
@@ -44,8 +45,9 @@ from app.models.score import GameMode
 
 from .hub import Client, Hub
 
+from httpx import HTTPError
 from sqlalchemy import update
-from sqlmodel import col, select
+from sqlmodel import col, exists, select
 from sqlmodel.ext.asyncio.session import AsyncSession
 
 GAMEPLAY_LOAD_TIMEOUT = 30
@@ -191,11 +193,25 @@ class MultiplayerHub(Hub[MultiplayerClientState]):
                 session.add(db_room)
                 await session.commit()
                 await session.refresh(db_room)
+
                 item = room.playlist[0]
                 item.owner_id = client.user_id
                 room.room_id = db_room.id
                 starts_at = db_room.starts_at or datetime.now(UTC)
+                beatmap_exists = await session.exec(
+                    select(exists().where(col(Beatmap.id) == item.beatmap_id))
+                )
+                if not beatmap_exists.one():
+                    fetcher = await get_fetcher()
+                    try:
+                        resp = await fetcher.get_beatmap(item.beatmap_id)
+                        await Beatmap.from_resp(session, resp)
+                    except HTTPError:
+                        raise InvokeException(
+                            "Failed to fetch beatmap, please retry later"
+                        )
                 await Playlist.add_to_db(item, db_room.id, session)
+
                 server_room = ServerMultiplayerRoom(
                     room=room,
                     category=RoomCategory.NORMAL,
@@ -372,6 +388,7 @@ class MultiplayerHub(Hub[MultiplayerClientState]):
         )
 
     async def validate_styles(self, room: ServerMultiplayerRoom):
+        fetcher = await get_fetcher()
         if not room.queue.current_item.freestyle:
             for user in room.room.users:
                 await self.change_user_style(
@@ -381,9 +398,12 @@ class MultiplayerHub(Hub[MultiplayerClientState]):
                     user,
                 )
         async with AsyncSession(engine) as session:
-            beatmap = await session.get(Beatmap, room.queue.current_item.beatmap_id)
-            if beatmap is None:
-                raise InvokeException("Beatmap not found")
+            try:
+                beatmap = await Beatmap.get_or_fetch(
+                    session, fetcher, bid=room.queue.current_item.beatmap_id
+                )
+            except HTTPError:
+                raise InvokeException("Current item beatmap not found")
             beatmap_ids = (
                 await session.exec(
                     select(Beatmap.id, Beatmap.mode).where(

app/signalr/hub/spectator.py

@@ -11,6 +11,7 @@ from app.database import Beatmap, User
 from app.database.score import Score
 from app.database.score_token import ScoreToken
 from app.dependencies.database import engine
+from app.dependencies.fetcher import get_fetcher
 from app.models.beatmap import BeatmapRankStatus
 from app.models.mods import mods_to_int
 from app.models.score import LegacyReplaySoloScoreInfo, ScoreStatistics
@@ -179,15 +180,13 @@ class SpectatorHub(Hub[StoreClientState]):
             return
         if state.beatmap_id is None or state.ruleset_id is None:
             return
+
+        fetcher = await get_fetcher()
         async with AsyncSession(engine) as session:
             async with session.begin():
-                beatmap = (
-                    await session.exec(
-                        select(Beatmap).where(Beatmap.id == state.beatmap_id)
-                    )
-                ).first()
-                if not beatmap:
-                    return
+                beatmap = await Beatmap.get_or_fetch(
+                    session, fetcher, bid=state.beatmap_id
+                )
                 user = (
                     await session.exec(select(User).where(User.id == user_id))
                 ).first()
@@ -237,16 +236,16 @@ class SpectatorHub(Hub[StoreClientState]):
         user_id = int(client.connection_id)
         store = self.get_or_create_state(client)
         score = store.score
-        assert store.beatmap_status is not None
-        assert store.state is not None
-        assert store.score is not None
-        if not score or not store.score_token:
+        if (
+            score is None
+            or store.score_token is None
+            or store.beatmap_status is None
+            or store.state is None
+        ):
             return
         if (
             BeatmapRankStatus.PENDING < store.beatmap_status <= BeatmapRankStatus.LOVED
-        ) and any(
-            k.is_hit() and v > 0 for k, v in store.score.score_info.statistics.items()
-        ):
+        ) and any(k.is_hit() and v > 0 for k, v in score.score_info.statistics.items()):
             await self._process_score(store, client)
         store.state = None
         store.beatmap_status = None