chuanqi/tools/append_server_user_store.py

"""Append ServerUserStore class to cq2_sdk store.py at the end."""
import io

target = r'D:\传奇\cq2_sdk\src\cq2_sdk\store.py'

addition = '''


# ----- ServerUserStore -----


class ServerUserStore:
    """CRUD for the ``server_users`` binding table.

    Each (server_id, user_id) pair is unique. ``is_main=1`` flags the
    user that the daily 00:00 zone-sync job will use to call
    996box's ``getTodayOpen`` for that server. Application layer
    enforces "at most one main per server" (MySQL 5.7 lacks partial
    unique indexes).
    """

    def __init__(self, conn: pymysql.connections.Connection) -> None:
        self.conn = conn

    def bind(self, server_id: int, user_id: int, server_name: str,
             username: str, is_main: int = 0) -> int:
        """Bind a user to a server. Idempotent. Returns the row id."""
        # 1. 若设主, 先把同 server 其它主清掉
        if is_main:
            with self.conn.cursor() as cur:
                cur.execute(
                    "UPDATE server_users SET is_main=0 "
                    "WHERE server_id=%s AND user_id<>%s AND is_main=1",
                    (server_id, user_id),
                )
        # 2. upsert
        sql = """
            INSERT INTO server_users (server_id, user_id, server_name, username, is_main)
            VALUES (%s, %s, %s, %s, %s)
            ON DUPLICATE KEY UPDATE
                server_name = VALUES(server_name),
                username    = VALUES(username),
                is_main     = VALUES(is_main)
        """
        with self.conn.cursor() as cur:
            cur.execute(sql, (server_id, user_id, server_name, username, is_main))
            cur.execute(
                "SELECT id FROM server_users WHERE server_id=%s AND user_id=%s",
                (server_id, user_id),
            )
            row = cur.fetchone()
            return row["id"] if row else 0

    def get_main(self, server_id: int) -> ServerUser | None:
        with self.conn.cursor() as cur:
            cur.execute(
                "SELECT * FROM server_users WHERE server_id=%s AND is_main=1 LIMIT 1",
                (server_id,),
            )
            row = cur.fetchone()
            return ServerUser.from_row(row) if row else None

    def list_servers_with_main(self) -> list[ServerUser]:
        with self.conn.cursor() as cur:
            cur.execute(
                "SELECT * FROM server_users WHERE is_main=1 ORDER BY server_id"
            )
            return [ServerUser.from_row(r) for r in cur.fetchall()]
'''

with open(target, 'a', encoding='utf-8') as f:
    f.write(addition)
print('appended ServerUserStore to', target)