diff --git a/utilities/utils.py b/utilities/utils.py index c709320..c3868b3 100644 --- a/utilities/utils.py +++ b/utilities/utils.py @@ -6,6 +6,7 @@ import nextcord import pandas as pd from sqlalchemy import distinct, select from sqlalchemy.orm import Session, aliased +from sqlmodel import col # pylint: disable=import-error from database.db import engine @@ -32,12 +33,12 @@ def autofill_categories(guild_id: int, group_id: int, category_begin: str) -> "d # search database for all category names subq = ( select(CategoryGroup) - .where(CategoryGroup.guild_id == guild_id) - .where(CategoryGroup.id == group_id) + .where(col(CategoryGroup.guild_id) == guild_id) + .where(col(CategoryGroup.id) == group_id) .subquery() ) cat_group_alias = aliased(CategoryGroup, subq) - stmt = select(Category).where(Category.group_id == cat_group_alias.id) + stmt = select(Category).where(col(Category.group_id) == cat_group_alias.id) categories_dict: dict[str, str] = {cat.name: str(cat.id) for cat in session.scalars(stmt)} # if string empty, return all category names @@ -45,7 +46,6 @@ def autofill_categories(guild_id: int, group_id: int, category_begin: str) -> "d return dict(list(categories_dict.items())[:25]) # else find all categories that start with category_begin - # TODO is it better to have sql handle the starts with or python? categories: dict[str, str] = { cat_name: cat_id for cat_name, cat_id in categories_dict.items() if category_begin.lower() in cat_name.lower() } @@ -65,7 +65,7 @@ def autofill_groups(guild_id: int, group_begin: str) -> "dict[str, str]": """ with Session(engine) as session: # Select all groups in guild. - stmt = select(CategoryGroup).where(CategoryGroup.guild_id == guild_id) + stmt = select(CategoryGroup).wherecol((CategoryGroup.guild_id) == guild_id) # dict has to have key values of same type for discord. group_dict: dict[str, str] = {group.name: str(group.id) for group in session.scalars(stmt)} @@ -102,16 +102,16 @@ def autofill_users_in_group_category(guild_id: int, group_id: int, category_id: # get all categories in the group subq = ( select(CategoryGroup) - .where(CategoryGroup.guild_id == guild_id) - .where(CategoryGroup.id == group_id) + .where(col(CategoryGroup.guild_id) == guild_id) + .where(col(CategoryGroup.id) == group_id) .subquery() ) category_group = aliased(CategoryGroup, subq) # If no category id, get all users in group. if category_id == 0: - subq = select(Category.id).where(Category.group_id == category_group.id) - stmt = select(distinct(Score.username)).where(Score.category_id.in_(subq)) + subq = select(col(Category.id)).where(Category.group_id == category_group.id) + stmt = select(distinct(col(Score.username))).where(col(Score.category_id).in_(subq)) else: # get all users in category_id @@ -148,8 +148,8 @@ def autofill_users_in_guild(guild_id: int, user_begin: str) -> "list[str]": list[str]: All users with submissions in the guild. """ with Session(engine) as session: - group_ids = select(CategoryGroup.id).where(CategoryGroup.guild_id == guild_id) - category_ids = select(Category.id).where(Category.group_id.in_(group_ids)) + group_ids = select(col(CategoryGroup.id)).where(CategoryGroup.guild_id == guild_id) + category_ids = select(Category.id).where(col(Category.group_id).in_(group_ids)) stmt = select(distinct(Score.username)).where(Score.category_id.in_(category_ids)) users_list: list[str] = list(session.scalars(stmt))