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)