diff --git a/.gitignore b/.gitignore index efd68a0a..924d8165 100644 --- a/.gitignore +++ b/.gitignore @@ -12,5 +12,6 @@ __pycache__/ db_data/ /config/ .ruff_cache/ +/app/karma_bot.session-journal .venv/ -venv/ \ No newline at end of file +venv/ diff --git a/app/config/main.py b/app/config/main.py index b16b0494..f1715dcf 100644 --- a/app/config/main.py +++ b/app/config/main.py @@ -5,13 +5,14 @@ import yaml from dotenv import load_dotenv -from app.models.config import Config, TgClientConfig +from app.models.config import Config from .db import load_db_config from .karmic_restriction import load_karmic_restriction_config from .log import load_log_config from .logging_config import logging_setup from .storage import load_storage +from .tg_client import load_tg_client_config from .webhook import load_webhook_config @@ -40,7 +41,9 @@ def load_config(config_dir: Path = None) -> Config: superusers=frozenset(config_file_data["superusers"]), log=log_config, dump_chat_id=config_file_data["dump_chat_id"], - tg_client=TgClientConfig(bot_token=_bot_token), + tg_client=load_tg_client_config( + config_file_data["tg_client_config"] | {"bot_token": _bot_token} + ), storage=load_storage(config_file_data["storage"]), report_karma_award=config_file_data.get("report_karma_award", 0), report_award_cleanup_delay=config_file_data.get( diff --git a/app/config/tg_client.py b/app/config/tg_client.py new file mode 100644 index 00000000..8fe2a3f9 --- /dev/null +++ b/app/config/tg_client.py @@ -0,0 +1,8 @@ +from app.models.config import TgClientConfig + + +def load_tg_client_config(config: dict) -> TgClientConfig: + return TgClientConfig( + bot_token=config["bot_token"], + request_interval=config["request_interval"], + ) diff --git a/app/handlers/base.py b/app/handlers/base.py index 6d89699a..3438d3c3 100644 --- a/app/handlers/base.py +++ b/app/handlers/base.py @@ -88,5 +88,5 @@ async def chat_migrate(message: types.Message, chat: Chat, chat_repo: ChatRepo): old_id = message.chat.id new_id = message.migrate_to_chat_id chat.chat_id = new_id - await chat_repo.update(chat) + await chat_repo.save(chat) logger.info(f"Migrate chat from {old_id} to {new_id}") diff --git a/app/handlers/change_karma.py b/app/handlers/change_karma.py index 885f2de2..e7031fe8 100644 --- a/app/handlers/change_karma.py +++ b/app/handlers/change_karma.py @@ -6,6 +6,7 @@ from app.filters import HasTargetFilter, KarmaFilter from app.infrastructure.database.models import Chat, User +from app.infrastructure.database.repo.user import UserRepo from app.models.config import Config from app.services.adaptive_trottle import AdaptiveThrottle from app.services.change_karma import cancel_karma_change, change_karma @@ -60,6 +61,7 @@ async def karma_change( target: User, config: Config, bot: Bot, + user_repo: UserRepo, ): try: result_change_karma = await change_karma( @@ -69,6 +71,7 @@ async def karma_change( how_change=karma["karma_change"], comment=karma["comment"], bot=bot, + user_repo=user_repo, ) except SubZeroKarma: return await message.reply("У Вас слишком мало кармы для этого") @@ -122,7 +125,10 @@ async def karma_change( @router.callback_query(kb.KarmaCancelCb.filter()) async def cancel_karma( - callback_query: types.CallbackQuery, callback_data: kb.KarmaCancelCb, bot: Bot + callback_query: types.CallbackQuery, + callback_data: kb.KarmaCancelCb, + bot: Bot, + user_repo: UserRepo, ): if callback_data.user_id != callback_query.from_user.id: return await callback_query.answer("Эта кнопка не для Вас", cache_time=3600) @@ -133,7 +139,7 @@ async def cancel_karma( else callback_data.moderator_event_id ) await cancel_karma_change( - callback_data.karma_event_id, rollback_karma, moderator_event_id, bot + callback_data.karma_event_id, rollback_karma, moderator_event_id, bot, user_repo ) await callback_query.answer("Вы отменили изменение кармы", show_alert=True) await callback_query.message.delete() diff --git a/app/handlers/karma.py b/app/handlers/karma.py index 12e37f12..afc03ba6 100644 --- a/app/handlers/karma.py +++ b/app/handlers/karma.py @@ -6,6 +6,7 @@ from app.infrastructure.database.models import Chat, User from app.infrastructure.database.repo.chat import ChatRepo +from app.infrastructure.database.repo.user import UserRepo from app.models.config import Config from app.services.karma import get_me_chat_info, get_me_info from app.services.karma import get_top as get_karma_top @@ -17,7 +18,9 @@ @router.message(Command("top", prefix="!"), F.chat.type == "private") -async def get_top_from_private(message: types.Message, user: User, chat_repo: ChatRepo): +async def get_top_from_private( + message: types.Message, user: User, chat_repo: ChatRepo, user_repo: UserRepo +): parts = message.text.split(maxsplit=1) if len(parts) > 1: chat = await chat_repo.get_by_id(chat_id=int(parts[1])) @@ -30,17 +33,23 @@ async def get_top_from_private(message: types.Message, user: User, chat_repo: Ch logger.info( "user {user} ask top karma of chat {chat}", user=user.tg_id, chat=chat.chat_id ) - text = await get_karma_top(chat, user) + text = await get_karma_top(chat, user, chat_repo=chat_repo, user_repo=user_repo) await message.reply(text, disable_web_page_preview=True) @router.message(Command("top", prefix="!")) -async def get_top(message: types.Message, chat: Chat, user: User): +async def get_top( + message: types.Message, + chat: Chat, + user: User, + chat_repo: ChatRepo, + user_repo: UserRepo, +): logger.info( "user {user} ask top karma of chat {chat}", user=user.tg_id, chat=chat.chat_id ) - text = await get_karma_top(chat, user) + text = await get_karma_top(chat, user, chat_repo=chat_repo, user_repo=user_repo) await message.reply(text, disable_web_page_preview=True) diff --git a/app/handlers/moderator.py b/app/handlers/moderator.py index c8a963be..2912d5aa 100644 --- a/app/handlers/moderator.py +++ b/app/handlers/moderator.py @@ -15,6 +15,7 @@ from app.handlers import keyboards as kb from app.infrastructure.database.models import Chat, ChatSettings, ReportStatus, User from app.infrastructure.database.repo.report import ReportRepo +from app.infrastructure.database.repo.user import UserRepo from app.models.config import Config from app.services.moderation import ( ban_user, @@ -253,10 +254,15 @@ async def get_info_about_user_private(message: types.Message): HasTargetFilter(can_be_same=True), ) async def get_info_about_user( - message: types.Message, chat: Chat, target: User, config: Config, bot: Bot + message: types.Message, + chat: Chat, + target: User, + config: Config, + bot: Bot, + user_repo: UserRepo, ): info = await get_user_info(target, chat, config.date_format) - target_karma = await target.get_karma(chat) + target_karma = await user_repo.get_karma(target, chat) if target_karma is None: target_karma = "пока не имеет кармы" information = f"Данные на {target.mention_link} ({target_karma}):\n" + "\n".join( @@ -319,6 +325,7 @@ async def approve_report_handler( config: Config, chat_settings: ChatSettings, report_repo: ReportRepo, + user_repo: UserRepo, ): logger.info( "Moderator {moderator} approved report {report}", @@ -338,6 +345,7 @@ async def approve_report_handler( chat=chat, reward_amount=config.report_karma_award, bot=bot, + user_repo=user_repo, ) message = await bot.edit_message_text( "{reporter} получил +{reward_amount} кармы " diff --git a/app/infrastructure/database/models/chat.py b/app/infrastructure/database/models/chat.py index 82276535..36ed4b71 100644 --- a/app/infrastructure/database/models/chat.py +++ b/app/infrastructure/database/models/chat.py @@ -3,25 +3,7 @@ from aiogram.utils.text_decorations import html_decoration as hd from tortoise import fields -from tortoise.exceptions import DoesNotExist from tortoise.models import Model -from tortoise.transactions import in_transaction - -from app.models.db.db import karma_filters -from app.utils.exceptions import NotHaveNeighbours - -SQL_PREV_NEXT = """ -SELECT - prev_user_id, - next_user_id -FROM (SELECT - user_id, - LAG(user_id) OVER(ORDER BY karma) prev_user_id, - LEAD(user_id) OVER(ORDER BY karma) next_user_id - FROM user_karma - WHERE chat_id = ?) -WHERE user_id = ? -""" class ChatType(str, Enum): @@ -45,21 +27,6 @@ class Chat(Model): class Meta: table = "chats" - @classmethod - async def create_from_tg_chat(cls, chat): - chat = await cls.create( - chat_id=chat.id, type_=chat.type, title=chat.title, username=chat.username - ) - return chat - - @classmethod - async def get_or_create_from_tg_chat(cls, chat): - try: - chat = await cls.get(chat_id=chat.id) - except DoesNotExist: - chat = await cls.create_from_tg_chat(chat=chat) - return chat - @property def mention(self): return ( @@ -80,51 +47,3 @@ def __str__(self): def __repr__(self): return str(self) - - # noinspection PyUnresolvedReferences - async def get_top_karma_list(self, limit: int = 15): - await self.fetch_related("user_karma") - users_karmas = ( - await self.user_karma.order_by(*karma_filters) - .limit(limit) - .prefetch_related("user") - .all() - ) - rez = [] - for user_karma in users_karmas: - user = user_karma.user - karma = user_karma.karma_round - rez.append((user, karma)) - - return rez - - # noinspection PyUnresolvedReferences - async def get_neighbours( - self, user - ) -> tuple["UserKarma", "UserKarma", "UserKarma"]: # noqa: F821 - prev_id, next_id = await get_neighbours_id(self.chat_id, user.id) - uk = ( - await self.user_karma.filter(user_id__in=(prev_id, next_id)) - .prefetch_related("user") - .order_by(*karma_filters) - .all() - ) - - user_uk = ( - await self.user_karma.filter(user=user).prefetch_related("user").first() - ) - return uk[0], user_uk, uk[1] - - -async def get_neighbours_id(chat_id, user_id) -> typing.Tuple[int, int]: - async with in_transaction() as conn: - neighbours = await conn.execute_query(SQL_PREV_NEXT, (chat_id, user_id)) - try: - rez = dict(neighbours[1][0]) - except IndexError: - raise NotHaveNeighbours - try: - rez = int(rez["prev_user_id"]), int(rez["next_user_id"]) - except TypeError: - raise NotHaveNeighbours - return rez diff --git a/app/infrastructure/database/models/user.py b/app/infrastructure/database/models/user.py index 568270bd..4e3664af 100644 --- a/app/infrastructure/database/models/user.py +++ b/app/infrastructure/database/models/user.py @@ -1,17 +1,9 @@ -from datetime import datetime from typing import TYPE_CHECKING -from aiogram import types from aiogram.utils.text_decorations import html_decoration as hd from tortoise import fields -from tortoise.exceptions import DoesNotExist from tortoise.models import Model -from app.infrastructure.database.models.chat import Chat -from app.models import dto -from app.models.common import TypeRestriction -from app.utils.exceptions import UserWithoutUserIdError - if TYPE_CHECKING: from app.infrastructure.database.models.moderator_actions import ModeratorEvent from app.infrastructure.database.models.user_karma import UserKarma @@ -30,62 +22,6 @@ class User(Model): class Meta: table = "users" - @classmethod - async def create_from_tg_user(cls, user: types.User): - user = await cls.create( - tg_id=user.id, - first_name=user.first_name, - last_name=user.last_name, - username=user.username, - is_bot=user.is_bot, - ) - - return user - - async def update_user_data(self, user_tg): - # TODO изучить фреймворк лучше - уверен есть встроенная функция для - # обновления только в случае расхождений - changed = False - - if self.tg_id is None and user_tg.id is not None: - changed = True - self.tg_id = user_tg.id - - if user_tg.first_name is not None: - if self.first_name != user_tg.first_name: - changed = True - self.first_name = user_tg.first_name - - if self.last_name != user_tg.last_name: - changed = True - self.last_name = user_tg.last_name - - if self.username != user_tg.username: - changed = True - self.username = user_tg.username - if self.is_bot is None and user_tg.is_bot is not None: - changed = True - self.is_bot = user_tg.is_bot - - if changed: - await self.save() - - @classmethod - async def get_or_create_from_tg_user(cls, user_tg: types.User | dto.TargetUser): - if user_tg.id is None: - try: - return await cls.get(username__iexact=user_tg.username) - except DoesNotExist: - raise UserWithoutUserIdError(username=user_tg.username) - - try: - user = await cls.get(tg_id=user_tg.id) - except DoesNotExist: - return await cls.create_from_tg_user(user=user_tg) - else: - await user.update_user_data(user_tg) - return user - @property def mention_link(self): return hd.link(hd.quote(self.fullname), self.link) @@ -108,37 +44,6 @@ def fullname(self): return " ".join((self.first_name, self.last_name)) return self.first_name or self.username or str(self.tg_id) or str(self.id) - async def get_uk(self, chat: Chat) -> "UserKarma": - return await self.karma.filter(chat=chat).first() - - async def get_karma(self, chat: Chat): - user_karma = await self.get_uk(chat) - if user_karma: - return user_karma.karma_round - return None - - async def set_karma(self, chat: Chat, karma: int): - user_karma = await self.karma.filter(chat=chat).first() - user_karma.karma = karma - await user_karma.save() - - async def get_number_in_top_karma(self, chat: Chat) -> int: - uk = await self.get_uk(chat) - return await uk.number_in_top() - - async def has_now_ro_db(self, chat: Chat): - my_restrictions = await self.my_restriction_events.filter( - chat=chat, type_restriction=TypeRestriction.ro.name - ).all() - for my_restriction in my_restrictions: - if ( - my_restriction.timedelta_restriction - and my_restriction.date + my_restriction.timedelta_restriction - > datetime.now() - ): - return True - return False - def to_json(self): return dict( id=self.id, diff --git a/app/infrastructure/database/repo/chat.py b/app/infrastructure/database/repo/chat.py index 049db8a5..6eeb71cf 100644 --- a/app/infrastructure/database/repo/chat.py +++ b/app/infrastructure/database/repo/chat.py @@ -2,6 +2,7 @@ from tortoise import BaseDBAsyncClient from tortoise.exceptions import DoesNotExist +from tortoise.transactions import in_transaction from app.infrastructure.database.models import UserKarma from app.infrastructure.database.models.chat import Chat @@ -20,7 +21,7 @@ def __init__(self, session: BaseDBAsyncClient | None = None): async def get_by_id(self, chat_id: int) -> Chat: return await Chat.get(chat_id=chat_id, using_db=self.session) - async def update(self, chat: Chat): + async def save(self, chat: Chat): await chat.save(using_db=self.session) async def create_from_tg_chat(self, chat) -> Chat: @@ -75,20 +76,22 @@ async def get_neighbours( return uk[0], user_uk, uk[1] async def get_neighbours_id(self, chat_id, user_id) -> Neighbours: - neighbours = await self.session.execute_query( - query=""" - SELECT prev_user_id, next_user_id - FROM ( - SELECT - user_id, - LAG(user_id) OVER(ORDER BY karma) prev_user_id, - LEAD(user_id) OVER(ORDER BY karma) next_user_id - FROM user_karma - WHERE chat_id = ? + # TODO: remove this new session when self.session won't be optional + async with in_transaction() as session: + neighbours = await session.execute_query( + """ + SELECT prev_user_id, next_user_id + FROM ( + SELECT + user_id, + LAG(user_id) OVER(ORDER BY karma) prev_user_id, + LEAD(user_id) OVER(ORDER BY karma) next_user_id + FROM user_karma + WHERE chat_id = ? + ) + WHERE user_id = ?""", + (chat_id, user_id), ) - WHERE user_id = ?""", - values=[chat_id, user_id], - ) try: rez = dict(neighbours[1][0]) except IndexError: diff --git a/app/infrastructure/database/repo/user.py b/app/infrastructure/database/repo/user.py new file mode 100644 index 00000000..f1010d72 --- /dev/null +++ b/app/infrastructure/database/repo/user.py @@ -0,0 +1,86 @@ +from aiogram import types +from tortoise import BaseDBAsyncClient +from tortoise.exceptions import DoesNotExist + +from app.infrastructure.database.models import Chat, User +from app.models import dto +from app.utils.exceptions import UserWithoutUserIdError + + +class UserRepo: + def __init__(self, session: BaseDBAsyncClient | None = None): + self.session = session + + async def get_by_id(self, user_id: int) -> User: + return await User.get(id=user_id, using_db=self.session) + + async def save(self, user: User): + await user.save(using_db=self.session) + + async def create_from_tg_user(self, user: types.User) -> User: + user = await User.create( + tg_id=user.id, + first_name=user.first_name, + last_name=user.last_name, + username=user.username, + is_bot=user.is_bot, + using_db=self.session, + ) + + return user + + async def update_user_data(self, user: User, tg_user: types.User): + changed = False + + if user.tg_id is None and tg_user.id is not None: + changed = True + user.tg_id = tg_user.id + + if tg_user.first_name is not None: + if user.first_name != tg_user.first_name: + changed = True + user.first_name = tg_user.first_name + + if user.last_name != tg_user.last_name: + changed = True + user.last_name = tg_user.last_name + + if user.username != tg_user.username: + changed = True + user.username = tg_user.username + if user.is_bot is None and tg_user.is_bot is not None: + changed = True + user.is_bot = tg_user.is_bot + + if changed: + await self.save(user) + + async def get_or_create_from_tg_user( + self, tg_user: types.User | dto.TargetUser + ) -> User: + if tg_user.id is None: + try: + return await User.get(username__iexact=tg_user.username) + except DoesNotExist: + raise UserWithoutUserIdError(username=tg_user.username) + + try: + user = await User.get(tg_id=tg_user.id) + except DoesNotExist: + return await self.create_from_tg_user(user=tg_user) + else: + await self.update_user_data(user, tg_user) + + return user + + async def get_karma(self, user: User, chat: Chat) -> float | None: + user_karma = await user.karma.filter(chat=chat).using_db(self.session).first() + if user_karma: + return user_karma.karma_round + return None + + async def get_number_in_top_karma(self, user: User, chat: Chat) -> int: + user_karma = await user.karma.filter(chat=chat).using_db(self.session).first() + return await user_karma.filter( + chat_id=user_karma.chat_id, karma__gte=user_karma.karma + ).count() diff --git a/app/middlewares/__init__.py b/app/middlewares/__init__.py index 26825e13..0473ea24 100644 --- a/app/middlewares/__init__.py +++ b/app/middlewares/__init__.py @@ -5,18 +5,17 @@ from app.middlewares.db_middleware import DBMiddleware from app.middlewares.fix_target_middleware import FixTargetMiddleware from app.models.config import Config -from app.utils.lock_factory import LockFactory +from app.services.user_getter import UserGetter from app.utils.log import Logger - logger = Logger(__name__) -def setup(dispatcher: Dispatcher, lock_factory: LockFactory, config: Config): +def setup(dispatcher: Dispatcher, user_getter: UserGetter, config: Config): logger.info("Configure middlewares...") - db_middleware_ = DBMiddleware(lock_factory) + db_middleware_ = DBMiddleware() dispatcher.update.outer_middleware.register(ConfigMiddleware(config)) dispatcher.errors.outer_middleware.register(ConfigMiddleware(config)) dispatcher.message.outer_middleware.register(db_middleware_) dispatcher.callback_query.outer_middleware.register(db_middleware_) - dispatcher.message.middleware.register(FixTargetMiddleware(tg_client_config=config.tg_client)) + dispatcher.message.middleware.register(FixTargetMiddleware(user_getter)) diff --git a/app/middlewares/db_middleware.py b/app/middlewares/db_middleware.py index 03ceb785..e9364ea9 100644 --- a/app/middlewares/db_middleware.py +++ b/app/middlewares/db_middleware.py @@ -7,21 +7,16 @@ from aiogram.types import TelegramObject from tortoise import BaseDBAsyncClient -from app.infrastructure.database.models import User from app.infrastructure.database.repo.chat import ChatRepo from app.infrastructure.database.repo.report import ReportRepo +from app.infrastructure.database.repo.user import UserRepo from app.services.settings import get_chat_settings -from app.utils.lock_factory import LockFactory from app.utils.log import Logger logger = Logger(__name__) class DBMiddleware(BaseMiddleware): - def __init__(self, lock_factory: LockFactory): - super(DBMiddleware, self).__init__() - self.lock_factory = lock_factory - async def __call__( self, handler: Callable[[TelegramObject, Dict[str, Any]], Awaitable[Any]], @@ -46,19 +41,20 @@ async def setup_chat( ): try: chat_repo = ChatRepo(session) + user_repo = UserRepo(session) report_repo = ReportRepo(session) - async with self.lock_factory.get_lock(user.id): - user = await User.get_or_create_from_tg_user(user) + user = await user_repo.get_or_create_from_tg_user(user) if chat and chat.type != "private": - async with self.lock_factory.get_lock(chat.id): - chat = await chat_repo.get_or_create_from_tg_chat(chat) - data["chat_settings"] = await get_chat_settings(chat=chat) + chat = await chat_repo.get_or_create_from_tg_chat(chat) + data["chat_settings"] = await get_chat_settings(chat=chat) except Exception as e: logger.exception("troubles with db", exc_info=e) raise e + data["user"] = user data["chat"] = chat data["chat_repo"] = chat_repo + data["user_repo"] = user_repo data["report_repo"] = report_repo diff --git a/app/middlewares/fix_target_middleware.py b/app/middlewares/fix_target_middleware.py index e61c9c51..93c8febc 100644 --- a/app/middlewares/fix_target_middleware.py +++ b/app/middlewares/fix_target_middleware.py @@ -1,16 +1,19 @@ -from typing import Callable, Any, Awaitable +import logging +from typing import Any, Awaitable, Callable from aiogram import BaseMiddleware from aiogram.types import TelegramObject -from app.models.config import TgClientConfig from app.services.find_target_user import get_db_user_by_tg_user +from app.services.user_getter import UserGetter + +logger = logging.getLogger(__name__) class FixTargetMiddleware(BaseMiddleware): - def __init__(self, tg_client_config: TgClientConfig): - super(FixTargetMiddleware, self).__init__() - self.tg_client_config = tg_client_config + def __init__(self, user_getter: UserGetter): + super().__init__() + self.user_getter = user_getter async def __call__( self, @@ -19,6 +22,10 @@ async def __call__( data: dict[str, Any], ) -> Any: if target := data.get("target", None): - target = await get_db_user_by_tg_user(target, self.tg_client_config) - data['target'] = target + logger.debug("Starting target lookup either in db or by pyrogram") + target = await get_db_user_by_tg_user( + target, self.user_getter, data["user_repo"] + ) + data["target"] = target + logger.debug("Target resolved") return await handler(event, data) diff --git a/app/models/config/tg_client.py b/app/models/config/tg_client.py index 5606b8a0..0be35246 100644 --- a/app/models/config/tg_client.py +++ b/app/models/config/tg_client.py @@ -4,5 +4,6 @@ @dataclass class TgClientConfig: bot_token: str - api_hash: str = 'eb06d4abfb49dc3eeb1aeb98ae0f581e' + api_hash: str = "eb06d4abfb49dc3eeb1aeb98ae0f581e" api_id: int = 6 + request_interval: int = 100 diff --git a/app/services/change_karma.py b/app/services/change_karma.py index 57665aae..24f647f9 100644 --- a/app/services/change_karma.py +++ b/app/services/change_karma.py @@ -10,6 +10,7 @@ User, UserKarma, ) +from app.infrastructure.database.repo.user import UserRepo from app.models.common import TypeRestriction from app.services.moderation import ( auto_restrict, @@ -35,6 +36,7 @@ async def change_karma( chat: Chat, how_change: float, bot: Bot, + user_repo: UserRepo, comment: str = "", is_reward: bool = False, ) -> ResultChangeKarma: @@ -83,6 +85,7 @@ async def change_karma( bot=bot, chat=chat, target=target_user, + user_repo=user_repo, using_db=conn, ) uk.karma = config.auto_restriction.after_restriction_karma @@ -90,7 +93,7 @@ async def change_karma( was_restricted = True else: count_auto_restrict = await get_count_auto_restrict( - target_user, chat, bot=bot + target_user, chat, bot=bot, user_repo=user_repo ) moderator_event = None was_restricted = False @@ -107,7 +110,11 @@ async def change_karma( async def cancel_karma_change( - karma_event_id: int, rollback_karma: float, moderator_event_id: int, bot: Bot + karma_event_id: int, + rollback_karma: float, + moderator_event_id: int, + bot: Bot, + user_repo: UserRepo, ): async with in_transaction() as conn: karma_event = await KarmaEvent.get(id_=karma_event_id) @@ -125,7 +132,7 @@ async def cancel_karma_change( await karma_event.delete(using_db=conn) if moderator_event_id is not None: moderator_event = await ModeratorEvent.get(id_=moderator_event_id) - restricted_user = await User.get(id=user_to_id) + restricted_user = await user_repo.get_by_id(user_id=user_to_id) if moderator_event.type_restriction == TypeRestriction.karmic_ro.name: await bot.restrict_chat_member( diff --git a/app/services/find_target_user.py b/app/services/find_target_user.py index 90b79088..45675622 100644 --- a/app/services/find_target_user.py +++ b/app/services/find_target_user.py @@ -1,12 +1,11 @@ from contextlib import suppress from aiogram import types -from pyrogram.errors import UsernameNotOccupied from tortoise.exceptions import MultipleObjectsReturned from app.infrastructure.database.models import User +from app.infrastructure.database.repo.user import UserRepo from app.models import dto -from app.models.config import TgClientConfig from app.services.user_getter import UserGetter from app.utils.exceptions import UserWithoutUserIdError from app.utils.log import Logger @@ -53,8 +52,6 @@ def has_target_user( if target_user is None: return False if not can_be_bot and target_user.is_bot: - # and not is_bot_username(target_user.username) - # don't check is_bot_username because user can have username like @user_bot return False if not can_be_same and is_one_user(author_user, target_user): return False @@ -115,38 +112,24 @@ def get_id_user(message: types.Message) -> dto.TargetUser | None: return None -def is_bot_username(username: str) -> bool: - """ - this function deprecated. user can use username like @alice_bot and it don't say that it is bot - """ - return username is not None and username[-3:] == "bot" - - async def get_db_user_by_tg_user( - target: dto.TargetUser, tg_client_config: TgClientConfig + target: dto.TargetUser, user_getter: UserGetter, user_repo: UserRepo ) -> User: - exception: Exception try: - target_user = await User.get_or_create_from_tg_user(target) - except MultipleObjectsReturned as e: - logger.warning( - "Strange, multiple username? check id={id}, username={username}", + return await user_repo.get_or_create_from_tg_user(target) + except MultipleObjectsReturned: + logger.error( + "Found multiple users with the same username: id={id}, username={username}", id=target.id, username=target.username, ) - exception = e - # In target can be user with only username - except UserWithoutUserIdError as e: - exception = e - else: - return target_user + except UserWithoutUserIdError: + logger.debug( + "User with username={username} not found in database", + username=target.username, + ) - try: - async with UserGetter(tg_client_config) as user_getter: - tg_user = await user_getter.get_user_by_username(target.username) + tg_user = await user_getter.get_user_by_username(target.username) + target_user = await user_repo.get_or_create_from_tg_user(tg_user) - target_user = await User.get_or_create_from_tg_user(tg_user) - # That username can be not valid - except (UsernameNotOccupied, IndexError): - raise exception return target_user diff --git a/app/services/karma.py b/app/services/karma.py index 3589175f..19066a6e 100644 --- a/app/services/karma.py +++ b/app/services/karma.py @@ -3,22 +3,26 @@ from aiogram.utils.markdown import hbold from app.infrastructure.database.models import Chat, User, UserKarma +from app.infrastructure.database.repo.chat import ChatRepo +from app.infrastructure.database.repo.user import UserRepo from app.utils.exceptions import NotHaveNeighbours -async def get_top(chat: Chat, user: User, limit: int = 15): - top_karmas = await chat.get_top_karma_list(limit) +async def get_top( + chat: Chat, user: User, user_repo: UserRepo, chat_repo: ChatRepo, limit: int = 15 +): + top_karmas = await chat_repo.get_top_karma_list(chat, limit) text_list = format_output( [(i, user, karma) for i, (user, karma) in enumerate(top_karmas, 1)] ) text = add_caption(text_list) try: - prev_uk, user_uk, next_uk = await chat.get_neighbours(user) + prev_uk, user_uk, next_uk = await chat_repo.get_neighbours(user, chat) except NotHaveNeighbours: return text user_ids = get_top_ids(top_karmas) - number_user_in_top = await user.get_number_in_top_karma(chat) + number_user_in_top = await user_repo.get_number_in_top_karma(user, chat) neighbours_karmas = [] if prev_uk.user.id not in user_ids: text = add_separator(text) diff --git a/app/services/moderation.py b/app/services/moderation.py index 85096ddd..1db1f77f 100644 --- a/app/services/moderation.py +++ b/app/services/moderation.py @@ -10,6 +10,7 @@ from app.config import moderation from app.config.main import load_config from app.infrastructure.database.models import Chat, ModeratorEvent, User +from app.infrastructure.database.repo.user import UserRepo from app.models.common import TypeRestriction from app.utils.exceptions import CantRestrict from app.utils.log import Logger @@ -160,14 +161,20 @@ async def user_has_now_ro(user: User, chat: Chat, bot: Bot) -> bool: async def auto_restrict( - target: User, chat: Chat, bot: Bot, using_db: TransactionContext | None = None + target: User, + chat: Chat, + bot: Bot, + user_repo: UserRepo, + using_db: TransactionContext | None = None, ) -> tuple[int, ModeratorEvent]: """ return count auto restrict """ - bot_user = await User.get_or_create_from_tg_user(await bot.me()) + bot_user = await user_repo.get_or_create_from_tg_user(await bot.me()) - count_auto_restrict = await get_count_auto_restrict(target, chat, bot_user=bot_user) + count_auto_restrict = await get_count_auto_restrict( + target, chat, user_repo=user_repo, bot_user=bot_user + ) logger.info( "auto restrict user {user} in chat {chat} for to negative karma. " "previous restrict count: {count}", @@ -196,6 +203,7 @@ async def auto_restrict( async def get_count_auto_restrict( target: User, chat: Chat, + user_repo: UserRepo, bot_user: User | None = None, bot: Bot | None = None, ) -> int: @@ -203,7 +211,7 @@ async def get_count_auto_restrict( bot is not None or bot_user is not None ), "One of bot and bot_user must be not None" if bot_user is None: - bot_user = await User.get_or_create_from_tg_user(await bot.me()) + bot_user = await user_repo.get_or_create_from_tg_user(await bot.me()) return await ModeratorEvent.filter( moderator=bot_user, user=target, diff --git a/app/services/report.py b/app/services/report.py index 0cf17a85..d4b0c085 100644 --- a/app/services/report.py +++ b/app/services/report.py @@ -6,6 +6,7 @@ from app.infrastructure.database.models import Chat, Report, ReportStatus, User from app.infrastructure.database.repo.report import ReportRepo +from app.infrastructure.database.repo.user import UserRepo from app.services.change_karma import change_karma from app.services.remove_message import delete_message_by_id from app.utils.types import ResultChangeKarma @@ -86,10 +87,14 @@ async def set_report_bot_reply( async def reward_reporter( - reporter_id: int, reward_amount: int, chat: Chat, bot: Bot + reporter_id: int, + reward_amount: int, + chat: Chat, + bot: Bot, + user_repo: UserRepo, ) -> ResultChangeKarma: - from_user = await User.get_or_create_from_tg_user(await bot.get_me()) - target_user = await User.get(id=reporter_id) + from_user = await user_repo.get_or_create_from_tg_user(await bot.get_me()) + target_user = await user_repo.get_by_id(user_id=reporter_id) return await change_karma( user=from_user, target_user=target_user, @@ -98,6 +103,7 @@ async def reward_reporter( bot=bot, comment="Report reward", is_reward=True, + user_repo=user_repo, ) diff --git a/app/services/restrict_call.py b/app/services/restrict_call.py index 6d37d867..e3eff041 100644 --- a/app/services/restrict_call.py +++ b/app/services/restrict_call.py @@ -1,25 +1,39 @@ # from https://gist.github.com/jorektheglitch/5ea4972d9cb87c2ec682604e53d1ff94 by @entressi import asyncio +import logging +from asyncio import Task +logger = logging.getLogger(__name__) -class RestrictCall: +class RestrictCall: def __init__(self, delay): self.delay = delay self.queue = asyncio.Queue() - self.bg_worker = asyncio.ensure_future(self.queue_processor()) + self.bg_worker: Task | None = None + + def start_worker(self): + self.bg_worker = asyncio.create_task(self.queue_processor()) + + def stop_worker(self): + self.bg_worker.cancel() async def queue_processor(self): while True: + logger.debug("Worker is waiting for events") can_go = await self.queue.get() can_go.set() + logger.debug("Worker fired event, sleeping for %s seconds", self.delay) await asyncio.sleep(self.delay) def __call__(self, func): async def wrapped(*args, **kwargs): can_go = asyncio.Event() + logger.debug("Invoking restricted function") await self.queue.put(can_go) + logger.debug("Waiting for the queue") await can_go.wait() + logger.debug("Calling restricted function") return await func(*args, **kwargs) return wrapped diff --git a/app/services/user_getter.py b/app/services/user_getter.py index 704cec0b..1c9ac384 100644 --- a/app/services/user_getter.py +++ b/app/services/user_getter.py @@ -1,18 +1,15 @@ import asyncio -import typing import pyrogram from aiogram.types import User from pyrogram import Client -from pyrogram.errors import RPCError, UsernameNotOccupied, FloodWait +from pyrogram.errors import FloodWait, UsernameNotOccupied from app.models.config import TgClientConfig from app.services.restrict_call import RestrictCall from app.utils.log import Logger - logger = Logger(__name__) -SLEEP_TIME = 100 class UserGetter: @@ -22,26 +19,17 @@ def __init__(self, client_config: TgClientConfig): bot_token=client_config.bot_token, api_id=client_config.api_id, api_hash=client_config.api_hash, - no_updates=True + no_updates=True, ) + self.restrict = RestrictCall(client_config.request_interval) + self.restrict_methods = ("get_users",) + self.patch_api_client() - async def get_user(self, username: str = None, fullname: str = None, chat_id: int = None) -> User: - async def try_by_name() -> typing.Optional[User]: - try: - return await self.get_user_by_fullname(chat_id, fullname) - except (IndexError, RPCError): - return None - - if username is not None: - try: - user_tg = await self.get_user_by_username(username) - except RPCError: - user_tg = await try_by_name() - else: - user_tg = await try_by_name() - return user_tg + def patch_api_client(self): + for method in self.restrict_methods: + patched = self.restrict(getattr(self._client_api_bot, method)) + setattr(self._client_api_bot, method, patched) - @RestrictCall(SLEEP_TIME) async def get_user_by_username(self, username: str) -> User: try: logger.info("get user of username {username}", username=username) @@ -52,28 +40,11 @@ async def get_user_by_username(self, username: str) -> User: raise except FloodWait as e: logger.error("Flood Wait {e}", e=e) - await asyncio.sleep(e.x) - raise IndexError - - return self.get_aiogram_user_by_pyrogram(user) - - @RestrictCall(SLEEP_TIME) - async def get_user_by_fullname(self, chat_id: int, fullname: str) -> User: - try: - logger.info("get user of name {name}", name=fullname) - chat_members = await self._client_api_bot.get_chat_members(chat_id=chat_id, query=fullname) - logger.info( - "found: {users}", - users=[self.get_user_dict_for_log(chat_member.user) for chat_member in chat_members] + await asyncio.sleep(e.value) + raise Exception( + "Username resolver encountered flood error. Waited for %s", e.value ) - user = chat_members[0].user - except IndexError: - logger.info("name not found {name}", name=fullname) - raise - except FloodWait as e: - logger.error("Flood Wait {e}", e=e) - await asyncio.sleep(e.x) - raise IndexError + return self.get_aiogram_user_by_pyrogram(user) @staticmethod @@ -98,10 +69,12 @@ def get_user_dict_for_log(user: pyrogram.types.User) -> dict: ) async def start(self): + self.restrict.start_worker() if not self._client_api_bot.is_connected: await self._client_api_bot.start() async def stop(self): + self.restrict.stop_worker() if self._client_api_bot.is_connected: await self._client_api_bot.stop() diff --git a/app/utils/cli.py b/app/utils/cli.py index ae941268..a733137b 100644 --- a/app/utils/cli.py +++ b/app/utils/cli.py @@ -1,18 +1,16 @@ # partially from https://github.com/aiogram/bot import argparse -import asyncio from aiogram import Bot, Dispatcher +from aiogram.fsm.storage.memory import SimpleEventIsolation import app +from app import handlers, middlewares +from app.models.config import Config from app.models.db import db -from app.utils.executor import on_startup_webhook, on_startup_notify -from app.utils.lock_factory import LockFactory +from app.services.user_getter import UserGetter +from app.utils.executor import on_startup_notify, on_startup_webhook from app.utils.log import Logger -from app import middlewares -from app import handlers -from app.models.config import Config - logger = Logger(__name__) PROGRAM_DESC = ( @@ -23,37 +21,43 @@ def create_parser(): - arg_parser = argparse.ArgumentParser(prog=app.__application_name__, description=PROGRAM_DESC, epilog=PROGRAM_EP) - arg_parser.add_argument('-p', '--polling', action='store_const', const=True, - help="Run tg bot with polling. Default use WebHook") - arg_parser.add_argument('-s', '--skip-updates', action='store_const', const=True, - help="Skip PENDING updates") + arg_parser = argparse.ArgumentParser( + prog=app.__application_name__, description=PROGRAM_DESC, epilog=PROGRAM_EP + ) + arg_parser.add_argument( + "-p", + "--polling", + action="store_const", + const=True, + help="Run tg bot with polling. Default use WebHook", + ) return arg_parser async def cli(config: Config): bot = Bot(config.bot_token, parse_mode="HTML") - dp = Dispatcher(storage=config.storage.create_storage()) + dp = Dispatcher( + storage=config.storage.create_storage(), events_isolation=SimpleEventIsolation() + ) parser = create_parser() namespace = parser.parse_args() await db.db_init(config.db) logger.debug(f"As application dir using: {config.app_dir}") - lock_factory = LockFactory() - middlewares.setup(dp, lock_factory, config) + user_getter = UserGetter(config.tg_client) + await user_getter.start() + middlewares.setup(dp, user_getter, config) logger.info("Configure handlers...") handlers.setup(dp, bot, config) await on_startup_notify(bot, config) try: - asyncio.create_task(lock_factory.check_and_clear()) if namespace.polling: logger.info("starting polling...") - await dp.start_polling(bot) + await dp.start_polling(bot, close_bot_session=True) else: logger.info("starting webhook...") await on_startup_webhook(bot, config.webhook) raise NotImplementedError("webhook are not implemented now") finally: await db.on_shutdown() - await bot.session.close() - + await user_getter.stop() diff --git a/app/utils/lock_factory.py b/app/utils/lock_factory.py deleted file mode 100644 index eafbf779..00000000 --- a/app/utils/lock_factory.py +++ /dev/null @@ -1,39 +0,0 @@ -import asyncio -import typing - -from app.utils.log import Logger - - -logger = Logger(__name__) - - -class LockFactory: - def __init__(self): - self._locks: typing.Optional[typing.Dict[typing.Any, asyncio.Lock]] = None - - def get_lock(self, id_: typing.Any): - if self._locks is None: - self._locks = {} - - # Это костыль, по хорошему эта таска должна создаваться в ините, - # для этого милваря должна создаваться после бота и диспетчера - - return self._locks.setdefault(id_, asyncio.Lock()) - - def clear(self): - if self._locks is None: - return - self._locks.clear() - - def clear_free(self): - if self._locks is None: - return - to_remove = [key for key, lock in self._locks.items() if not lock.locked()] - for key in to_remove: - del self._locks[key] - logger.debug("remove lock for {key}", key=key) - - async def check_and_clear(self, cool_down: int = 1800): - while True: - await asyncio.sleep(cool_down) - self.clear_free() diff --git a/config_dist/bot-config.yaml b/config_dist/bot-config.yaml index 05878222..f6758924 100644 --- a/config_dist/bot-config.yaml +++ b/config_dist/bot-config.yaml @@ -12,3 +12,6 @@ storage: dump_chat_id: -123456 log_chat_id: -123456 report_karma_award: 0 + +tg_client_config: + request_interval: 100 diff --git a/docs/deploy_manual.md b/docs/deploy_manual.md index a28a8dcf..628b599f 100644 --- a/docs/deploy_manual.md +++ b/docs/deploy_manual.md @@ -29,9 +29,13 @@ * -s - skip updates that accumulated on tg servers * To create tables in database run script:\ -```PYTHONPATH=. python migrations/01_initialize.py``` +```shell +PYTHONPATH=. python migrations/01_initialize.py +``` * for run concrete sql migration:\ -```PYTHONPATH=. python migrations/migrate.py 05_add_report_table.sql``` +```shell +PYTHONPATH=. python migrations/migrate.py 05_add_report_table.sql +``` # Karmabot deploy manual with a docker: