Refactor scores, stats and channels repos to build queries with sqlalchemy core 1.4 (osuAkatsuki#630)

* bugfixes

Co-authored-git stasby: James Wilson <[email protected]>

* begin to refactor

Co-authored-by: James Wilson <[email protected]>

* channels, stats

Co-authored-by: James Wilson <[email protected]>

* resolve dependencies + add sqlalchemy type stubs

* sqlalchemy-stubs

* le bugfix de la typing

* _POLEASE

* type fixes

* type fixes

* remove non working test makefile stuff

* fixes

* re-add pymypysql/aiomysql

---------

Co-authored-by: James Wilson <[email protected]>

Author: cmyui

Makefile

@@ -22,12 +22,6 @@ test:
 	docker-compose -f docker-compose.test.yml up -d bancho-test mysql-test redis-test
 	docker-compose -f docker-compose.test.yml exec -T bancho-test /srv/root/scripts/run-tests.sh
 
-test-local:
-	poetry run pytest -vv tests/
-
-test-dbg:
-	poetry run pytest -vv --pdb -s tests/
-
 lint:
 	poetry run pre-commit run --all-files
 

app/api/domains/osu.py

@@ -928,7 +928,7 @@ async def osuSubmitModularSelector(
             # update global & country ranking
             stats.rank = await score.player.update_rank(score.mode)
 
-    await stats_repo.update(
+    await stats_repo.partial_update(
         score.player.id,
         score.mode.value,
         plays=stats_updates.get("plays", UNSET),
@@ -1388,7 +1388,11 @@ async def getScores(
         return Response("\n".join(response_lines).encode())
 
     if personal_best_score_row is not None:
-        user_clan = await clans_repo.fetch_one(id=player.clan_id)
+        user_clan = (
+            await clans_repo.fetch_one(id=player.clan_id)
+            if player.clan_id is not None
+            else None
+        )
         display_name = (
             f"[{user_clan['tag']}] {player.name}"
             if user_clan is not None

app/commands.py

@@ -867,7 +867,11 @@ async def user(ctx: Context) -> str | None:
         else "False"
     )
 
-    user_clan = await clans_repo.fetch_one(id=player.clan_id)
+    user_clan = (
+        await clans_repo.fetch_one(id=player.clan_id)
+        if player.clan_id is not None
+        else None
+    )
     display_name = (
         f"[{user_clan['tag']}] {player.name}" if user_clan is not None else player.name
     )

app/repositories/__init__.py

@@ -0,0 +1,23 @@
+from __future__ import annotations
+
+from sqlalchemy.dialects.mysql.mysqldb import MySQLDialect_mysqldb
+from sqlalchemy.orm import DeclarativeMeta
+from sqlalchemy.orm import registry
+
+mapper_registry = registry()
+
+
+class Base(metaclass=DeclarativeMeta):
+    __abstract__ = True
+
+    registry = mapper_registry
+    metadata = mapper_registry.metadata
+
+    __init__ = mapper_registry.constructor
+
+
+class MySQLDialect(MySQLDialect_mysqldb):
+    default_paramstyle = "named"
+
+
+DIALECT = MySQLDialect()

app/repositories/achievements.py

@@ -1,6 +1,5 @@
 from __future__ import annotations
 
-import textwrap
 from collections.abc import Callable
 from typing import TYPE_CHECKING
 from typing import Any
@@ -10,24 +9,45 @@
 import app.state.services
 from app._typing import UNSET
 from app._typing import _UnsetSentinel
+from app.repositories import DIALECT
+from app.repositories import Base
 
 if TYPE_CHECKING:
     from app.objects.score import Score
 
-# +-------+--------------+------+-----+---------+----------------+
-# | Field | Type         | Null | Key | Default | Extra          |
-# +-------+--------------+------+-----+---------+----------------+
-# | id    | int          | NO   | PRI | NULL    | auto_increment |
-# | file  | varchar(128) | NO   | UNI | NULL    |                |
-# | name  | varchar(128) | NO   | UNI | NULL    |                |
-# | desc  | varchar(256) | NO   | UNI | NULL    |                |
-# | cond  | varchar(64)  | NO   |     | NULL    |                |
-# +-------+--------------+------+-----+---------+----------------+
-
-READ_PARAMS = textwrap.dedent(
-    """\
-        id, file, name, `desc`, cond
-    """,
+from sqlalchemy import Column
+from sqlalchemy import Index
+from sqlalchemy import Integer
+from sqlalchemy import String
+from sqlalchemy import delete
+from sqlalchemy import func
+from sqlalchemy import insert
+from sqlalchemy import select
+from sqlalchemy import update
+
+
+class AchievementsTable(Base):
+    __tablename__ = "achievements"
+
+    id = Column("id", Integer, primary_key=True)
+    file = Column("file", String(128), nullable=False)
+    name = Column("name", String(128, collation="utf8"), nullable=False)
+    desc = Column("desc", String(256, collation="utf8"), nullable=False)
+    cond = Column("cond", String(64), nullable=False)
+
+    __table_args__ = (
+        Index("achievements_desc_uindex", desc, unique=True),
+        Index("achievements_file_uindex", file, unique=True),
+        Index("achievements_name_uindex", name, unique=True),
+    )
+
+
+READ_PARAMS = (
+    AchievementsTable.id,
+    AchievementsTable.file,
+    AchievementsTable.name,
+    AchievementsTable.desc,
+    AchievementsTable.cond,
 )
 
 
@@ -39,41 +59,27 @@ class Achievement(TypedDict):
     cond: Callable[[Score, int], bool]
 
 
-class AchievementUpdateFields(TypedDict, total=False):
-    file: str
-    name: str
-    desc: str
-    cond: str
-
-
 async def create(
     file: str,
     name: str,
     desc: str,
     cond: str,
 ) -> Achievement:
     """Create a new achievement."""
-    query = """\
-        INSERT INTO achievements (file, name, desc, cond)
-             VALUES (:file, :name, :desc, :cond)
-    """
-    params: dict[str, Any] = {
-        "file": file,
-        "name": name,
-        "desc": desc,
-        "cond": cond,
-    }
-    rec_id = await app.state.services.database.execute(query, params)
-
-    query = f"""\
-        SELECT {READ_PARAMS}
-          FROM achievements
-         WHERE id = :id
-    """
-    params = {
-        "id": rec_id,
-    }
-    rec = await app.state.services.database.fetch_one(query, params)
+    insert_stmt = insert(AchievementsTable).values(
+        file=file,
+        name=name,
+        desc=desc,
+        cond=cond,
+    )
+    compiled = insert_stmt.compile(dialect=DIALECT)
+
+    rec_id = await app.state.services.database.execute(str(compiled), compiled.params)
+
+    select_stmt = select(READ_PARAMS).where(AchievementsTable.id == rec_id)
+    compiled = select_stmt.compile(dialect=DIALECT)
+
+    rec = await app.state.services.database.fetch_one(str(compiled), compiled.params)
     assert rec is not None
 
     achievement = dict(rec._mapping)
@@ -89,17 +95,15 @@ async def fetch_one(
     if id is None and name is None:
         raise ValueError("Must provide at least one parameter.")
 
-    query = f"""\
-        SELECT {READ_PARAMS}
-          FROM achievements
-         WHERE id = COALESCE(:id, id)
-            OR name = COALESCE(:name, name)
-    """
-    params: dict[str, Any] = {
-        "id": id,
-        "name": name,
-    }
-    rec = await app.state.services.database.fetch_one(query, params)
+    select_stmt = select(READ_PARAMS)
+
+    if id is not None:
+        select_stmt = select_stmt.where(AchievementsTable.id == id)
+    if name is not None:
+        select_stmt = select_stmt.where(AchievementsTable.name == name)
+
+    compiled = select_stmt.compile(dialect=DIALECT)
+    rec = await app.state.services.database.fetch_one(str(compiled), compiled.params)
 
     if rec is None:
         return None
@@ -111,13 +115,10 @@ async def fetch_one(
 
 async def fetch_count() -> int:
     """Fetch the number of achievements."""
-    query = """\
-        SELECT COUNT(*) AS count
-          FROM achievements
-    """
-    params: dict[str, Any] = {}
+    select_stmt = select(func.count().label("count")).select_from(AchievementsTable)
+    compiled = select_stmt.compile(dialect=DIALECT)
 
-    rec = await app.state.services.database.fetch_one(query, params)
+    rec = await app.state.services.database.fetch_one(str(compiled), compiled.params)
     assert rec is not None
     return cast(int, rec._mapping["count"])
 
@@ -127,21 +128,16 @@ async def fetch_many(
     page_size: int | None = None,
 ) -> list[Achievement]:
     """Fetch a list of achievements."""
-    query = f"""\
-        SELECT {READ_PARAMS}
-          FROM achievements
-    """
-    params: dict[str, Any] = {}
-
+    select_stmt = select(READ_PARAMS)
     if page is not None and page_size is not None:
-        query += """\
-            LIMIT :limit
-           OFFSET :offset
-        """
-        params["page_size"] = page_size
-        params["offset"] = (page - 1) * page_size
+        select_stmt = select_stmt.limit(page_size).offset((page - 1) * page_size)
 
-    records = await app.state.services.database.fetch_all(query, params)
+    compiled = select_stmt.compile(dialect=DIALECT)
+
+    records = await app.state.services.database.fetch_all(
+        str(compiled),
+        compiled.params,
+    )
 
     achievements: list[dict[str, Any]] = []
 
@@ -153,74 +149,50 @@ async def fetch_many(
     return cast(list[Achievement], achievements)
 
 
-async def update(
+async def partial_update(
     id: int,
     file: str | _UnsetSentinel = UNSET,
     name: str | _UnsetSentinel = UNSET,
     desc: str | _UnsetSentinel = UNSET,
     cond: str | _UnsetSentinel = UNSET,
 ) -> Achievement | None:
     """Update an existing achievement."""
-    update_fields: AchievementUpdateFields = {}
+    update_stmt = update(AchievementsTable).where(AchievementsTable.id == id)
     if not isinstance(file, _UnsetSentinel):
-        update_fields["file"] = file
+        update_stmt = update_stmt.values(file=file)
     if not isinstance(name, _UnsetSentinel):
-        update_fields["name"] = name
+        update_stmt = update_stmt.values(name=name)
     if not isinstance(desc, _UnsetSentinel):
-        update_fields["desc"] = desc
+        update_stmt = update_stmt.values(desc=desc)
     if not isinstance(cond, _UnsetSentinel):
-        update_fields["cond"] = cond
-
-    query = f"""\
-        UPDATE achievements
-           SET {",".join(f"{k} = COALESCE(:{k}, {k})" for k in update_fields)}
-         WHERE id = :id
-    """
-    params: dict[str, Any] = {
-        "id": id,
-    } | update_fields
-    await app.state.services.database.execute(query, params)
-
-    query = f"""\
-        SELECT {READ_PARAMS}
-          FROM achievements
-         WHERE id = :id
-    """
-    params = {
-        "id": id,
-    }
-    rec = await app.state.services.database.fetch_one(query, params)
+        update_stmt = update_stmt.values(cond=cond)
+
+    compiled = update_stmt.compile(dialect=DIALECT)
+    await app.state.services.database.execute(str(compiled), compiled.params)
+
+    select_stmt = select(READ_PARAMS).where(AchievementsTable.id == id)
+    compiled = select_stmt.compile(dialect=DIALECT)
+    rec = await app.state.services.database.fetch_one(str(compiled), compiled.params)
     assert rec is not None
 
     achievement = dict(rec._mapping)
     achievement["cond"] = eval(f'lambda score, mode_vn: {rec["cond"]}')
     return cast(Achievement, achievement)
 
 
-async def delete(
+async def delete_one(
     id: int,
 ) -> Achievement | None:
     """Delete an existing achievement."""
-    query = f"""\
-        SELECT {READ_PARAMS}
-          FROM achievements
-         WHERE id = :id
-    """
-    params: dict[str, Any] = {
-        "id": id,
-    }
-    rec = await app.state.services.database.fetch_one(query, params)
+    select_stmt = select(READ_PARAMS).where(AchievementsTable.id == id)
+    compiled = select_stmt.compile(dialect=DIALECT)
+    rec = await app.state.services.database.fetch_one(str(compiled), compiled.params)
     if rec is None:
         return None
 
-    query = """\
-        DELETE FROM achievements
-              WHERE id = :id
-    """
-    params = {
-        "id": id,
-    }
-    await app.state.services.database.execute(query, params)
+    delete_stmt = delete(AchievementsTable).where(AchievementsTable.id == id)
+    compiled = delete_stmt.compile(dialect=DIALECT)
+    await app.state.services.database.execute(str(compiled), compiled.params)
 
     achievement = dict(rec._mapping)
     achievement["cond"] = eval(f'lambda score, mode_vn: {rec["cond"]}')