diff --git a/core/config.py b/core/config.py index dda43aa..e05323b 100644 --- a/core/config.py +++ b/core/config.py @@ -1,5 +1,9 @@ -import logging, os -from typing import Any +import logging +import os +import ssl +from typing import Any, Union + +from typing_extensions import Optional class ServerConfig: def __init__(self, parent_config: "CoreConfig") -> None: @@ -175,12 +179,60 @@ class DatabaseConfig: return CoreConfig.get_config_field( self.__config, "core", "database", "protocol", default="mysql" ) - + @property - def ssl_enabled(self) -> str: + def ssl_enabled(self) -> bool: return CoreConfig.get_config_field( self.__config, "core", "database", "ssl_enabled", default=False ) + + @property + def ssl_cafile(self) -> Optional[str]: + return CoreConfig.get_config_field( + self.__config, "core", "database", "ssl_cafile", default=None + ) + + @property + def ssl_capath(self) -> Optional[str]: + return CoreConfig.get_config_field( + self.__config, "core", "database", "ssl_capath", default=None + ) + + @property + def ssl_cert(self) -> Optional[str]: + return CoreConfig.get_config_field( + self.__config, "core", "database", "ssl_cert", default=None + ) + + @property + def ssl_key(self) -> Optional[str]: + return CoreConfig.get_config_field( + self.__config, "core", "database", "ssl_key", default=None + ) + + @property + def ssl_key_password(self) -> Optional[str]: + return CoreConfig.get_config_field( + self.__config, "core", "database", "ssl_key_password", default=None + ) + + @property + def ssl_verify_identity(self) -> bool: + return CoreConfig.get_config_field( + self.__config, "core", "database", "ssl_verify_identity", default=True + ) + + @property + def ssl_verify_cert(self) -> Optional[Union[str, bool]]: + return CoreConfig.get_config_field( + self.__config, "core", "database", "ssl_verify_cert", default=None + ) + + @property + def ssl_ciphers(self) -> Optional[str]: + return CoreConfig.get_config_field( + self.__config, "core", "database", "ssl_ciphers", default=None + ) @property def sha2_password(self) -> bool: @@ -208,6 +260,53 @@ class DatabaseConfig: self.__config, "core", "database", "memcached_host", default="localhost" ) + def create_ssl_context_if_enabled(self): + if not self.ssl_enabled: + return + + no_ca = ( + self.ssl_cafile is None + and self.ssl_capath is None + ) + + ctx = ssl.create_default_context( + cafile=self.ssl_cafile, + capath=self.ssl_capath, + ) + ctx.check_hostname = not no_ca and self.ssl_verify_identity + + if self.ssl_verify_cert is None: + ctx.verify_mode = ssl.CERT_NONE if no_ca else ssl.CERT_REQUIRED + elif isinstance(self.ssl_verify_cert, bool): + ctx.verify_mode = ( + ssl.CERT_REQUIRED + if self.ssl_verify_cert + else ssl.CERT_NONE + ) + elif isinstance(self.ssl_verify_cert, str): + value = self.ssl_verify_cert.lower() + + if value in ("none", "0", "false", "no"): + ctx.verify_mode = ssl.CERT_NONE + elif value == "optional": + ctx.verify_mode = ssl.CERT_OPTIONAL + elif value in ("required", "1", "true", "yes"): + ctx.verify_mode = ssl.CERT_REQUIRED + else: + ctx.verify_mode = ssl.CERT_NONE if no_ca else ssl.CERT_REQUIRED + + if self.ssl_cert: + ctx.load_cert_chain( + self.ssl_cert, + self.ssl_key, + self.ssl_key_password, + ) + + if self.ssl_ciphers: + ctx.set_ciphers(self.ssl_ciphers) + + return ctx + class FrontendConfig: def __init__(self, parent_config: "CoreConfig") -> None: self.__config = parent_config diff --git a/core/data/alembic/env.py b/core/data/alembic/env.py index f2a8182..b175ee6 100644 --- a/core/data/alembic/env.py +++ b/core/data/alembic/env.py @@ -1,14 +1,18 @@ from __future__ import with_statement import asyncio +import os +from pathlib import Path import threading from logging.config import fileConfig +import yaml from alembic import context from sqlalchemy import pool from sqlalchemy.engine import Connection from sqlalchemy.ext.asyncio import async_engine_from_config +from core.config import CoreConfig from core.data.schema.base import metadata # this is the Alembic Config object, which provides @@ -74,8 +78,18 @@ async def run_async_migrations() -> None: for override in overrides: ini_section[override] = overrides[override] + core_config = CoreConfig() + + with (Path("../../..") / os.environ["ARTEMIS_CFG_DIR"] / "core.yaml").open(encoding="utf-8") as f: + core_config.update(yaml.safe_load(f)) + connectable = async_engine_from_config( - ini_section, prefix="sqlalchemy.", poolclass=pool.NullPool + ini_section, + poolclass=pool.NullPool, + connect_args={ + "charset": "utf8mb4", + "ssl": core_config.database.create_ssl_context_if_enabled(), + } ) async with connectable.connect() as connection: diff --git a/core/data/database.py b/core/data/database.py index 0095a20..170665e 100644 --- a/core/data/database.py +++ b/core/data/database.py @@ -1,11 +1,12 @@ import logging import os import secrets +import ssl import string import warnings from hashlib import sha256 from logging.handlers import TimedRotatingFileHandler -from typing import ClassVar, Optional +from typing import Any, ClassVar, Optional import alembic.config import bcrypt @@ -35,12 +36,20 @@ class Data: if self.config.database.sha2_password: passwd = sha256(self.config.database.password.encode()).digest() - self.__url = f"{self.config.database.protocol}+aiomysql://{self.config.database.username}:{passwd.hex()}@{self.config.database.host}:{self.config.database.port}/{self.config.database.name}?charset=utf8mb4&ssl={str(self.config.database.ssl_enabled).lower()}" + self.__url = f"{self.config.database.protocol}+aiomysql://{self.config.database.username}:{passwd.hex()}@{self.config.database.host}:{self.config.database.port}/{self.config.database.name}" else: - self.__url = f"{self.config.database.protocol}+aiomysql://{self.config.database.username}:{self.config.database.password}@{self.config.database.host}:{self.config.database.port}/{self.config.database.name}?charset=utf8mb4&ssl={str(self.config.database.ssl_enabled).lower()}" + self.__url = f"{self.config.database.protocol}+aiomysql://{self.config.database.username}:{self.config.database.password}@{self.config.database.host}:{self.config.database.port}/{self.config.database.name}" if Data.engine is MISSING: - Data.engine = create_async_engine(self.__url, pool_recycle=3600, isolation_level="AUTOCOMMIT") + Data.engine = create_async_engine( + self.__url, + pool_recycle=3600, + isolation_level="AUTOCOMMIT", + connect_args={ + "charset": "utf8mb4", + "ssl": self.config.database.create_ssl_context_if_enabled(), + }, + ) self.__engine = Data.engine if Data.session is MISSING: diff --git a/dbutils.py b/dbutils.py index 9080afc..154df0a 100644 --- a/dbutils.py +++ b/dbutils.py @@ -9,7 +9,7 @@ import yaml from core.config import CoreConfig from core.data import Data -if __name__ == "__main__": +async def main(): parser = argparse.ArgumentParser(description="Database utilities") parser.add_argument( "--config", "-c", type=str, help="Config folder to use", default="config" @@ -44,10 +44,8 @@ if __name__ == "__main__": data = Data(cfg) - loop = asyncio.get_event_loop() - if args.action == "create": - loop.run_until_complete(data.create_database()) + await data.create_database() elif args.action == "upgrade": data.schema_upgrade(args.version) @@ -59,16 +57,20 @@ if __name__ == "__main__": data.schema_downgrade(args.version) elif args.action == "create-owner": - loop.run_until_complete(data.create_owner(args.email, args.access_code)) + await data.create_owner(args.email, args.access_code) elif args.action == "migrate": - loop.run_until_complete(data.migrate()) + await data.migrate() elif args.action == "create-revision": - loop.run_until_complete(data.create_revision(args.message)) + await data.create_revision(args.message) elif args.action == "create-autorevision": - loop.run_until_complete(data.create_revision_auto(args.message)) + await data.create_revision_auto(args.message) else: logging.getLogger("database").info(f"Unknown action {args.action}") + + +if __name__ == "__main__": + asyncio.run(main()) diff --git a/titles/chuni/base.py b/titles/chuni/base.py index 2410ef0..37bcbb3 100644 --- a/titles/chuni/base.py +++ b/titles/chuni/base.py @@ -1,16 +1,16 @@ -import logging +import itertools import json +import logging from datetime import datetime, timedelta -from time import strftime +from typing import Any, Dict, List import pytz -from typing import Dict, Any, List from core.config import CoreConfig +from titles.chuni.config import ChuniConfig from titles.chuni.const import ChuniConstants, ItemKind from titles.chuni.database import ChuniData -from titles.chuni.config import ChuniConfig -SCORE_BUFFER = {} + class ChuniBase: def __init__(self, core_cfg: CoreConfig, game_cfg: ChuniConfig) -> None: @@ -277,35 +277,39 @@ class ChuniBase: } async def handle_get_user_character_api_request(self, data: Dict) -> Dict: - characters = await self.data.item.get_characters(data["userId"]) - if characters is None: + user_id = int(data["userId"]) + next_idx = int(data["nextIndex"]) + max_ct = int(data["maxCount"]) + + # add one to the limit so we know if there's a next page of items + rows = await self.data.item.get_characters( + user_id, limit=max_ct + 1, offset=next_idx + ) + + if rows is None or len(rows) == 0: return { - "userId": data["userId"], + "userId": user_id, "length": 0, "nextIndex": -1, "userCharacterList": [], } character_list = [] - next_idx = int(data["nextIndex"]) - max_ct = int(data["maxCount"]) - for x in range(next_idx, len(characters)): - tmp = characters[x]._asdict() - tmp.pop("user") + for row in rows[:max_ct]: + tmp = row._asdict() tmp.pop("id") + tmp.pop("user") + character_list.append(tmp) - if len(character_list) >= max_ct: - break - - if len(characters) >= next_idx + max_ct: + if len(rows) > max_ct: next_idx += max_ct else: next_idx = -1 return { - "userId": data["userId"], + "userId": user_id, "length": len(character_list), "nextIndex": next_idx, "userCharacterList": character_list, @@ -335,29 +339,31 @@ class ChuniBase: } async def handle_get_user_course_api_request(self, data: Dict) -> Dict: - user_course_list = await self.data.score.get_courses(data["userId"]) - if user_course_list is None: + user_id = int(data["userId"]) + next_idx = int(data["nextIndex"]) + max_ct = int(data["maxCount"]) + + rows = await self.data.score.get_courses( + user_id, limit=max_ct + 1, offset=next_idx + ) + + if rows is None or len(rows) == 0: return { - "userId": data["userId"], + "userId": user_id, "length": 0, "nextIndex": -1, "userCourseList": [], } course_list = [] - next_idx = int(data.get("nextIndex", 0)) - max_ct = int(data.get("maxCount", 300)) - for x in range(next_idx, len(user_course_list)): - tmp = user_course_list[x]._asdict() + for row in rows[:max_ct]: + tmp = row._asdict() tmp.pop("user") tmp.pop("id") course_list.append(tmp) - if len(user_course_list) >= max_ct: - break - - if len(user_course_list) >= next_idx + max_ct: + if len(rows) > max_ct: next_idx += max_ct else: next_idx = -1 @@ -425,75 +431,94 @@ class ChuniBase: } async def handle_get_user_rival_music_api_request(self, data: Dict) -> Dict: - rival_id = data["rivalId"] - next_index = int(data["nextIndex"]) - max_count = int(data["maxCount"]) - user_rival_music_list = [] + user_id = int(data["userId"]) + rival_id = int(data["rivalId"]) + next_idx = int(data["nextIndex"]) + max_ct = int(data["maxCount"]) + rival_levels = [int(x["level"]) for x in data["userRivalMusicLevelList"]] # Fetch all the rival music entries for the user - all_entries = await self.data.score.get_rival_music(rival_id) + rows = await self.data.score.get_scores( + rival_id, + levels=rival_levels, + limit=max_ct + 1, + offset=next_idx, + ) - # Process the entries based on max_count and nextIndex - for music in all_entries: - music_id = music["musicId"] - level = music["level"] - score = music["scoreMax"] - rank = music["scoreRank"] + if rows is None or len(rows) == 0: + return { + "userId": user_id, + "rivalId": rival_id, + "nextIndex": -1, + "userRivalMusicList": [], + } - # Create a music entry for the current music_id if it's unique - music_entry = next((entry for entry in user_rival_music_list if entry["musicId"] == music_id), None) - if music_entry is None: - music_entry = { - "musicId": music_id, - "length": 0, - "userRivalMusicDetailList": [] - } - user_rival_music_list.append(music_entry) + music_details = [x._asdict() for x in rows] + returned_music_details_count = 0 + music_list = [] - # Create a level entry for the current level if it's unique or has a higher score - level_entry = next((entry for entry in music_entry["userRivalMusicDetailList"] if entry["level"] == level), None) - if level_entry is None: - level_entry = { - "level": level, - "scoreMax": score, - "scoreRank": rank - } - music_entry["userRivalMusicDetailList"].append(level_entry) - elif score > level_entry["scoreMax"]: - level_entry["scoreMax"] = score - level_entry["scoreRank"] = rank + # note that itertools.groupby will only work on sorted keys, which is already sorted by + # the query in get_scores + for music_id, details_iter in itertools.groupby(music_details, key=lambda x: x["musicId"]): + details: list[dict[Any, Any]] = [ + {"level": d["level"], "scoreMax": d["scoreMax"]} + for d in details_iter + ] - # Calculate the length for each "musicId" by counting the unique levels - for music_entry in user_rival_music_list: - music_entry["length"] = len(music_entry["userRivalMusicDetailList"]) + music_list.append({"musicId": music_id, "length": len(details), "userMusicDetailList": details}) + returned_music_details_count += len(details) - # Prepare the result dictionary with user rival music data - result = { - "userId": data["userId"], - "rivalId": data["rivalId"], - "nextIndex": str(next_index + len(user_rival_music_list[next_index: next_index + max_count]) if max_count <= len(user_rival_music_list[next_index: next_index + max_count]) else -1), - "userRivalMusicList": user_rival_music_list[next_index: next_index + max_count] + if len(music_list) >= max_ct: + break + + # if we returned fewer PBs than we originally asked for from the database, that means + # we queried for the PBs of max_ct + 1 songs. + if returned_music_details_count < len(rows): + next_idx += max_ct + else: + next_idx = -1 + + return { + "userId": user_id, + "rivalId": rival_id, + "length": len(music_list), + "nextIndex": next_idx, + "userRivalMusicList": music_list, } - return result - async def handle_get_user_favorite_item_api_request(self, data: Dict) -> Dict: + user_id = int(data["userId"]) + next_idx = int(data["nextIndex"]) + max_ct = int(data["maxCount"]) + kind = int(data["kind"]) + is_all_favorite_item = str(data["isAllFavoriteItem"]) == "true" + user_fav_item_list = [] # still needs to be implemented on WebUI # 1: Music, 2: User, 3: Character - fav_list = await self.data.item.get_all_favorites( - data["userId"], self.version, fav_kind=int(data["kind"]) + rows = await self.data.item.get_all_favorites( + user_id, + self.version, + fav_kind=kind, + limit=max_ct + 1, + offset=next_idx, ) - if fav_list is not None: - for fav in fav_list: + + if rows is not None: + for fav in rows[:max_ct]: user_fav_item_list.append({"id": fav["favId"]}) + if rows is None or len(rows) <= max_ct: + next_idx = -1 + else: + next_idx += max_ct + return { - "userId": data["userId"], + "userId": user_id, "length": len(user_fav_item_list), - "kind": data["kind"], - "nextIndex": -1, + "kind": kind, + "nextIndex": next_idx, "userFavoriteItemList": user_fav_item_list, } @@ -505,36 +530,39 @@ class ChuniBase: return {"userId": data["userId"], "length": 0, "userFavoriteMusicList": []} async def handle_get_user_item_api_request(self, data: Dict) -> Dict: - kind = int(int(data["nextIndex"]) / 10000000000) - next_idx = int(int(data["nextIndex"]) % 10000000000) - user_item_list = 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 user_item_list is None or len(user_item_list) == 0: + 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 or len(rows) == 0: return { - "userId": data["userId"], + "userId": user_id, "nextIndex": -1, "itemKind": kind, "userItemList": [], } items: List[Dict[str, Any]] = [] - for i in range(next_idx, len(user_item_list)): - tmp = user_item_list[i]._asdict() + + for row in rows[:max_ct]: + tmp = row._asdict() tmp.pop("user") tmp.pop("id") items.append(tmp) - if len(items) >= int(data["maxCount"]): - break - xout = kind * 10000000000 + next_idx + len(items) - - if len(items) < int(data["maxCount"]): - next_idx = 0 + if len(rows) > max_ct: + next_idx = kind * 10000000000 + next_idx + max_ct else: - next_idx = xout + next_idx = -1 return { - "userId": data["userId"], + "userId": user_id, "nextIndex": next_idx, "itemKind": kind, "length": len(items), @@ -586,62 +614,55 @@ class ChuniBase: } async def handle_get_user_music_api_request(self, data: Dict) -> Dict: - music_detail = await self.data.score.get_scores(data["userId"]) - if music_detail is None: + user_id = int(data["userId"]) + next_idx = int(data["nextIndex"]) + max_ct = int(data["maxCount"]) + + rows = await self.data.score.get_scores( + user_id, limit=max_ct + 1, offset=next_idx + ) + + if rows is None or len(rows) == 0: return { - "userId": data["userId"], + "userId": user_id, "length": 0, "nextIndex": -1, "userMusicList": [], # 240 } - song_list = [] - next_idx = int(data["nextIndex"]) - max_ct = int(data["maxCount"]) + music_details = [x._asdict() for x in rows] + returned_music_details_count = 0 + music_list = [] - for x in range(next_idx, len(music_detail)): - found = False - tmp = music_detail[x]._asdict() - tmp.pop("user") - tmp.pop("id") + # note that itertools.groupby will only work on sorted keys, which is already sorted by + # the query in get_scores + for _music_id, details_iter in itertools.groupby(music_details, key=lambda x: x["musicId"]): + details: list[dict[Any, Any]] = [] - for song in song_list: - score_buf = SCORE_BUFFER.get(str(data["userId"])) or [] - if song["userMusicDetailList"][0]["musicId"] == tmp["musicId"]: - found = True - song["userMusicDetailList"].append(tmp) - song["length"] = len(song["userMusicDetailList"]) - score_buf.append(tmp["musicId"]) - SCORE_BUFFER[str(data["userId"])] = score_buf + for d in details_iter: + d.pop("id") + d.pop("user") - score_buf = SCORE_BUFFER.get(str(data["userId"])) or [] - if not found and tmp["musicId"] not in score_buf: - song_list.append({"length": 1, "userMusicDetailList": [tmp]}) - score_buf.append(tmp["musicId"]) - SCORE_BUFFER[str(data["userId"])] = score_buf + details.append(d) - if len(song_list) >= max_ct: + music_list.append({"length": len(details), "userMusicDetailList": details}) + returned_music_details_count += len(details) + + if len(music_list) >= max_ct: break - - for songIdx in range(len(song_list)): - for recordIdx in range(x+1, len(music_detail)): - if song_list[songIdx]["userMusicDetailList"][0]["musicId"] == music_detail[recordIdx]["musicId"]: - music = music_detail[recordIdx]._asdict() - music.pop("user") - music.pop("id") - song_list[songIdx]["userMusicDetailList"].append(music) - song_list[songIdx]["length"] += 1 - - if len(song_list) >= max_ct: - next_idx += len(song_list) + + # if we returned fewer PBs than we originally asked for from the database, that means + # we queried for the PBs of max_ct + 1 songs. + if returned_music_details_count < len(rows): + next_idx += max_ct else: next_idx = -1 - SCORE_BUFFER[str(data["userId"])] = [] + return { - "userId": data["userId"], - "length": len(song_list), + "userId": user_id, + "length": len(music_list), "nextIndex": next_idx, - "userMusicList": song_list, # 240 + "userMusicList": music_list, } async def handle_get_user_option_api_request(self, data: Dict) -> Dict: diff --git a/titles/chuni/const.py b/titles/chuni/const.py index 68d7056..45fd498 100644 --- a/titles/chuni/const.py +++ b/titles/chuni/const.py @@ -1,4 +1,4 @@ -from enum import Enum +from enum import Enum, IntEnum class ChuniConstants: @@ -81,12 +81,31 @@ class ChuniConstants: return cls.VERSION_NAMES[ver] -class MapAreaConditionType(Enum): - UNLOCKED = 0 +class MapAreaConditionType(IntEnum): + """Condition types for the GetGameMapAreaConditionApi endpoint. Incomplete. + + For the MAP_CLEARED/MAP_AREA_CLEARED/TROPHY_OBTAINED conditions, the conditionId + is the map/map area/trophy. + + For the RANK_*/ALL_JUSTICE conditions, the conditionId is songId * 100 + difficultyId. + For example, Halcyon [ULTIMA] would be 173 * 100 + 4 = 17304. + """ + + ALWAYS_UNLOCKED = 0 + MAP_CLEARED = 1 MAP_AREA_CLEARED = 2 + TROPHY_OBTAINED = 3 + RANK_SSS = 19 + RANK_SSP = 20 + RANK_SS = 21 + RANK_SP = 22 + RANK_S = 23 + + ALL_JUSTICE = 28 + class MapAreaConditionLogicalOperator(Enum): AND = 1 @@ -102,11 +121,36 @@ class AvatarCategory(Enum): FRONT = 6 BACK = 7 -class ItemKind(Enum): +class ItemKind(IntEnum): NAMEPLATE = 1 + + FRAME = 2 + """ + "Frame" is the background for the gauge/score/max combo display + shown during gameplay. This item cannot be equipped (as of LUMINOUS) + and is hardcoded to the current game's version. + """ + TROPHY = 3 + SKILL = 4 + TICKET = 5 + """A statue is also a ticket.""" + PRESENT = 6 + MUSIC_UNLOCK = 7 MAP_ICON = 8 SYSTEM_VOICE = 9 - AVATAR_ACCESSORY = 11 \ No newline at end of file + SYMBOL_CHAT = 10 + AVATAR_ACCESSORY = 11 + + ULTIMA_UNLOCK = 12 + """This only applies to ULTIMA difficulties that are *not* unlocked by + SS-ing EXPERT+MASTER. + """ + + +class FavoriteItemKind(IntEnum): + MUSIC = 1 + RIVAL = 2 + CHARACTER = 3 diff --git a/titles/chuni/new.py b/titles/chuni/new.py index 3d3fb98..15d2b6c 100644 --- a/titles/chuni/new.py +++ b/titles/chuni/new.py @@ -4,12 +4,14 @@ from random import randint from typing import Dict import pytz + from core.config import CoreConfig from core.utils import Utils -from titles.chuni.const import ChuniConstants -from titles.chuni.database import ChuniData from titles.chuni.base import ChuniBase from titles.chuni.config import ChuniConfig +from titles.chuni.const import ChuniConstants +from titles.chuni.database import ChuniData + class ChuniNew(ChuniBase): ITEM_TYPE = {"character": 20, "story": 21, "card": 22} @@ -285,35 +287,37 @@ class ChuniNew(ChuniBase): } async def handle_get_user_printed_card_api_request(self, data: Dict) -> Dict: - user_print_list = await self.data.item.get_user_print_states( - data["userId"], has_completed=True + user_id = int(data["userId"]) + next_idx = int(data["nextIndex"]) + max_ct = int(data["maxCount"]) + + rows = await self.data.item.get_user_print_states( + user_id, + has_completed=True, + limit=max_ct + 1, + offset=next_idx, ) - if user_print_list is None: + if rows is None or len(rows) == 0: return { - "userId": data["userId"], + "userId": user_id, "length": 0, "nextIndex": -1, "userPrintedCardList": [], } print_list = [] - next_idx = int(data["nextIndex"]) - max_ct = int(data["maxCount"]) - for x in range(next_idx, len(user_print_list)): - tmp = user_print_list[x]._asdict() + for row in rows[:max_ct]: + tmp = row._asdict() print_list.append(tmp["cardId"]) - if len(print_list) >= max_ct: - break - - if len(print_list) >= max_ct: - next_idx = next_idx + max_ct + if len(rows) > max_ct: + next_idx += max_ct else: next_idx = -1 return { - "userId": data["userId"], + "userId": user_id, "length": len(print_list), "nextIndex": next_idx, "userPrintedCardList": print_list, diff --git a/titles/chuni/schema/item.py b/titles/chuni/schema/item.py index 92910da..93dcf86 100644 --- a/titles/chuni/schema/item.py +++ b/titles/chuni/schema/item.py @@ -1,22 +1,22 @@ from typing import Dict, List, Optional + from sqlalchemy import ( - Table, Column, - UniqueConstraint, PrimaryKeyConstraint, + Table, + UniqueConstraint, and_, delete, ) -from sqlalchemy.types import Integer, String, TIMESTAMP, Boolean, JSON -from sqlalchemy.engine.base import Connection -from sqlalchemy.schema import ForeignKey -from sqlalchemy.sql import func, select 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 JSON, TIMESTAMP, Boolean, Integer, String from core.data.schema import BaseData, metadata -character = Table( +character: Table = Table( "chuni_item_character", metadata, Column("id", Integer, primary_key=True, nullable=False), @@ -40,7 +40,7 @@ character = Table( mysql_charset="utf8mb4", ) -item = Table( +item: Table = Table( "chuni_item_item", metadata, Column("id", Integer, primary_key=True, nullable=False), @@ -141,7 +141,7 @@ gacha = Table( mysql_charset="utf8mb4", ) -print_state = Table( +print_state: Table = Table( "chuni_item_print_state", metadata, Column("id", Integer, primary_key=True, nullable=False), @@ -210,7 +210,7 @@ login_bonus = Table( mysql_charset="utf8mb4", ) -favorite = Table( +favorite: Table = Table( "chuni_item_favorite", metadata, Column("id", Integer, primary_key=True, nullable=False), @@ -379,9 +379,14 @@ class ChuniItemData(BaseData): return True if len(result.all()) else False async def get_all_favorites( - self, user_id: int, version: int, fav_kind: int = 1 + self, + user_id: int, + version: int, + fav_kind: int = 1, + limit: Optional[int] = None, + offset: Optional[int] = None, ) -> Optional[List[Row]]: - sql = favorite.select( + sql = select(favorite).where( and_( favorite.c.version == version, favorite.c.user == user_id, @@ -389,6 +394,13 @@ class ChuniItemData(BaseData): ) ) + if limit is not None or offset is not None: + sql = sql.order_by(favorite.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 @@ -488,9 +500,18 @@ class ChuniItemData(BaseData): return None return result.fetchone() - async def get_characters(self, user_id: int) -> Optional[List[Row]]: + async def get_characters( + self, user_id: int, limit: Optional[int] = None, offset: Optional[int] = None + ) -> Optional[List[Row]]: sql = select(character).where(character.c.user == user_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 @@ -509,13 +530,26 @@ class ChuniItemData(BaseData): return None return result.lastrowid - async def get_items(self, user_id: int, kind: int = None) -> Optional[List[Row]]: - if kind is None: - sql = select(item).where(item.c.user == user_id) - else: - sql = select(item).where( - and_(item.c.user == user_id, item.c.itemKind == kind) - ) + async def get_items( + self, + user_id: int, + kind: Optional[int] = None, + limit: Optional[int] = None, + offset: Optional[int] = None, + ) -> Optional[List[Row]]: + cond = item.c.user == user_id + + if kind is not None: + cond &= item.c.itemKind == 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: @@ -609,15 +643,26 @@ class ChuniItemData(BaseData): return result.lastrowid async def get_user_print_states( - self, aime_id: int, has_completed: bool = False + self, + aime_id: int, + has_completed: bool = False, + limit: Optional[int] = None, + offset: Optional[int] = None, ) -> Optional[List[Row]]: - sql = print_state.select( + sql = select(print_state).where( and_( print_state.c.user == aime_id, print_state.c.hasCompleted == has_completed, ) ) + if limit is not None or offset is not None: + sql = sql.order_by(print_state.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 diff --git a/titles/chuni/schema/score.py b/titles/chuni/schema/score.py index 308afa8..ab6766a 100644 --- a/titles/chuni/schema/score.py +++ b/titles/chuni/schema/score.py @@ -1,16 +1,17 @@ from typing import Dict, List, Optional -from sqlalchemy import Table, Column, UniqueConstraint, PrimaryKeyConstraint, and_ -from sqlalchemy.types import Integer, String, TIMESTAMP, Boolean, JSON, BigInteger -from sqlalchemy.engine.base import Connection -from sqlalchemy.schema import ForeignKey -from sqlalchemy.engine import Row -from sqlalchemy.sql import func, select + +from sqlalchemy import Column, Table, UniqueConstraint from sqlalchemy.dialects.mysql import insert -from sqlalchemy.sql.expression import exists +from sqlalchemy.engine import Row +from sqlalchemy.schema import ForeignKey +from sqlalchemy.sql import func, select +from sqlalchemy.types import Boolean, Integer, String + from core.data.schema import BaseData, metadata + from ..config import ChuniConfig -course = Table( +course: Table = Table( "chuni_score_course", metadata, Column("id", Integer, primary_key=True, nullable=False), @@ -41,7 +42,7 @@ course = Table( mysql_charset="utf8mb4", ) -best_score = Table( +best_score: Table = Table( "chuni_score_best", metadata, Column("id", Integer, primary_key=True, nullable=False), @@ -229,9 +230,21 @@ class ChuniRomVersion(): return -1 class ChuniScoreData(BaseData): - async def get_courses(self, aime_id: int) -> Optional[Row]: + async def get_courses( + self, + aime_id: int, + limit: Optional[int] = None, + offset: Optional[int] = None, + ) -> Optional[List[Row]]: sql = select(course).where(course.c.user == aime_id) + if limit is not None or offset is not None: + sql = sql.order_by(course.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 @@ -249,8 +262,45 @@ class ChuniScoreData(BaseData): return None return result.lastrowid - async def get_scores(self, aime_id: int) -> Optional[Row]: - sql = select(best_score).where(best_score.c.user == aime_id) + async def get_scores( + self, + aime_id: int, + levels: Optional[list[int]] = None, + limit: Optional[int] = None, + offset: Optional[int] = None, + ) -> Optional[List[Row]]: + condition = best_score.c.user == aime_id + + if levels is not None: + condition &= best_score.c.level.in_(levels) + + if limit is None and offset is None: + sql = ( + select(best_score) + .where(condition) + .order_by(best_score.c.musicId.asc(), best_score.c.level.asc()) + ) + else: + subq = ( + select(best_score.c.musicId) + .distinct() + .where(condition) + .order_by(best_score.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(best_score) + .join(subq, best_score.c.musicId == subq.c.musicId) + .where(condition) + .order_by(best_score.c.musicId, best_score.c.level) + ) result = await self.execute(sql) if result is None: @@ -360,11 +410,3 @@ class ChuniScoreData(BaseData): rows = result.fetchall() return [dict(row) for row in rows] - - async def get_rival_music(self, rival_id: int) -> Optional[List[Dict]]: - sql = select(best_score).where(best_score.c.user == rival_id) - - result = await self.execute(sql) - if result is None: - return None - return result.fetchall() diff --git a/titles/mai2/base.py b/titles/mai2/base.py index 1c9d91c..add3ed2 100644 --- a/titles/mai2/base.py +++ b/titles/mai2/base.py @@ -1,16 +1,17 @@ -from datetime import datetime, timedelta -from typing import Any, Dict, List +import itertools import logging from base64 import b64decode -from os import path, stat, remove, mkdir, access, W_OK -from PIL import ImageFile -from random import randint +from datetime import datetime, timedelta +from os import W_OK, access, mkdir, path +from typing import Any, Dict, List import pytz + from core.config import CoreConfig from core.utils import Utils -from .const import Mai2Constants + from .config import Mai2Config +from .const import Mai2Constants from .database import Mai2Data @@ -444,23 +445,22 @@ class Mai2Base: return {"userId": data["userId"], "userOption": options_dict} async def handle_get_user_card_api_request(self, data: Dict) -> Dict: - user_cards = await self.data.item.get_cards(data["userId"]) - if user_cards is None: - return {"userId": data["userId"], "nextIndex": 0, "userCardList": []} + user_id = int(data["userId"]) + next_idx = int(data["nextIndex"]) + max_ct = int(data["maxCount"]) - max_ct = data["maxCount"] - next_idx = data["nextIndex"] - start_idx = next_idx - end_idx = max_ct + start_idx + user_cards = await self.data.item.get_cards( + user_id, limit=max_ct + 1, offset=next_idx + ) - if len(user_cards[start_idx:]) > max_ct: - next_idx += max_ct - else: - next_idx = 0 + if user_cards is None or len(user_cards) == 0: + return {"userId": user_id, "nextIndex": 0, "userCardList": []} card_list = [] - for card in user_cards: + + for card in user_cards[:max_ct]: tmp = card._asdict() + tmp.pop("id") tmp.pop("user") tmp["startDate"] = datetime.strftime( @@ -469,12 +469,18 @@ class Mai2Base: tmp["endDate"] = datetime.strftime( tmp["endDate"], Mai2Constants.DATE_TIME_FORMAT ) + card_list.append(tmp) + if len(user_cards) > max_ct: + next_idx += max_ct + else: + next_idx = 0 + return { - "userId": data["userId"], + "userId": user_id, "nextIndex": next_idx, - "userCardList": card_list[start_idx:end_idx], + "userCardList": card_list, } async def handle_get_user_charge_api_request(self, data: Dict) -> Dict: @@ -536,28 +542,35 @@ class Mai2Base: return { "userId": data.get("userId", 0), "userBossData": boss_lst} async def handle_get_user_item_api_request(self, data: Dict) -> Dict: - kind = int(data["nextIndex"] / 10000000000) - next_idx = int(data["nextIndex"] % 10000000000) - user_item_list = await self.data.item.get_items(data["userId"], kind) + user_id: int = data["userId"] + kind: int = data["nextIndex"] // 10000000000 + next_idx: int = data["nextIndex"] % 10000000000 + max_ct: int = data["maxCount"] + rows = await self.data.item.get_items(user_id, kind, limit=max_ct, offset=next_idx) + + if rows is None or len(rows) == 0: + return { + "userId": user_id, + "nextIndex": 0, + "itemKind": kind, + "userItemList": [], + } items: List[Dict[str, Any]] = [] - for i in range(next_idx, len(user_item_list)): - tmp = user_item_list[i]._asdict() + + for row in rows[:max_ct]: + tmp = row._asdict() tmp.pop("user") tmp.pop("id") items.append(tmp) - if len(items) >= int(data["maxCount"]): - break - xout = kind * 10000000000 + next_idx + len(items) - - if len(items) < int(data["maxCount"]): - next_idx = 0 + if len(rows) > max_ct: + next_idx = kind * 10000000000 + next_idx + max_ct else: - next_idx = xout + next_idx = 0 return { - "userId": data["userId"], + "userId": user_id, "nextIndex": next_idx, "itemKind": kind, "userItemList": items, @@ -675,77 +688,90 @@ class Mai2Base: return {"length": 0, "userPortraitList": []} async def handle_get_user_friend_season_ranking_api_request(self, data: Dict) -> Dict: - friend_season_ranking = await self.data.item.get_friend_season_ranking(data["userId"]) - if friend_season_ranking is None: + user_id: int = data["userId"] + next_idx: int = data["nextIndex"] + max_ct: int = data["maxCount"] + + rows = await self.data.item.get_friend_season_ranking( + user_id, limit=max_ct + 1, offset=next_idx + ) + + if rows is None: return { - "userId": data["userId"], + "userId": user_id, "nextIndex": 0, "userFriendSeasonRankingList": [], } friend_season_ranking_list = [] - next_idx = int(data["nextIndex"]) - max_ct = int(data["maxCount"]) - for x in range(next_idx, len(friend_season_ranking)): - tmp = friend_season_ranking[x]._asdict() - tmp.pop("user") + for row in rows[:max_ct]: + tmp = row._asdict() + tmp.pop("id") + tmp.pop("user") tmp["recordDate"] = datetime.strftime( tmp["recordDate"], f"{Mai2Constants.DATE_TIME_FORMAT}.0" ) + friend_season_ranking_list.append(tmp) - if len(friend_season_ranking_list) >= max_ct: - break - - if len(friend_season_ranking) >= next_idx + max_ct: + if len(rows) > max_ct: next_idx += max_ct else: next_idx = 0 return { - "userId": data["userId"], + "userId": user_id, "nextIndex": next_idx, "userFriendSeasonRankingList": friend_season_ranking_list, } async def handle_get_user_map_api_request(self, data: Dict) -> Dict: - maps = await self.data.item.get_maps(data["userId"]) - if maps is None: + user_id: int = data["userId"] + next_idx: int = data["nextIndex"] + max_ct: int = data["maxCount"] + + rows = await self.data.item.get_maps( + user_id, limit=max_ct + 1, offset=next_idx, + ) + + if rows is None: return { - "userId": data["userId"], + "userId": user_id, "nextIndex": 0, "userMapList": [], } map_list = [] - next_idx = int(data["nextIndex"]) - max_ct = int(data["maxCount"]) - for x in range(next_idx, len(maps)): - tmp = maps[x]._asdict() + for row in rows[:max_ct]: + tmp = row._asdict() tmp.pop("user") tmp.pop("id") map_list.append(tmp) - if len(map_list) >= max_ct: - break - - if len(maps) >= next_idx + max_ct: + if len(rows) > max_ct: next_idx += max_ct else: next_idx = 0 return { - "userId": data["userId"], + "userId": user_id, "nextIndex": next_idx, "userMapList": map_list, } async def handle_get_user_login_bonus_api_request(self, data: Dict) -> Dict: - login_bonuses = await self.data.item.get_login_bonuses(data["userId"]) - if login_bonuses is None: + user_id: int = data["userId"] + next_idx: int = data["nextIndex"] + max_ct: int = data["maxCount"] + + rows = await self.data.item.get_login_bonuses( + user_id, limit=max_ct + 1, offset=next_idx + ) + + if rows is None: return { "userId": data["userId"], "nextIndex": 0, @@ -753,25 +779,20 @@ class Mai2Base: } login_bonus_list = [] - next_idx = int(data["nextIndex"]) - max_ct = int(data["maxCount"]) - for x in range(next_idx, len(login_bonuses)): - tmp = login_bonuses[x]._asdict() + for row in rows[:max_ct]: + tmp = row._asdict() tmp.pop("user") tmp.pop("id") login_bonus_list.append(tmp) - if len(login_bonus_list) >= max_ct: - break - - if len(login_bonuses) >= next_idx + max_ct: + if len(rows) > max_ct: next_idx += max_ct else: next_idx = 0 return { - "userId": data["userId"], + "userId": user_id, "nextIndex": next_idx, "userLoginBonusList": login_bonus_list, } @@ -805,42 +826,54 @@ class Mai2Base: return {"userId": data["userId"], "userGradeStatus": grade_stat, "length": 0, "userGradeList": []} async def handle_get_user_music_api_request(self, data: Dict) -> Dict: - user_id = data.get("userId", 0) - next_index = data.get("nextIndex", 0) - max_ct = data.get("maxCount", 50) - upper_lim = next_index + max_ct - music_detail_list = [] + user_id: int = data.get("userId", 0) + next_idx: int = data.get("nextIndex", 0) + max_ct: int = data.get("maxCount", 50) if user_id <= 0: self.logger.warning("handle_get_user_music_api_request: Could not find userid in data, or userId is 0") return {} - songs = await self.data.score.get_best_scores(user_id, is_dx=False) - if songs is None: + rows = await self.data.score.get_best_scores( + user_id, is_dx=False, limit=max_ct + 1, offset=next_idx + ) + + if rows is None: self.logger.debug("handle_get_user_music_api_request: get_best_scores returned None!") return { - "userId": data["userId"], - "nextIndex": 0, - "userMusicList": [], - } + "userId": user_id, + "nextIndex": 0, + "userMusicList": [], + } - num_user_songs = len(songs) + music_details = [row._asdict() for row in rows] + returned_count = 0 + music_list = [] - for x in range(next_index, upper_lim): - if num_user_songs <= x: + 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({"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 = 0 - tmp = songs[x]._asdict() - tmp.pop("id") - tmp.pop("user") - music_detail_list.append(tmp) - - next_index = 0 if len(music_detail_list) < max_ct or num_user_songs == upper_lim else upper_lim - self.logger.info(f"Send songs {next_index}-{upper_lim} ({len(music_detail_list)}) out of {num_user_songs} for user {user_id} (next idx {next_index})") return { - "userId": data["userId"], - "nextIndex": next_index, - "userMusicList": [{"userMusicDetailList": music_detail_list}], + "userId": user_id, + "nextIndex": next_idx, + "userMusicList": music_list, } async def handle_upload_user_portrait_api_request(self, data: Dict) -> Dict: @@ -925,30 +958,52 @@ class Mai2Base: async def handle_get_user_favorite_item_api_request(self, data: Dict) -> Dict: user_id = data.get("userId", 0) kind = data.get("kind", 0) # 1 is fav music, 2 is rival user IDs - next_index = data.get("nextIndex", 0) + next_idx = data.get("nextIndex", 0) max_ct = data.get("maxCount", 100) # always 100 is_all = data.get("isAllFavoriteItem", False) # always false + + empty_resp = { + "userId": user_id, + "kind": kind, + "nextIndex": 0, + "userFavoriteItemList": [], + } + + if not user_id or kind not in (1, 2): + return empty_resp + id_list: List[Dict] = [] - if user_id: - if kind == 1: - fav_music = await self.data.item.get_fav_music(user_id) - if fav_music: - for fav in fav_music: - id_list.append({"orderId": fav["orderId"] or 0, "id": fav["musicId"]}) - if len(id_list) >= 100: # Lazy but whatever - break - - elif kind == 2: - rivals = await self.data.profile.get_rivals_game(user_id) - if rivals: - for rival in rivals: - id_list.append({"orderId": 0, "id": rival["rival"]}) + if kind == 1: + rows = await self.data.item.get_fav_music( + user_id, limit=max_ct + 1, offset=next_idx + ) + + if rows is None: + return empty_resp + + for row in rows[:max_ct]: + id_list.append({"orderId": row["orderId"] or 0, "id": row["musicId"]}) + elif kind == 2: + rows = await self.data.profile.get_rivals_game( + user_id, limit=max_ct + 1, offset=next_idx + ) + + if rows is None: + return empty_resp + + for row in rows[:max_ct]: + id_list.append({"orderId": 0, "id": row["rival"]}) + + if rows is None or len(rows) <= max_ct: + next_idx = 0 + else: + next_idx += max_ct return { "userId": user_id, "kind": kind, - "nextIndex": 0, + "nextIndex": next_idx, "userFavoriteItemList": id_list, } @@ -964,5 +1019,4 @@ class Mai2Base: """ return {"userId": data["userId"], "userRecommendSelectionMusicIdList": []} async def handle_get_user_score_ranking_api_request(self, data: Dict) ->Dict: - - return {"userId": data["userId"], "userScoreRanking": []} \ No newline at end of file + return {"userId": data["userId"], "userScoreRanking": []} diff --git a/titles/mai2/dx.py b/titles/mai2/dx.py index 66bf914..b37a3f4 100644 --- a/titles/mai2/dx.py +++ b/titles/mai2/dx.py @@ -1,8 +1,9 @@ -from typing import Any, List, Dict +import itertools from datetime import datetime, timedelta -import pytz -import json from random import randint +from typing import Any, Dict, List + +import pytz from core.config import CoreConfig from core.utils import Utils @@ -309,83 +310,112 @@ class Mai2DX(Mai2Base): return {"userId": data["userId"], "userOption": options_dict} async def handle_get_user_card_api_request(self, data: Dict) -> Dict: - user_cards = await self.data.item.get_cards(data["userId"]) - if user_cards is None: - return {"userId": data["userId"], "nextIndex": 0, "userCardList": []} + user_id: int = data["userId"] + next_idx: int = data["nextIndex"] + max_ct: int = data["maxCount"] + + rows = await self.data.item.get_cards(user_id, limit=max_ct + 1, offset=next_idx) + + if rows is None: + return {"userId": user_id, "nextIndex": 0, "userCardList": []} - max_ct = data["maxCount"] - next_idx = data["nextIndex"] - start_idx = next_idx - end_idx = max_ct + start_idx + card_list = [] - if len(user_cards[start_idx:]) > max_ct: + for row in rows[:max_ct]: + card = row._asdict() + card.pop("id") + card.pop("user") + card["startDate"] = datetime.strftime( + card["startDate"], Mai2Constants.DATE_TIME_FORMAT + ) + card["endDate"] = datetime.strftime( + card["endDate"], Mai2Constants.DATE_TIME_FORMAT + ) + card_list.append(card) + + if len(rows) > max_ct: next_idx += max_ct else: next_idx = 0 - card_list = [] - for card in user_cards: - tmp = card._asdict() - tmp.pop("id") - tmp.pop("user") - tmp["startDate"] = datetime.strftime( - tmp["startDate"], Mai2Constants.DATE_TIME_FORMAT - ) - tmp["endDate"] = datetime.strftime( - tmp["endDate"], Mai2Constants.DATE_TIME_FORMAT - ) - card_list.append(tmp) - return { "userId": data["userId"], "nextIndex": next_idx, - "userCardList": card_list[start_idx:end_idx], + "userCardList": card_list, } async def handle_get_user_item_api_request(self, data: Dict) -> Dict: - kind = data["nextIndex"] // 10000000000 - next_idx = data["nextIndex"] % 10000000000 + user_id: int = data["userId"] + next_idx: int = data["nextIndex"] + max_ct: int = data["maxCount"] + + kind = next_idx // 10000000000 + next_idx = next_idx % 10000000000 + items: List[Dict[str, Any]] = [] if kind == 4: # presents - user_pres_list = await self.data.item.get_presents_by_version_user(self.version, data["userId"]) - if user_pres_list: - self.logger.debug(f"Found {len(user_pres_list)} possible presents") - for present in user_pres_list: - if (present['startDate'] and present['startDate'].timestamp() > datetime.now().timestamp()): - self.logger.debug(f"Present {present['id']} distribution hasn't started yet (begins {present['startDate']})") - continue # present period hasn't started yet, move onto the next one - - if (present['endDate'] and present['endDate'].timestamp() < datetime.now().timestamp()): - self.logger.warn(f"Present {present['id']} ended on {present['endDate']} and should be removed") - continue # present period ended, move onto the next one - - test = await self.data.item.get_item(data["userId"], present['itemKind'], present['itemId']) - if not test: # Don't send presents for items the user already has - pres_id = present['itemKind'] * 1000000 - pres_id += present['itemId'] - items.append({"itemId": pres_id, "itemKind": 4, "stock": present['stock'], "isValid": True}) - self.logger.info(f"Give user {data['userId']} {present['stock']}x item {present['itemId']} (kind {present['itemKind']}) as present") + rows = await self.data.item.get_presents_by_version_user( + version=self.version, + user_id=user_id, + exclude_owned=True, + exclude_not_in_present_period=True, + limit=max_ct + 1, + offset=next_idx, + ) + if rows is None: + return { + "userId": user_id, + "nextIndex": 0, + "itemKind": kind, + "userItemList": [], + } + + for row in rows[:max_ct]: + self.logger.info( + f"Give user {user_id} {row['stock']}x item {row['itemId']} (kind {row['itemKind']}) as present" + ) + + items.append( + { + "itemId": row["itemKind"] * 1000000 + row["itemId"], + "itemKind": kind, + "stock": row["stock"], + "isValid": True, + } + ) else: - user_item_list = await self.data.item.get_items(data["userId"], kind) - for i in range(next_idx, len(user_item_list)): - tmp = user_item_list[i]._asdict() - tmp.pop("user") - tmp.pop("id") - items.append(tmp) - if len(items) >= int(data["maxCount"]): - break + rows = await self.data.item.get_items( + user_id=user_id, + item_kind=kind, + limit=max_ct + 1, + offset=next_idx, + ) - xout = kind * 10000000000 + next_idx + len(items) + if rows is None: + return { + "userId": user_id, + "nextIndex": 0, + "itemKind": kind, + "userItemList": [], + } - if len(items) < int(data["maxCount"]): + for row in rows[:max_ct]: + item = row._asdict() + + item.pop("id") + item.pop("user") + + items.append(item) + + if len(rows) > max_ct: + next_idx = kind * 10000000000 + next_idx + max_ct + else: next_idx = 0 - else: - next_idx = xout return { - "userId": data["userId"], + "userId": user_id, "nextIndex": next_idx, "itemKind": kind, "userItemList": items, @@ -491,103 +521,115 @@ class Mai2DX(Mai2Base): return {"length": 0, "userPortraitList": []} async def handle_get_user_friend_season_ranking_api_request(self, data: Dict) -> Dict: - friend_season_ranking = await self.data.item.get_friend_season_ranking(data["userId"]) - if friend_season_ranking is None: + user_id: int = data["userId"] + next_idx: int = data["nextIndex"] + max_ct: int = data["maxCount"] + + rows = await self.data.item.get_friend_season_ranking( + user_id, limit=max_ct + 1, offset=next_idx + ) + + if rows is None: return { - "userId": data["userId"], + "userId": user_id, "nextIndex": 0, "userFriendSeasonRankingList": [], } friend_season_ranking_list = [] - next_idx = int(data["nextIndex"]) - max_ct = int(data["maxCount"]) - for x in range(next_idx, len(friend_season_ranking)): - tmp = friend_season_ranking[x]._asdict() - tmp.pop("user") - tmp.pop("id") - tmp["recordDate"] = datetime.strftime( - tmp["recordDate"], f"{Mai2Constants.DATE_TIME_FORMAT}.0" + for row in rows[:max_ct]: + friend_season_ranking = row._asdict() + + friend_season_ranking.pop("user") + friend_season_ranking.pop("id") + friend_season_ranking["recordDate"] = datetime.strftime( + friend_season_ranking["recordDate"], f"{Mai2Constants.DATE_TIME_FORMAT}.0" ) - friend_season_ranking_list.append(tmp) + + friend_season_ranking_list.append(friend_season_ranking) - if len(friend_season_ranking_list) >= max_ct: - break - - if len(friend_season_ranking) >= next_idx + max_ct: + if len(rows) > max_ct: next_idx += max_ct else: next_idx = 0 return { - "userId": data["userId"], + "userId": user_id, "nextIndex": next_idx, "userFriendSeasonRankingList": friend_season_ranking_list, } async def handle_get_user_map_api_request(self, data: Dict) -> Dict: - maps = await self.data.item.get_maps(data["userId"]) - if maps is None: + user_id: int = data["userId"] + next_idx: int = data["nextIndex"] + max_ct: int = data["maxCount"] + + rows = await self.data.item.get_maps( + user_id, limit=max_ct + 1, offset=next_idx + ) + + if rows is None: return { - "userId": data["userId"], + "userId": user_id, "nextIndex": 0, "userMapList": [], } map_list = [] - next_idx = int(data["nextIndex"]) - max_ct = int(data["maxCount"]) - for x in range(next_idx, len(maps)): - tmp = maps[x]._asdict() - tmp.pop("user") - tmp.pop("id") - map_list.append(tmp) + for row in rows[:max_ct]: + map = row._asdict() + + map.pop("user") + map.pop("id") + + map_list.append(map) - if len(map_list) >= max_ct: - break - - if len(maps) >= next_idx + max_ct: + if len(rows) > max_ct: next_idx += max_ct else: next_idx = 0 return { - "userId": data["userId"], + "userId": user_id, "nextIndex": next_idx, "userMapList": map_list, } async def handle_get_user_login_bonus_api_request(self, data: Dict) -> Dict: - login_bonuses = await self.data.item.get_login_bonuses(data["userId"]) - if login_bonuses is None: + user_id: int = data["userId"] + next_idx: int = data["nextIndex"] + max_ct: int = data["maxCount"] + + rows = await self.data.item.get_login_bonuses( + user_id, limit=max_ct + 1, offset=next_idx + ) + + if rows is None: return { - "userId": data["userId"], + "userId": user_id, "nextIndex": 0, "userLoginBonusList": [], } login_bonus_list = [] - next_idx = int(data["nextIndex"]) - max_ct = int(data["maxCount"]) - for x in range(next_idx, len(login_bonuses)): - tmp = login_bonuses[x]._asdict() - tmp.pop("user") - tmp.pop("id") - login_bonus_list.append(tmp) + for row in rows[:max_ct]: + login_bonus = row._asdict() + + login_bonus.pop("user") + login_bonus.pop("id") + + login_bonus_list.append(login_bonus) - if len(login_bonus_list) >= max_ct: - break - - if len(login_bonuses) >= next_idx + max_ct: + if len(rows) > max_ct: next_idx += max_ct else: next_idx = 0 return { - "userId": data["userId"], + "userId": user_id, "nextIndex": next_idx, "userLoginBonusList": login_bonus_list, } @@ -619,46 +661,62 @@ class Mai2DX(Mai2Base): } async def handle_get_user_rival_music_api_request(self, data: Dict) -> Dict: - user_id = data.get("userId", 0) - rival_id = data.get("rivalId", 0) - next_index = data.get("nextIndex", 0) - max_ct = 100 - upper_lim = next_index + max_ct - rival_music_list: Dict[int, List] = {} + user_id: int = data["userId"] + rival_id: int = data["rivalId"] + next_idx: int = data["nextIndex"] + max_ct: int = 100 + levels: list[int] = [x["level"] for x in data["userRivalMusicLevelList"]] - songs = await self.data.score.get_best_scores(rival_id) - if songs is None: + rows = await self.data.score.get_best_scores( + rival_id, + is_dx=True, + limit=max_ct + 1, + offset=next_idx, + levels=levels, + ) + + if rows is None: self.logger.debug("handle_get_user_rival_music_api_request: get_best_scores returned None!") + return { "userId": user_id, "rivalId": rival_id, "nextIndex": 0, "userRivalMusicList": [] # musicId userRivalMusicDetailList -> level achievement deluxscoreMax } + + music_details = [x._asdict() for x in rows] + returned_count = 0 + music_list = [] - num_user_songs = len(songs) + for music_id, details_iter in itertools.groupby(music_details, key=lambda x: x["musicId"]): + details: list[dict[Any, Any]] = [] - for x in range(next_index, upper_lim): - if x >= num_user_songs: + for d in details_iter: + details.append( + { + "level": d["level"], + "achievement": d["achievement"], + "deluxscoreMax": d["deluxscoreMax"], + } + ) + + music_list.append({"musicId": music_id, "userRivalMusicDetailList": details}) + returned_count += len(details) + + if len(music_list) >= max_ct: break - tmp = songs[x]._asdict() - if tmp['musicId'] in rival_music_list: - rival_music_list[tmp['musicId']].append([{"level": tmp['level'], 'achievement': tmp['achievement'], 'deluxscoreMax': tmp['deluxscoreMax']}]) - - else: - if len(rival_music_list) >= max_ct: - break - rival_music_list[tmp['musicId']] = [{"level": tmp['level'], 'achievement': tmp['achievement'], 'deluxscoreMax': tmp['deluxscoreMax']}] - - next_index = 0 if len(rival_music_list) < max_ct or num_user_songs == upper_lim else upper_lim - self.logger.info(f"Send rival {rival_id} songs {next_index}-{upper_lim} ({len(rival_music_list)}) out of {num_user_songs} for user {user_id} (next idx {next_index})") + if returned_count < len(rows): + next_idx += max_ct + else: + next_idx = 0 return { "userId": user_id, "rivalId": rival_id, - "nextIndex": next_index, - "userRivalMusicList": [{"musicId": x, "userRivalMusicDetailList": y} for x, y in rival_music_list.items()] + "nextIndex": next_idx, + "userRivalMusicList": music_list, } async def handle_get_user_new_item_api_request(self, data: Dict) -> Dict: @@ -674,42 +732,55 @@ class Mai2DX(Mai2Base): } async def handle_get_user_music_api_request(self, data: Dict) -> Dict: - user_id = data.get("userId", 0) - next_index = data.get("nextIndex", 0) - max_ct = data.get("maxCount", 50) - upper_lim = next_index + max_ct - music_detail_list = [] + user_id: int = data.get("userId", 0) + next_idx: int = data.get("nextIndex", 0) + max_ct: int = data.get("maxCount", 50) if user_id <= 0: self.logger.warning("handle_get_user_music_api_request: Could not find userid in data, or userId is 0") return {} - songs = await self.data.score.get_best_scores(user_id) - if songs is None: + rows = await self.data.score.get_best_scores( + user_id, is_dx=True, limit=max_ct + 1, offset=next_idx + ) + + if rows is None: self.logger.debug("handle_get_user_music_api_request: get_best_scores returned None!") + return { - "userId": data["userId"], - "nextIndex": 0, - "userMusicList": [], - } + "userId": user_id, + "nextIndex": 0, + "userMusicList": [], + } - num_user_songs = len(songs) + music_details = [row._asdict() for row in rows] + returned_count = 0 + music_list = [] - for x in range(next_index, upper_lim): - if num_user_songs <= x: + 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({"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 = 0 - tmp = songs[x]._asdict() - tmp.pop("id") - tmp.pop("user") - music_detail_list.append(tmp) - - next_index = 0 if len(music_detail_list) < max_ct or num_user_songs == upper_lim else upper_lim - self.logger.info(f"Send songs {next_index}-{upper_lim} ({len(music_detail_list)}) out of {num_user_songs} for user {user_id} (next idx {next_index})") return { - "userId": data["userId"], - "nextIndex": next_index, - "userMusicList": [{"userMusicDetailList": music_detail_list}], + "userId": user_id, + "nextIndex": next_idx, + "userMusicList": music_list, } async def handle_user_login_api_request(self, data: Dict) -> Dict: @@ -812,39 +883,43 @@ class Mai2DX(Mai2Base): return {"length": len(selling_card_list), "sellingCardList": selling_card_list} 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"] + next_idx: int = data["nextIndex"] + max_ct: int = data["maxCount"] + + rows = await self.data.item.get_cards( + user_id, limit=max_ct + 1, offset=next_idx + ) + + if rows is None: return {"returnCode": 1, "length": 0, "nextIndex": 0, "userCardList": []} - max_ct = data["maxCount"] - next_idx = data["nextIndex"] - start_idx = next_idx - end_idx = max_ct + start_idx + card_list = [] - if len(user_cards[start_idx:]) > max_ct: + for row in rows[:max_ct]: + card = row._asdict() + + card.pop("id") + card.pop("user") + card["startDate"] = datetime.strftime( + card["startDate"], Mai2Constants.DATE_TIME_FORMAT + ) + card["endDate"] = datetime.strftime( + card["endDate"], Mai2Constants.DATE_TIME_FORMAT + ) + + card_list.append(card) + + if len(rows) > max_ct: next_idx += max_ct else: next_idx = 0 - card_list = [] - for card in user_cards: - tmp = card._asdict() - tmp.pop("id") - tmp.pop("user") - - tmp["startDate"] = datetime.strftime( - tmp["startDate"], Mai2Constants.DATE_TIME_FORMAT - ) - tmp["endDate"] = datetime.strftime( - tmp["endDate"], Mai2Constants.DATE_TIME_FORMAT - ) - card_list.append(tmp) - return { "returnCode": 1, - "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_item_api_request(self, data: Dict) -> Dict: diff --git a/titles/mai2/schema/item.py b/titles/mai2/schema/item.py index 87ddca4..3b7d8d4 100644 --- a/titles/mai2/schema/item.py +++ b/titles/mai2/schema/item.py @@ -1,15 +1,16 @@ -from core.data.schema import BaseData, metadata - from datetime import datetime -from typing import Optional, Dict, List -from sqlalchemy import Table, Column, UniqueConstraint, PrimaryKeyConstraint, and_, or_ -from sqlalchemy.types import Integer, String, TIMESTAMP, Boolean, JSON, BIGINT, INTEGER -from sqlalchemy.schema import ForeignKey -from sqlalchemy.sql import func, select +from typing import Dict, List, Optional + +from sqlalchemy import Column, Table, UniqueConstraint, and_, or_ 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 BIGINT, INTEGER, JSON, TIMESTAMP, Boolean, Integer, String -character = Table( +from core.data.schema import BaseData, metadata + +character: Table = Table( "mai2_item_character", metadata, Column("id", Integer, primary_key=True, nullable=False), @@ -27,7 +28,7 @@ character = Table( mysql_charset="utf8mb4", ) -card = Table( +card: Table = Table( "mai2_item_card", metadata, Column("id", Integer, primary_key=True, nullable=False), @@ -46,7 +47,7 @@ card = Table( mysql_charset="utf8mb4", ) -item = Table( +item: Table = Table( "mai2_item_item", metadata, Column("id", Integer, primary_key=True, nullable=False), @@ -63,7 +64,7 @@ item = Table( mysql_charset="utf8mb4", ) -map = Table( +map: Table = Table( "mai2_item_map", metadata, Column("id", Integer, primary_key=True, nullable=False), @@ -81,7 +82,7 @@ map = Table( mysql_charset="utf8mb4", ) -login_bonus = Table( +login_bonus: Table = Table( "mai2_item_login_bonus", metadata, Column("id", Integer, primary_key=True, nullable=False), @@ -98,7 +99,7 @@ login_bonus = Table( mysql_charset="utf8mb4", ) -friend_season_ranking = Table( +friend_season_ranking: Table = Table( "mai2_item_friend_season_ranking", metadata, Column("id", Integer, primary_key=True, nullable=False), @@ -134,7 +135,7 @@ favorite = Table( mysql_charset="utf8mb4", ) -fav_music = Table( +fav_music: Table = Table( "mai2_item_favorite_music", metadata, Column("id", Integer, primary_key=True, nullable=False), @@ -199,7 +200,7 @@ print_detail = Table( mysql_charset="utf8mb4", ) -present = Table( +present: Table = Table( "mai2_item_present", metadata, Column('id', BIGINT, primary_key=True, nullable=False), @@ -239,13 +240,26 @@ class Mai2ItemData(BaseData): return None return result.lastrowid - async def get_items(self, user_id: int, item_kind: int = None) -> Optional[List[Row]]: - if item_kind is None: - sql = item.select(item.c.user == user_id) - else: - sql = item.select( - and_(item.c.user == user_id, item.c.itemKind == item_kind) - ) + async def get_items( + self, + user_id: int, + item_kind: Optional[int] = None, + limit: Optional[int] = None, + offset: Optional[int] = None, + ) -> Optional[List[Row]]: + cond = item.c.user == user_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: @@ -296,8 +310,20 @@ class Mai2ItemData(BaseData): return None return result.lastrowid - async def get_login_bonuses(self, user_id: int) -> Optional[List[Row]]: - sql = login_bonus.select(login_bonus.c.user == user_id) + async def get_login_bonuses( + self, + user_id: int, + limit: Optional[int] = None, + offset: Optional[int] = None, + ) -> Optional[List[Row]]: + sql = select(login_bonus).where(login_bonus.c.user == user_id) + + if limit is not None or offset is not None: + sql = sql.order_by(login_bonus.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: @@ -347,8 +373,20 @@ class Mai2ItemData(BaseData): return None return result.lastrowid - async def get_maps(self, user_id: int) -> Optional[List[Row]]: - sql = map.select(map.c.user == user_id) + async def get_maps( + self, + user_id: int, + limit: Optional[int] = None, + offset: Optional[int] = None, + ) -> Optional[List[Row]]: + sql = select(map).where(map.c.user == user_id) + + if limit is not None or offset is not None: + sql = sql.order_by(map.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: @@ -424,8 +462,20 @@ class Mai2ItemData(BaseData): return None return result.fetchone() - async def get_friend_season_ranking(self, user_id: int) -> Optional[Row]: - sql = friend_season_ranking.select(friend_season_ranking.c.user == user_id) + async def get_friend_season_ranking( + self, + user_id: int, + limit: Optional[int] = None, + offset: Optional[int] = None, + ) -> Optional[List[Row]]: + sql = select(friend_season_ranking).where(friend_season_ranking.c.user == user_id) + + if limit is not None or offset is not None: + sql = sql.order_by(friend_season_ranking.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: @@ -480,8 +530,23 @@ class Mai2ItemData(BaseData): return None return result.fetchall() - async def get_fav_music(self, user_id: int) -> Optional[List[Row]]: - result = await self.execute(fav_music.select(fav_music.c.user == user_id)) + async def get_fav_music( + self, + user_id: int, + limit: Optional[int] = None, + offset: Optional[int] = None, + ) -> Optional[List[Row]]: + sql = select(fav_music).where(fav_music.c.user == user_id) + + if limit is not None or offset is not None: + sql = sql.order_by(fav_music.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: return result.fetchall() @@ -537,13 +602,24 @@ class Mai2ItemData(BaseData): return None return result.lastrowid - async def get_cards(self, user_id: int, kind: int = None) -> Optional[Row]: - if kind is None: - sql = card.select(card.c.user == user_id) - else: - sql = card.select(and_(card.c.user == user_id, card.c.cardKind == kind)) + async def get_cards( + self, + user_id: int, + kind: Optional[int] = None, + limit: Optional[int] = None, + offset: Optional[int] = None, + ) -> Optional[List[Row]]: + condition = card.c.user == user_id - sql = sql.order_by(card.c.startDate.desc()) + if kind is not None: + condition &= card.c.cardKind == kind + + sql = select(card).where(condition).order_by(card.c.startDate.desc(), card.c.id.asc()) + + 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: @@ -634,13 +710,46 @@ class Mai2ItemData(BaseData): if result: return result.fetchall() - async def get_presents_by_version_user(self, ver: int = None, user_id: int = None) -> Optional[List[Row]]: - result = await self.execute(present.select( - and_( - or_(present.c.user == user_id, present.c.user == None), - or_(present.c.version == ver, present.c.version == None) + async def get_presents_by_version_user( + self, + version: Optional[int] = None, + user_id: Optional[int] = None, + exclude_owned: bool = False, + exclude_not_in_present_period: bool = False, + limit: Optional[int] = None, + offset: Optional[int] = None, + ) -> Optional[List[Row]]: + sql = select(present) + condition = ( + ((present.c.user == user_id) | present.c.user.is_(None)) + & ((present.c.version == version) | present.c.version.is_(None)) + ) + + # Do an anti-join with the mai2_item_item table to exclude any + # items the users have already owned. + if exclude_owned: + sql = sql.join( + item, + (present.c.itemKind == item.c.itemKind) + & (present.c.itemId == item.c.itemId) ) - )) + condition &= (item.c.itemKind.is_(None) & item.c.itemId.is_(None)) + + if exclude_not_in_present_period: + condition &= (present.c.startDate.is_(None) | (present.c.startDate <= func.now())) + condition &= (present.c.endDate.is_(None) | (present.c.endDate >= func.now())) + + sql = sql.where(condition) + + if limit is not None or offset is not None: + sql = sql.order_by(present.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: return result.fetchall() diff --git a/titles/mai2/schema/profile.py b/titles/mai2/schema/profile.py index ede0adf..ebf04a4 100644 --- a/titles/mai2/schema/profile.py +++ b/titles/mai2/schema/profile.py @@ -1,15 +1,26 @@ -from core.data.schema import BaseData, metadata -from titles.mai2.const import Mai2Constants +from datetime import datetime +from typing import Dict, List, Optional from uuid import uuid4 -from typing import Optional, Dict, List -from sqlalchemy import Table, Column, UniqueConstraint, PrimaryKeyConstraint, and_ -from sqlalchemy.types import Integer, String, TIMESTAMP, Boolean, JSON, BigInteger, SmallInteger, VARCHAR, INTEGER +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.engine import Row -from sqlalchemy.dialects.mysql import insert -from datetime import datetime +from sqlalchemy.types import ( + INTEGER, + JSON, + TIMESTAMP, + VARCHAR, + BigInteger, + Boolean, + Integer, + SmallInteger, + String, +) + +from core.data.schema import BaseData, metadata +from titles.mai2.const import Mai2Constants detail = Table( "mai2_profile_detail", @@ -495,7 +506,7 @@ consec_logins = Table( mysql_charset="utf8mb4", ) -rival = Table( +rival: Table = Table( "mai2_user_rival", metadata, Column("id", Integer, primary_key=True, nullable=False), @@ -908,8 +919,23 @@ class Mai2ProfileData(BaseData): if result: return result.fetchall() - async def get_rivals_game(self, user_id: int) -> Optional[List[Row]]: - result = await self.execute(rival.select(and_(rival.c.user == user_id, rival.c.show == True)).limit(3)) + async def get_rivals_game( + self, + user_id: int, + limit: Optional[int] = None, + offset: Optional[int] = None, + ) -> Optional[List[Row]]: + sql = select(rival).where((rival.c.user == user_id) & rival.c.show.is_(True)) + + if limit is not None or offset is not None: + sql = sql.order_by(rival.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: return result.fetchall() diff --git a/titles/mai2/schema/score.py b/titles/mai2/schema/score.py index e376216..cbe7448 100644 --- a/titles/mai2/schema/score.py +++ b/titles/mai2/schema/score.py @@ -1,15 +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, BigInteger + +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.engine import Row -from sqlalchemy.dialects.mysql import insert +from sqlalchemy.types import JSON, BigInteger, Boolean, Integer, String from core.data.schema import BaseData, metadata -from core.data import cached -best_score = Table( +best_score: Table = Table( "mai2_score_best", metadata, Column("id", Integer, primary_key=True, nullable=False), @@ -272,7 +272,7 @@ playlog_old = Table( mysql_charset="utf8mb4", ) -best_score_old = Table( +best_score_old: Table = Table( "maimai_score_best", metadata, Column("id", Integer, primary_key=True, nullable=False), @@ -313,22 +313,55 @@ class Mai2ScoreData(BaseData): return None return result.lastrowid - @cached(2) - async def get_best_scores(self, user_id: int, song_id: int = None, is_dx: bool = True) -> Optional[List[Row]]: + async def get_best_scores( + self, + user_id: int, + song_id: Optional[int] = None, + is_dx: bool = True, + limit: Optional[int] = None, + offset: Optional[int] = None, + levels: Optional[list[int]] = None, + ) -> Optional[List[Row]]: if is_dx: - sql = best_score.select( - and_( - best_score.c.user == user_id, - (best_score.c.musicId == song_id) if song_id is not None else True, - ) - ).order_by(best_score.c.musicId).order_by(best_score.c.level) + table = best_score else: - sql = best_score_old.select( - and_( - best_score_old.c.user == user_id, - (best_score_old.c.musicId == song_id) if song_id is not None else True, - ) - ).order_by(best_score.c.musicId).order_by(best_score.c.level) + table = best_score_old + + cond = table.c.user == user_id + + if song_id is not None: + cond &= table.c.musicId == song_id + + if levels is not None: + cond &= table.c.level.in_(levels) + + if limit is None and offset is None: + sql = ( + select(table) + .where(cond) + .order_by(table.c.musicId, table.c.level) + ) + else: + subq = ( + select(table.c.musicId) + .distinct() + .where(cond) + .order_by(table.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(table) + .join(subq, table.c.musicId == subq.c.musicId) + .where(cond) + .order_by(table.c.musicId, table.c.level) + ) result = await self.execute(sql) if result is None: diff --git a/titles/ongeki/base.py b/titles/ongeki/base.py index e454dc3..1bebb4d 100644 --- a/titles/ongeki/base.py +++ b/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 diff --git a/titles/ongeki/bright.py b/titles/ongeki/bright.py index 690a118..5c95af3 100644 --- a/titles/ongeki/bright.py +++ b/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: diff --git a/titles/ongeki/schema/item.py b/titles/ongeki/schema/item.py index ca2de1f..274e16d 100644 --- a/titles/ongeki/schema/item.py +++ b/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: diff --git a/titles/ongeki/schema/score.py b/titles/ongeki/schema/score.py index 6867133..178cf29 100644 --- a/titles/ongeki/schema/score.py +++ b/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: