use SQL's limit/offset pagination for nextIndex/maxCount requests (#185)

Instead of retrieving the entire list of items/characters/scores/etc. at once (and even store them in memory), use SQL's `LIMIT ... OFFSET ...` pagination so we only take what we need.

Currently only CHUNITHM uses this, but this will also affect maimai DX and O.N.G.E.K.I. once the PR is ready.

Also snuck in a fix for CHUNITHM/maimai DX's `GetUserRivalMusicApi` to respect the `userRivalMusicLevelList` sent by the client.

### How this works

Say we have a `GetUserCharacterApi` request:

```json
{
    "userId": 10000,
    "maxCount": 700,
    "nextIndex": 0
}
```

Instead of getting the entire character list from the database (which can be very large if the user force unlocked everything), add limit/offset to the query:

```python
select(character)
.where(character.c.user == user_id)
.order_by(character.c.id.asc())
.limit(max_count + 1)
.offset(next_index)
```

The query takes `maxCount + 1` items from the database to determine if there is more items than can be returned:

```python
rows = ...

if len(rows) > max_count:
    # return only max_count rows
    next_index += max_count
else:
    # return everything left
    next_index = -1
```

This has the benefit of not needing to load everything into memory (and also having to store server state, as seen in the [`SCORE_BUFFER` list](2274b42358/titles/chuni/base.py (L13)).)

Reviewed-on: https://gitea.tendokyu.moe/Hay1tsme/artemis/pulls/185
Co-authored-by: beerpsi <beerpsi@duck.com>
Co-committed-by: beerpsi <beerpsi@duck.com>

titles/ongeki/base.py

@@ -1,16 +1,16 @@
-from datetime import date, datetime, timedelta
-from typing import Any, Dict, List
+import itertools
 import json
 import logging
+from datetime import datetime, timedelta
 from enum import Enum
+from typing import Any, Dict, List
 
 import pytz
+
 from core.config import CoreConfig
-from core.data.cache import cached
+from titles.ongeki.config import OngekiConfig
 from titles.ongeki.const import OngekiConstants
-from titles.ongeki.config import OngekiConfig
 from titles.ongeki.database import OngekiData
-from titles.ongeki.config import OngekiConfig
 
 
 class OngekiBattleGrade(Enum):
@@ -500,57 +500,93 @@ class OngekiBase:
         }
 
     async def handle_get_user_music_api_request(self, data: Dict) -> Dict:
-        song_list = await self.util_generate_music_list(data["userId"])
-        max_ct = data["maxCount"]
-        next_idx = data["nextIndex"]
-        start_idx = next_idx
-        end_idx = max_ct + start_idx
+        user_id: int = data["userId"]
+        next_idx: int = data["nextIndex"]
+        max_ct: int = data["maxCount"]
 
-        if len(song_list[start_idx:]) > max_ct:
+        rows = await self.data.score.get_best_scores(
+            user_id, limit=max_ct + 1, offset=next_idx
+        )
+
+        if rows is None:
+            return {
+                "userId": user_id,
+                "length": 0,
+                "nextIndex": 0,
+                "userMusicList": [],
+            }
+
+        music_details = [row._asdict() for row in rows]
+        returned_count = 0
+        music_list = []
+
+        for _music_id, details_iter in itertools.groupby(music_details, key=lambda d: d["musicId"]):
+            details: list[dict[Any, Any]] = []
+
+            for d in details_iter:
+                d.pop("id")
+                d.pop("user")
+                
+                details.append(d)
+
+            music_list.append({"length": len(details), "userMusicDetailList": details})
+            returned_count += len(details)
+
+            if len(music_list) >= max_ct:
+                break
+        
+        if returned_count < len(rows):
             next_idx += max_ct
-
         else:
-            next_idx = -1
+            next_idx = 0
 
         return {
-            "userId": data["userId"],
-            "length": len(song_list[start_idx:end_idx]),
+            "userId": user_id,
+            "length": len(music_list),
             "nextIndex": next_idx,
-            "userMusicList": song_list[start_idx:end_idx],
+            "userMusicList": music_list,
         }
 
     async def handle_get_user_item_api_request(self, data: Dict) -> Dict:
-        kind = data["nextIndex"] / 10000000000
-        p = await self.data.item.get_items(data["userId"], kind)
+        user_id: int = data["userId"]
+        next_idx: int = data["nextIndex"]
+        max_ct: int = data["maxCount"]
 
-        if p is None:
+        kind = next_idx // 10000000000
+        next_idx = next_idx % 10000000000
+
+        rows = await self.data.item.get_items(
+            user_id, kind, limit=max_ct + 1, offset=next_idx
+        )
+
+        if rows is None:
             return {
-                "userId": data["userId"],
-                "nextIndex": -1,
+                "userId": user_id,
+                "nextIndex": 0,
                 "itemKind": kind,
+                "length": 0,
                 "userItemList": [],
             }
 
         items: List[Dict[str, Any]] = []
-        for i in range(data["nextIndex"] % 10000000000, len(p)):
-            if len(items) > data["maxCount"]:
-                break
-            tmp = p[i]._asdict()
-            tmp.pop("user")
-            tmp.pop("id")
-            items.append(tmp)
 
-        xout = kind * 10000000000 + (data["nextIndex"] % 10000000000) + len(items)
+        for row in rows[:max_ct]:
+            item = row._asdict()
+            
+            item.pop("id")
+            item.pop("user")
+            
+            items.append(item)
 
-        if len(items) < data["maxCount"] or data["maxCount"] == 0:
-            nextIndex = 0
+        if len(rows) > max_ct:
+            next_idx = kind * 10000000000 + next_idx + max_ct
         else:
-            nextIndex = xout
+            next_idx = 0
 
         return {
-            "userId": data["userId"],
-            "nextIndex": int(nextIndex),
-            "itemKind": int(kind),
+            "userId": user_id,
+            "nextIndex": next_idx,
+            "itemKind": kind,
             "length": len(items),
             "userItemList": items,
         }
@@ -1143,43 +1179,56 @@ class OngekiBase:
         """
         Added in Bright
         """
-        rival_id = data["rivalUserId"]
-        next_idx = data["nextIndex"]
-        max_ct = data["maxCount"]
-        music = self.handle_get_user_music_api_request(
-            {"userId": rival_id, "nextIndex": next_idx, "maxCount": max_ct}
+        user_id: int = data["userId"]
+        rival_id: int = data["rivalUserId"]
+        next_idx: int = data["nextIndex"]
+        max_ct: int = data["maxCount"]
+
+        rows = await self.data.score.get_best_scores(
+            rival_id, limit=max_ct + 1, offset=next_idx
         )
 
-        for song in music["userMusicList"]:
-            song["userRivalMusicDetailList"] = song["userMusicDetailList"]
-            song.pop("userMusicDetailList")
+        if rows is None:
+            return {
+                "userId": user_id,
+                "rivalUserId": rival_id,
+                "nextIndex": 0,
+                "length": 0,
+                "userRivalMusicList": [],
+            }
+        
+        music_details = [row._asdict() for row in rows]
+        returned_count = 0
+        music_list = []
+
+        for _music_id, details_iter in itertools.groupby(music_details, key=lambda d: d["musicId"]):
+            details: list[dict[Any, Any]] = []
+
+            for d in details_iter:
+                d.pop("id")
+                d.pop("user")
+                d.pop("playCount")
+                d.pop("isLock")
+                d.pop("clearStatus")
+                d.pop("isStoryWatched")
+
+                details.append(d)
+
+            music_list.append({"length": len(details), "userRivalMusicDetailList": details})
+            returned_count += len(details)
+
+            if len(music_list) >= max_ct:
+                break
+        
+        if returned_count < len(rows):
+            next_idx += max_ct
+        else:
+            next_idx = 0
+
         return {
-            "userId": data["userId"],
+            "userId": user_id,
             "rivalUserId": rival_id,
-            "length": music["length"],
-            "nextIndex": music["nextIndex"],
-            "userRivalMusicList": music["userMusicList"],
+            "nextIndex": next_idx,
+            "length": len(music_list),
+            "userRivalMusicList": music_list,
         }
-
-    @cached(2)
-    async def util_generate_music_list(self, user_id: int) -> List:
-        music_detail = await self.data.score.get_best_scores(user_id)
-        song_list = []
-
-        for md in music_detail:
-            found = False
-            tmp = md._asdict()
-            tmp.pop("user")
-            tmp.pop("id")
-
-            for song in song_list:
-                if song["userMusicDetailList"][0]["musicId"] == tmp["musicId"]:
-                    found = True
-                    song["userMusicDetailList"].append(tmp)
-                    song["length"] = len(song["userMusicDetailList"])
-                    break
-
-            if not found:
-                song_list.append({"length": 1, "userMusicDetailList": [tmp]})
-
-        return song_list

titles/ongeki/bright.py

@@ -1,13 +1,11 @@
-from datetime import date, datetime, timedelta
-from typing import Any, Dict
+from datetime import datetime
 from random import randint
-import pytz
-import json
+from typing import Dict
 
 from core.config import CoreConfig
 from titles.ongeki.base import OngekiBase
-from titles.ongeki.const import OngekiConstants
 from titles.ongeki.config import OngekiConfig
+from titles.ongeki.const import OngekiConstants
 
 
 class OngekiBright(OngekiBase):
@@ -62,66 +60,72 @@ class OngekiBright(OngekiBase):
         return {"returnCode": 1}
 
     async def handle_cm_get_user_card_api_request(self, data: Dict) -> Dict:
-        user_cards = await self.data.item.get_cards(data["userId"])
-        if user_cards is None:
+        user_id: int = data["userId"]
+        max_ct: int = data["maxCount"]
+        next_idx: int = data["nextIndex"]
+
+        rows = await self.data.item.get_cards(
+            user_id, limit=max_ct + 1, offset=next_idx
+        )
+
+        if rows is None:
             return {}
+        
+        card_list = []
 
-        max_ct = data["maxCount"]
-        next_idx = data["nextIndex"]
-        start_idx = next_idx
-        end_idx = max_ct + start_idx
-
-        if len(user_cards[start_idx:]) > max_ct:
+        for row in rows[:max_ct]:
+            card = row._asdict()
+            card.pop("id")
+            card.pop("user")
+            card_list.append(card)
+        
+        if len(rows) > max_ct:
             next_idx += max_ct
         else:
-            next_idx = -1
-
-        card_list = []
-        for card in user_cards:
-            tmp = card._asdict()
-            tmp.pop("id")
-            tmp.pop("user")
-            card_list.append(tmp)
+            next_idx = 0
 
         return {
             "userId": data["userId"],
-            "length": len(card_list[start_idx:end_idx]),
+            "length": len(card_list),
             "nextIndex": next_idx,
-            "userCardList": card_list[start_idx:end_idx],
+            "userCardList": card_list,
         }
 
     async def handle_cm_get_user_character_api_request(self, data: Dict) -> Dict:
-        user_characters = await self.data.item.get_characters(data["userId"])
-        if user_characters is None:
+        user_id: int = data["userId"]
+        max_ct: int = data["maxCount"]
+        next_idx: int = data["nextIndex"]
+
+        rows = await self.data.item.get_characters(
+            user_id, limit=max_ct + 1, offset=next_idx
+        )
+
+        if rows is None:
             return {
-                "userId": data["userId"],
+                "userId": user_id,
                 "length": 0,
                 "nextIndex": 0,
                 "userCharacterList": [],
             }
 
-        max_ct = data["maxCount"]
-        next_idx = data["nextIndex"]
-        start_idx = next_idx
-        end_idx = max_ct + start_idx
+        character_list = []
 
-        if len(user_characters[start_idx:]) > max_ct:
+        for row in rows[:max_ct]:
+            character = row._asdict()
+            character.pop("id")
+            character.pop("user")
+            character_list.append(character)
+
+        if len(rows) > max_ct:
             next_idx += max_ct
         else:
-            next_idx = -1
-
-        character_list = []
-        for character in user_characters:
-            tmp = character._asdict()
-            tmp.pop("id")
-            tmp.pop("user")
-            character_list.append(tmp)
+            next_idx = 0
 
         return {
             "userId": data["userId"],
-            "length": len(character_list[start_idx:end_idx]),
+            "length": len(character_list),
             "nextIndex": next_idx,
-            "userCharacterList": character_list[start_idx:end_idx],
+            "userCharacterList": character_list,
         }
 
     async def handle_get_user_gacha_api_request(self, data: Dict) -> Dict:

titles/ongeki/schema/item.py

@@ -1,15 +1,16 @@
-from datetime import date, datetime, timedelta
-from typing import Dict, Optional, List
-from sqlalchemy import Table, Column, UniqueConstraint, PrimaryKeyConstraint, and_
-from sqlalchemy.types import Integer, String, TIMESTAMP, Boolean, JSON
-from sqlalchemy.schema import ForeignKey
-from sqlalchemy.engine import Row
-from sqlalchemy.sql import func, select
+from datetime import datetime
+from typing import Dict, List, Optional
+
+from sqlalchemy import Column, Table, UniqueConstraint, and_
 from sqlalchemy.dialects.mysql import insert
+from sqlalchemy.engine import Row
+from sqlalchemy.schema import ForeignKey
+from sqlalchemy.sql import func, select
+from sqlalchemy.types import TIMESTAMP, Boolean, Integer, String
 
 from core.data.schema import BaseData, metadata
 
-card = Table(
+card: Table = Table(
     "ongeki_user_card",
     metadata,
     Column("id", Integer, primary_key=True, nullable=False),
@@ -45,7 +46,7 @@ deck = Table(
     mysql_charset="utf8mb4",
 )
 
-character = Table(
+character: Table = Table(
     "ongeki_user_character",
     metadata,
     Column("id", Integer, primary_key=True, nullable=False),
@@ -130,7 +131,7 @@ memorychapter = Table(
     mysql_charset="utf8mb4",
 )
 
-item = Table(
+item: Table = Table(
     "ongeki_user_item",
     metadata,
     Column("id", Integer, primary_key=True, nullable=False),
@@ -351,9 +352,18 @@ class OngekiItemData(BaseData):
             return None
         return result.lastrowid
 
-    async def get_cards(self, aime_id: int) -> Optional[List[Dict]]:
+    async def get_cards(
+        self, aime_id: int, limit: Optional[int] = None, offset: Optional[int] = None
+    ) -> Optional[List[Row]]:
         sql = select(card).where(card.c.user == aime_id)
 
+        if limit is not None or offset is not None:
+            sql = sql.order_by(card.c.id)
+        if limit is not None:
+            sql = sql.limit(limit)
+        if offset is not None:
+            sql = sql.offset(offset)
+
         result = await self.execute(sql)
         if result is None:
             return None
@@ -371,9 +381,18 @@ class OngekiItemData(BaseData):
             return None
         return result.lastrowid
 
-    async def get_characters(self, aime_id: int) -> Optional[List[Dict]]:
+    async def get_characters(
+        self, aime_id: int, limit: Optional[int] = None, offset: Optional[int] = None
+    ) -> Optional[List[Row]]:
         sql = select(character).where(character.c.user == aime_id)
 
+        if limit is not None or offset is not None:
+            sql = sql.order_by(character.c.id)
+        if limit is not None:
+            sql = sql.limit(limit)
+        if offset is not None:
+            sql = sql.offset(offset)
+
         result = await self.execute(sql)
         if result is None:
             return None
@@ -479,13 +498,26 @@ class OngekiItemData(BaseData):
             return None
         return result.fetchone()
 
-    async def get_items(self, aime_id: int, item_kind: int = None) -> Optional[List[Dict]]:
-        if item_kind is None:
-            sql = select(item).where(item.c.user == aime_id)
-        else:
-            sql = select(item).where(
-                and_(item.c.user == aime_id, item.c.itemKind == item_kind)
-            )
+    async def get_items(
+        self,
+        aime_id: int,
+        item_kind: Optional[int] = None,
+        limit: Optional[int] = None,
+        offset: Optional[int] = None, 
+    ) -> Optional[List[Row]]:
+        cond = item.c.user == aime_id
+
+        if item_kind is not None:
+            cond &= item.c.itemKind == item_kind
+
+        sql = select(item).where(cond)
+
+        if limit is not None or offset is not None:
+            sql = sql.order_by(item.c.id)
+        if limit is not None:
+            sql = sql.limit(limit)
+        if offset is not None:
+            sql = sql.offset(offset)
 
         result = await self.execute(sql)
         if result is None:

titles/ongeki/schema/score.py

@@ -1,13 +1,15 @@
 from typing import Dict, List, Optional
-from sqlalchemy import Table, Column, UniqueConstraint, PrimaryKeyConstraint, and_
-from sqlalchemy.types import Integer, String, TIMESTAMP, Boolean, JSON, Float
-from sqlalchemy.schema import ForeignKey
-from sqlalchemy.sql import func, select
+
+from sqlalchemy import Column, Table, UniqueConstraint
 from sqlalchemy.dialects.mysql import insert
+from sqlalchemy.engine import Row
+from sqlalchemy.schema import ForeignKey
+from sqlalchemy.sql import select
+from sqlalchemy.types import TIMESTAMP, Boolean, Float, Integer, String
 
 from core.data.schema import BaseData, metadata
 
-score_best = Table(
+score_best: Table = Table(
     "ongeki_score_best",
     metadata,
     Column("id", Integer, primary_key=True, nullable=False),
@@ -149,8 +151,41 @@ class OngekiScoreData(BaseData):
             return None
         return result.lastrowid
 
-    async def get_best_scores(self, aime_id: int) -> Optional[List[Dict]]:
-        sql = select(score_best).where(score_best.c.user == aime_id)
+    async def get_best_scores(
+        self,
+        aime_id: int,
+        limit: Optional[int] = None,
+        offset: Optional[int] = None,
+    ) -> Optional[List[Row]]:
+        cond = score_best.c.user == aime_id
+
+        if limit is None and offset is None:
+            sql = (
+                select(score_best)
+                .where(cond)
+                .order_by(score_best.c.musicId, score_best.c.level)
+            )
+        else:
+            subq = (
+                select(score_best.c.musicId)
+                .distinct()
+                .where(cond)
+                .order_by(score_best.c.musicId)
+            )
+
+            if limit is not None:
+                subq = subq.limit(limit)
+            if offset is not None:
+                subq = subq.offset(offset)
+            
+            subq = subq.subquery()
+
+            sql = (
+                select(score_best)
+                .join(subq, score_best.c.musicId == subq.c.musicId)
+                .where(cond)
+                .order_by(score_best.c.musicId, score_best.c.level)
+            )
 
         result = await self.execute(sql)
         if result is None: