Commit 531784d
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