Commit 531784d

benny-dou <60535774+benny-dou@users.noreply.github.com>
2025-08-09 06:26:17
fix(history): fix `get_user` function to handle usernames and UIDs correctly
1 parent 47076b9
Changed files (4)
src/history/query.py
@@ -7,14 +7,23 @@ from glom import glom
 from loguru import logger
 from pyrogram.client import Client
 from pyrogram.parser.markdown import BLOCKQUOTE_EXPANDABLE_DELIM, BLOCKQUOTE_EXPANDABLE_END_DELIM
-from pyrogram.types import Message
+from pyrogram.types import Message, User
 
 from config import HISTORY, PREFIX, TZ, cache
 from database.d1 import query_d1
 from database.turso import turso_exec, turso_parse_resp
 from history.d1 import get_d1_chatinfo, save_chatinfo_to_d1
-from history.turso import get_turso_chatinfo, get_user, save_chatinfo_to_turso
-from history.utils import TURSO_KWARGS, check_save_history, filter_response, generate_query, get_chat, is_admin, list_chat_ids
+from history.turso import get_turso_chatinfo, save_chatinfo_to_turso
+from history.utils import (
+    TURSO_KWARGS,
+    check_save_history,
+    filter_response,
+    generate_query,
+    get_chat,
+    get_user_from_chat,
+    is_admin,
+    list_chat_ids,
+)
 from llm.utils import convert_html
 from messages.parser import parse_chat, parse_msg
 from messages.progress import modify_progress
@@ -217,7 +226,7 @@ async def query_history(
     if user:
         # 由于username可以修改, 我们优先使用UID进行匹配
         real_cid = cinfo["chandle"] if cinfo.get("chandle") else cinfo["cid"] if cinfo["ctype"] in ["BOT", "PRIVATE"] else f"-100{cinfo['cid']}"
-        if uid := await get_uid_by_username(client, real_cid, user):
+        if uid := await get_uid_by_username(client, real_cid, user, engine):
             sql += f" AND T.uid = {uid}"
         else:
             sql += f" AND T.user = '{user}'"
@@ -251,7 +260,7 @@ async def query_history(
     return {"texts": texts.strip(), "full_texts": full_texts.strip(), "count": count}
 
 
-async def get_uid_by_username(client: Client, chat_id: str | int, username: str) -> int:
+async def get_uid_by_username(client: Client, chat_id: str | int, username: str, engine: str = HISTORY.QUERY_ENGINE) -> int:
     """Get Telegram user id by username.
 
     Support formats of `username`:
@@ -259,6 +268,73 @@ async def get_uid_by_username(client: Client, chat_id: str | int, username: str)
     """
     if cache.get(f"get_uid_by_username-{chat_id}-{username}"):
         return cache.get(f"get_uid_by_username-{chat_id}-{username}")
-    user = await get_user(client, to_int(username), chat_id)
+    user = await get_user(client, to_int(username), chat_id, engine)
     cache.set(f"get_uid_by_username-{username}", user.id, ttl=0)
     return user.id
+
+
+async def get_user(client: Client, uid: int | str, cid: int | str = "", engine: str = HISTORY.QUERY_ENGINE) -> User:
+    try:
+        found = await client.get_users(to_int(uid))
+        if not isinstance(found, User):
+            return User(id=0)
+        # check if this user is really in this chat
+        # this step is important because:
+        # the `uid` could be a fullname like "Tom", but the handle "@Tom" is occupied by another user
+        found = await get_user_from_chat(client, found.id, cid)
+        if found.id != 0:
+            return found
+    except Exception as e:
+        logger.warning(e)
+
+    user = await get_user_from_chat(client, uid, cid)
+    if user.id == 0:  # this uid is not in this chat
+        users = await get_turso_userinfo_by_uid(uid, cid) if engine == "turso" else await get_d1_userinfo_by_uid(uid, cid)
+        for user_id, chat_id in users:  # check if this user is still in this chat
+            found = await get_user_from_chat(client, user_id, chat_id)
+            if found.id != 0:
+                return found
+    return User(id=0)
+
+
+async def get_turso_userinfo_by_uid(uid: int | str, cid: int | str = "") -> list[tuple[int, int]]:
+    """Get user info by uid from turso.
+
+    Returns:
+        [(uid, cid)]
+    """
+    uid = to_int(uid)
+    cond = f"uid = {uid}" if isinstance(uid, int) else f"handle = '{uid}' OR name = '{uid}'"
+    resp = await turso_exec([{"type": "execute", "stmt": {"sql": f"SELECT cid,uid FROM userinfo WHERE {cond};"}}], retry=2, silent=True, **TURSO_KWARGS)
+    parsed = turso_parse_resp(resp)
+    if cid:
+        parsed = [x for x in parsed if slim_cid(x["cid"]) == slim_cid(cid)]
+    res = []
+    for info in parsed:
+        cid = int(info["cid"])
+        if chat := await get_turso_chatinfo(cid):
+            real_cid = int(cid) if chat["ctype"] in ["PRIVATE", "BOT"] else int(f"-100{cid}")
+            res.append((int(info["uid"]), real_cid))
+    return res
+
+
+async def get_d1_userinfo_by_uid(uid: int | str, cid: int | str = "") -> list[tuple[int, int]]:
+    """Get user info by uid from D1.
+
+    Returns:
+        [(uid, cid)]
+    """
+    uid = to_int(uid)
+    cond = f"uid = {uid}" if isinstance(uid, int) else f"handle = '{uid}' OR name = '{uid}'"
+
+    resp = await query_d1(f"SELECT cid,uid FROM userinfo WHERE {cond};", db_name=HISTORY.D1_DATABASE, silent=True)
+    parsed = glom(resp, "result.0.results", default=[])
+    if cid:
+        parsed = [x for x in parsed if slim_cid(x["cid"]) == slim_cid(cid)]
+    res = []
+    for info in parsed:
+        cid = int(info["cid"])
+        if chat := await get_d1_chatinfo(cid):
+            real_cid = int(cid) if chat["ctype"] in ["PRIVATE", "BOT"] else int(f"-100{cid}")
+            res.append((int(info["uid"]), real_cid))
+    return res
src/history/turso.py
@@ -10,12 +10,11 @@ from zoneinfo import ZoneInfo
 from glom import Coalesce, flatten, glom
 from loguru import logger
 from pyrogram.client import Client
-from pyrogram.errors import PeerIdInvalid, UsernameNotOccupied
-from pyrogram.types import Message, User
+from pyrogram.types import Message
 
 from config import DOWNLOAD_DIR, HISTORY, TZ, cache, cutter
 from database.turso import insert_statement, turso_create_table, turso_exec, turso_parse_resp
-from history.utils import CHAT_COLUMNS, MSG_COLUMNS, MSG_INDEXES, TURSO_KWARGS, USER_COLUMNS, USER_INDEXES, check_save_history, fine_grained_check, get_chat, get_user_from_chat
+from history.utils import CHAT_COLUMNS, MSG_COLUMNS, MSG_INDEXES, TURSO_KWARGS, USER_COLUMNS, USER_INDEXES, check_save_history, fine_grained_check, get_chat
 from messages.parser import parse_chat, parse_msg
 from utils import i_am_bot, nowdt, slim_cid, to_int, true
 
@@ -384,40 +383,3 @@ async def save_userinfo_to_turso(client: Client, minfo: dict) -> dict[str, str]:
         cache.set(f"turso-user-{uid}-{cid}", records, ttl=0)
         await turso_exec([insert_statement("userinfo", records, update_on_conflict="id")], retry=2, **TURSO_KWARGS)
     return records
-
-
-async def get_user(client: Client, uid: int | str, cid: int | str = "") -> User:
-    try:
-        user = await client.get_users(to_int(uid))
-        if isinstance(user, User):
-            return user
-    except (PeerIdInvalid, UsernameNotOccupied):
-        user = await get_user_from_chat(client, uid, cid)
-        if user.id == 0:  # this uid is not in this chat
-            users = await get_userinfo_by_uid(uid, cid)
-            for user_id, chat_id in users:  # check if this user is still in this chat
-                user = await get_user_from_chat(client, user_id, chat_id)
-                if user.id != 0:
-                    return user
-    return User(id=0)
-
-
-async def get_userinfo_by_uid(uid: int | str, cid: int | str = "") -> list[tuple[int, int]]:
-    """Get user info by uid from turso.
-
-    Returns:
-        [(uid, cid)]
-    """
-    uid = to_int(uid)
-    cond = f"uid = {uid}" if isinstance(uid, int) else f"handle = '{uid}' OR name = '{uid}'"
-    resp = await turso_exec([{"type": "execute", "stmt": {"sql": f"SELECT cid,uid FROM userinfo WHERE {cond};"}}], retry=2, silent=True, **TURSO_KWARGS)
-    parsed = turso_parse_resp(resp)
-    if cid:
-        parsed = [x for x in parsed if slim_cid(x["cid"]) == slim_cid(cid)]
-    res = []
-    for info in parsed:
-        cid = int(info["cid"])
-        if chat := await get_turso_chatinfo(cid):
-            real_cid = int(cid) if chat["ctype"] in ["PRIVATE", "BOT"] else int(f"-100{cid}")
-            res.append((int(info["uid"]), real_cid))
-    return res
src/history/utils.py
@@ -149,8 +149,11 @@ def is_admin(uid: int) -> bool:
     return any(slim_cid(admin) == slim_cid(uid) for admin in strings_list(TID.HISTORY_ADMIN))
 
 
+@cache.memoize(ttl=10)
 async def get_user_from_chat(client: Client, uid: int | str, cid: int | str) -> User:
     user = User(id=0)
+    if any(char not in f"{string.ascii_letters}_{string.digits}" for char in str(uid)):
+        return user
     try:  # get chat member directly
         chat_member = await client.get_chat_member(to_int(cid), to_int(uid))
         user = chat_member.user
src/quotly/quotly.py
@@ -10,7 +10,7 @@ from pyrogram.client import Client
 from pyrogram.types import Message
 
 from config import PREFIX
-from history.turso import get_user
+from history.query import get_user
 from messages.sender import send2tg
 from messages.utils import equal_prefix, set_reaction, startswith_prefix
 from quotly.api import generate_from_api