From c4d3c921bd15ed4584ceafe1e91c21ac84d1bdc3 Mon Sep 17 00:00:00 2001 From: Zakhar Izmaylov Date: Thu, 23 Jan 2025 06:34:45 +0300 Subject: [PATCH] Refactor code to utilize aiofiles for asynchronous file handling across multiple modules. Added a new function to remove blocked users from the database and updated user blocking/unblocking logic to use this new function. Improved database interaction in admin panel handlers by passing session objects instead of creating new connections. This enhances performance and maintains consistency in file operations. --- backup.py | 1 + database.py | 13 +++++++++++++ handlers/admin/admin_panel.py | 22 +++++++--------------- handlers/instructions/instructions.py | 1 + handlers/keys/keys.py | 1 + handlers/profile.py | 2 +- handlers/start.py | 1 + handlers/user.py | 7 +++---- 8 files changed, 28 insertions(+), 20 deletions(-) diff --git a/backup.py b/backup.py index 326dd183..c0a55e0a 100644 --- a/backup.py +++ b/backup.py @@ -1,6 +1,7 @@ import os import subprocess from datetime import datetime + import aiofiles from aiogram.types import BufferedInputFile diff --git a/database.py b/database.py index e226f8d5..6c3f1958 100644 --- a/database.py +++ b/database.py @@ -43,6 +43,19 @@ async def add_blocked_user(tg_id: int, conn: asyncpg.Connection): ) +async def remove_blocked_user(tg_id: int | list[int], conn: asyncpg.Connection): + """ + Удаляет пользователя или список пользователей из списка заблокированных. + + :param tg_id: ID пользователя Telegram или список ID + :param conn: Подключение к базе данных + """ + if isinstance(tg_id, list): + await conn.execute("DELETE FROM blocked_users WHERE tg_id = ANY($1)", tg_id) + else: + await conn.execute("DELETE FROM blocked_users WHERE tg_id = $1", tg_id) + + async def init_db(file_path: str = "assets/schema.sql"): with open(file_path, mode="r") as file: sql_content = file.read() diff --git a/handlers/admin/admin_panel.py b/handlers/admin/admin_panel.py index 71dba4e4..81d52d52 100644 --- a/handlers/admin/admin_panel.py +++ b/handlers/admin/admin_panel.py @@ -3,7 +3,6 @@ from datetime import datetime from io import BytesIO from typing import Any -import asyncpg from aiogram import F, Router, types from aiogram.filters import Command from aiogram.fsm.context import FSMContext @@ -13,8 +12,7 @@ from aiogram.utils.keyboard import InlineKeyboardBuilder from backup import backup_database from bot import bot -from config import DATABASE_URL -from database import delete_user_data +from database import delete_user_data, remove_blocked_user from filters.admin import IsAdminFilter from logger import logger @@ -443,10 +441,9 @@ async def handle_ban_user(callback_query: types.CallbackQuery): @router.callback_query(F.data == "export_to_csv") -async def export_banned_users_to_csv(callback_query: types.CallbackQuery): - conn = await asyncpg.connect(DATABASE_URL) +async def export_banned_users_to_csv(callback_query: types.CallbackQuery, session: Any): try: - banned_users = await conn.fetch("SELECT tg_id, blocked_at FROM blocked_users") + banned_users = await session.fetch("SELECT tg_id, blocked_at FROM blocked_users") import csv import io @@ -476,15 +473,12 @@ async def export_banned_users_to_csv(callback_query: types.CallbackQuery): text=f"Ошибка при выгрузке CSV: {e}", reply_markup=builder.as_markup(), ) - finally: - await conn.close() @router.callback_query(F.data == "delete_banned_users") -async def delete_banned_users(callback_query: types.CallbackQuery): - conn = await asyncpg.connect(DATABASE_URL) +async def delete_banned_users(callback_query: types.CallbackQuery, session: Any): try: - blocked_users = await conn.fetch("SELECT tg_id FROM blocked_users") + blocked_users = await session.fetch("SELECT tg_id FROM blocked_users") blocked_ids = [record["tg_id"] for record in blocked_users] if not blocked_ids: @@ -492,9 +486,9 @@ async def delete_banned_users(callback_query: types.CallbackQuery): return for tg_id in blocked_ids: - await delete_user_data(conn, tg_id) + await delete_user_data(session, tg_id) - await conn.execute("DELETE FROM blocked_users WHERE tg_id = ANY($1)", blocked_ids) + await remove_blocked_user(blocked_ids, session) builder = InlineKeyboardBuilder() builder.row(InlineKeyboardButton(text="⬅️ Назад", callback_data="bot_management")) @@ -509,5 +503,3 @@ async def delete_banned_users(callback_query: types.CallbackQuery): text=f"Ошибка при удалении записей: {e}", reply_markup=builder.as_markup(), ) - finally: - await conn.close() diff --git a/handlers/instructions/instructions.py b/handlers/instructions/instructions.py index 2686c24d..9663531d 100644 --- a/handlers/instructions/instructions.py +++ b/handlers/instructions/instructions.py @@ -1,5 +1,6 @@ import os from typing import Any + import aiofiles import asyncpg from aiogram import F, Router, types diff --git a/handlers/keys/keys.py b/handlers/keys/keys.py index 7b8261df..a0b4c6b5 100644 --- a/handlers/keys/keys.py +++ b/handlers/keys/keys.py @@ -3,6 +3,7 @@ import locale import os from datetime import datetime, timedelta from typing import Any + import aiofiles import asyncpg import pytz diff --git a/handlers/profile.py b/handlers/profile.py index 6e7a3126..82023a45 100644 --- a/handlers/profile.py +++ b/handlers/profile.py @@ -1,9 +1,9 @@ import os from typing import Any +import aiofiles import asyncpg from aiogram import F, Router, types -import aiofiles from aiogram.fsm.context import FSMContext from aiogram.types import BufferedInputFile, InlineKeyboardButton from aiogram.utils.keyboard import InlineKeyboardBuilder diff --git a/handlers/start.py b/handlers/start.py index 866890e4..a48f02d6 100644 --- a/handlers/start.py +++ b/handlers/start.py @@ -1,5 +1,6 @@ import os from typing import Any + import aiofiles from aiogram import F, Router from aiogram.filters import Command diff --git a/handlers/user.py b/handlers/user.py index 5fc65f9c..bc5e590d 100644 --- a/handlers/user.py +++ b/handlers/user.py @@ -4,6 +4,7 @@ from aiogram import Router from aiogram.filters.chat_member_updated import KICKED, MEMBER, ChatMemberUpdatedFilter from aiogram.types import ChatMemberUpdated +from database import add_blocked_user, remove_blocked_user from logger import logger router = Router() @@ -12,12 +13,10 @@ router = Router() @router.my_chat_member(ChatMemberUpdatedFilter(member_status_changed=KICKED)) async def user_blocked_bot(event: ChatMemberUpdated, session: Any): logger.info(f"User {event.from_user.id} blocked the bot.") - await session.execute( - "INSERT INTO blocked_users (tg_id) VALUES ($1) ON CONFLICT (tg_id) DO NOTHING", event.from_user.id - ) + await add_blocked_user(event.from_user.id, session) @router.my_chat_member(ChatMemberUpdatedFilter(member_status_changed=MEMBER)) async def user_unblocked_bot(event: ChatMemberUpdated, session: Any): logger.info(f"User {event.from_user.id} unblocked the bot.") - await session.execute("DELETE FROM blocked_users WHERE tg_id = $1", event.from_user.id) + await remove_blocked_user(event.from_user.id, session)