zhenxun_bot/zhenxun/models/group_member_info.py

from typing import ClassVar

from tortoise import fields

from zhenxun.configs.config import BotConfig
from zhenxun.services.db_context import Model


class GroupInfoUser(Model):
    id = fields.IntField(pk=True, generated=True, auto_increment=True)
    """自增id"""
    user_id = fields.CharField(255)
    """用户id"""
    user_name = fields.CharField(255, default="")
    """用户昵称"""
    group_id = fields.CharField(255)
    """群聊id"""
    user_join_time = fields.DatetimeField(null=True)
    """用户入群时间"""
    uid = fields.BigIntField(null=True)
    """用户uid"""
    platform = fields.CharField(255, null=True, description="平台")
    """平台"""

    class Meta:  # pyright: ignore [reportIncompatibleVariableOverride]
        table = "group_info_users"
        table_description = "群员信息数据表"
        unique_together = ("user_id", "group_id")
        indexes: ClassVar = [("group_id",), ("user_id",)]

    @classmethod
    async def get_all_uid(cls, group_id: str) -> set[str]:
        """获取该群所有用户id

        参数:
            group_id: 群号
        """
        return set(
            await cls.filter(group_id=group_id).values_list("user_id", flat=True)
        )  # type: ignore

    @classmethod
    async def get_user_all_group(cls, user_id: str) -> list[str]:
        """获取该用户所在的所有群聊

        参数:
            user_id: 用户id
        """
        return list(
            await cls.filter(user_id=user_id).values_list("group_id", flat=True)
        )  # type: ignore

    @classmethod
    async def _run_script(cls):
        db_type = (BotConfig.get_sql_type() or "").lower()
        scripts = ["ALTER TABLE group_info_users DROP COLUMN nickname;"]

        if "postgres" in db_type:
            scripts.extend(
                [
                    (
                        "ALTER TABLE group_info_users ADD COLUMN IF NOT EXISTS "
                        "platform character varying(255);"
                    ),
                    (
                        "ALTER TABLE group_info_users ALTER COLUMN user_id "
                        "TYPE character varying(255) USING user_id::character varying;"
                    ),
                    (
                        "ALTER TABLE group_info_users ALTER COLUMN group_id "
                        "TYPE character varying(255) USING group_id::character varying;"
                    ),
                    (
                        "ALTER TABLE group_info_users ALTER COLUMN platform "
                        "TYPE character varying(255) USING platform::character varying;"
                    ),
                ]
            )
        elif "mysql" in db_type:
            scripts.extend(
                [
                    (
                        "ALTER TABLE group_info_users ADD COLUMN "
                        "platform VARCHAR(255) NULL;"
                    ),
                    (
                        "ALTER TABLE group_info_users MODIFY COLUMN "
                        "user_id VARCHAR(255) NOT NULL;"
                    ),
                    (
                        "ALTER TABLE group_info_users MODIFY COLUMN "
                        "group_id VARCHAR(255) NOT NULL;"
                    ),
                    (
                        "ALTER TABLE group_info_users MODIFY COLUMN "
                        "platform VARCHAR(255) NULL;"
                    ),
                ]
            )
        elif "sqlite" in db_type:
            scripts.append(
                "ALTER TABLE group_info_users ADD COLUMN platform VARCHAR(255);"
            )

        return scripts