Реализация взаимодействия с БД через Middleware в Telegram-ботах / Хабр

Вступление

Когда я начинал писать своих первых ботов с использованием базы данных, их код был очень плохим: он расходовал лишние ресурсы, а также была плохая архитектура проекта. Поэтому я хочу поделиться с вами своими знаниями, чтобы вы не наступали на те грабли, на которые наступал я. В проекте бота, который будет использован в качестве примера в данной статье, я использовал такие технологии, как aiogram, SQLAlchemy, alembic и Docker. В качестве СУБД выступает PostgreSQL. Приятного чтения!

Коротко о проекте на примере которого мы будем рассматривать работу с БД

Этот проект является Telegram-ботом для чтения книги, который рассматривается в курсе от Михаила Крыжановского на платформе Stepik. Я решил использовать его бота как основу для своего проекта, поскольку придумать что-то своё с нуля оказалось сложным, особенно когда ещё нужно было объяснить всё.

В курсе Михаил использовал обычный словарь в качестве примера базы данных для хранения информации. Однако, я решил улучшить это и использовать настоящую базу данных в виде PostgreSQL.

Ссылка на курс Михаила

Ссылка на Telegram-бота

Структура проекта

BookBot ?
├── alembic ?
│ ├── versions ?
│ │ └── 001_version.py
│ ├── env.py
│ ├── README
│ └── script.py.mako
├── app ?
│ ├── book ?
│ │ └── book.txt
│ ├── config ?
│ │ ├── init.py
│ │ └── config.py
│ ├── database ?
│ │ ├── init.py
│ │ ├── base.py
│ │ ├── models.py
│ │ └── requests.py
│ ├── filters ?
│ │ ├── init.py
│ │ └── filters.py
│ ├── handlers ?
│ │ ├── init.py
│ │ └── user_handlers.py
│ ├── keyboards ?
│ │ ├── init.py
│ │ ├── bookmarks_kb.py
│ │ ├── pagination_kb.py
│ │ └── set_menu.py
│ ├── lexicon ?
│ │ ├── init.py
│ │ └── lexicon.py
│ ├── middlewares ?
│ │ ├── init.py
│ │ └── db.py
│ ├── services ?
│ │ ├── init.py
│ │ └── file_handling.py
│ ├── init.py
│ └── main.py
├── README.md
├── alembic.ini
├── .env
├── .dockerignore
├── Dockerfile
├── docker-compose.yml
├── .env.example
└── .gitignore

Что же нового в структуре проекта в отличии от того что было в курсе Михаила

Во-первых, весь код бота был вынесен в папку "app" для большего удобства. Также в корне проекта добавилась папка "alembic" и файл "alembic.ini". Все это необходимо для работы с Alembic и выполнения миграций базы данных. Далее расположена папка "database" внутри папки "app", в которой содержатся все файлы для работы с базой данных в нашем боте. Также была добавлена папка "middlewares", в которой находится файл "db.py". По названию можно понять, что в этом файле содержится пользовательская middleware, благодаря которой мы сможем получать объект класса "Database" в наших хэндлерах и фильтрах для работы с нашей БД.

Какие преимущества мы получаем используя Middleware для работы с БД

Решение проблемы DRY (Don't Repeat Yourself): Наш код становится более эффективным и поддерживаемым благодаря тому, что мы выносим логику доступа к базе данных в пользовательскую Middleware. Вместо того чтобы повторять одну и ту же логику в каждом хэндлере.
Структурированность кода: Использование Middleware позволяет нам лучше организовать наш код. Логика доступа к базе данных вынесена в отдельный класс, что делает код более читаемым и легко понятным.

Разбор кода

Давайте начнем разбирать код Middleware, который будет передавать соединение с базой данных в обработчики и фильтры.

Пример Middleware

class DatabaseMiddleware(BaseMiddleware):
    def __init__(self, session: async_sessionmaker[AsyncSession]) -> None:
        self.session = session

    async def __call__(
        self, 
        handler: Callable[[TelegramObject, Dict[str, Any]], Awaitable[Any]],
        event: TelegramObject, 
        data: Dict[str, Any]) -> Any:

        async with self.session() as session:
            db = Database(session=session)
            data['db'] = db
            return await handler(event, data)

Рассмотрим __init__. Как мы можем заметить, при регистрации нашей Middleware передается сессия базы данных.
Далее нас ждет асинхронный магический метод __call__. Где я использую асинхронный контекстный менеджер async with self.session() as session, чтобы гарантировать правильное открытие и закрытие сессии базы данных и управление ресурсами.
Внутри контекста сессии я создаю объект db класса Database, который принимает сессию в качестве параметра. Этот объект будет использоваться для взаимодействия с базой данных.
Я добавляю объект db в словарь data по ключу 'db', чтобы передать его в другие компоненты приложения, такие как хэндлеры и фильтры. Таким образом, они смогут использовать об��ект db для работы с базой данных. После чего мы возвращаем результат работы хэндлера.

Если вы хоть раз использовали SQLAlchemy то вы знаете что нужно создать engine который обеспечивает связь между вашим приложением и БД. Так где же его лучше создать? Как по мне одним из самых хороших вариантов будет в функции main, в нашей точки в хода в программу.

Пример функции main

async def main():
    config: Config = load_config()
    engine = create_async_engine(url=config.db.url, echo=True)
    session = async_sessionmaker(engine, expire_on_commit=False)

    bot = Bot(token=config.tg_bot.token, parse_mode='HTML')
    dp = Dispatcher()
    dp.include_router(user_handlers.router)
    dp.update.middleware(DatabaseMiddleware(session=session))
    await set_main_menu(bot=bot)
    await bot.delete_webhook(drop_pending_updates=True)
    await dp.start_polling(bot)

В коде функции main ничего особо не изменилось, кроме добавления новых переменных engine и session. Мы также зарегистрировали нашу Middleware в диспетчере, передавая переменную session.

Как это всё работает

Давайте разбирать на примере хэндлера который обрабатывает команду /start

Пример хэндлера

@router.message(CommandStart())
async def process_start_command(message: Message, db: Database):
    await message.answer(LEXICON[message.text])
    await db.add_user(
        id=message.from_user.id
    )

Как вы можете заметить, теперь мы можем сразу получить объект класса Database, так как он является аргументом функции, и использовать его. Это действительно прекрасно, так как нам не нужно каждый раз писать код в хэндлере, который отвечает за установку соединения с базой данных. Мы сразу получаем готовый объект и можем работать с базой данных, используя его методы.

Бонус

Может быть, некоторые из вас слышали о aiogram_dialog от Tishka17, которая является надстройкой над aiogram и предоставляет совершенно новый способ написания Telegram-ботов с использованием диалогов. В aiogram_dialog при реализации диалога мы работаем с окнами, которые представляют собой сообщения, отправляемые пользователю. Окна создаются с использованием различных виджетов.

Пример диалога

begin = Dialog(
    Window(
        Format(text='Привет {name}!'),
        Button(
            text=Const('?Меню?'),
            id='menu',
            on_click=go_to_menu
        ),
        getter=get_user_name,
        state=BeginningSG.begin
    )
    
)

В данном примере давайте представим что нам надо брать имя пользователя из БД, а не из события в котором есть его имя.

Теперь давайте рассмотрим, как передать информацию из базы данных в наше окно при использовании aiogram_dialog. Для этого используются геттеры, в которых мы определяем данные, которые хотим передать в окно. Получение сессии с базой данных не отличается от предыдущего примера, за исключением одного момента.

Пример гетера

async def get_user_name(event_from_user: User, dialog_manager: DialogManager, **kwargs):
    db: Database = dialog_manager.middleware_data.get('db')
    user = await db.get_user_data(
        user_id=event_from_user.id
    )
    return {'name': user.name}

Как можно заметить, получение объекта класса Database теперь происходит через dialog_manager. Такой же подход мы используем, когда обрабатываем нажатие на кнопку (on_click), и нам нужно занести данные в БД.

Однако, когда пользователь начинает общение с нашим ботом, всё начинается с команды /start, и в этом хэндлере мы указываем запуск нашего диалога и добавление пользователя в БД. То есть dialog_manager присутствует и в хэндлере, который обрабатывает команду /start. Таким образом, чтобы снова получить объект класса Database, нам нужно использовать dialog_manager? Ответ: нет!

Пример хэндлера

@router.message(CommandStart())
async def start(message: Message, db: Database, dialog_manager: DialogManager):
    await db.add_user(
        user_id=message.from_user.id,
        join_date=message.date
    )
    await dialog_manager.start(state=BeginningSG.begin, mode=StartMode.RESET_STACK)

Как можно заметить получение объекта класса Database ничем не отличается из примера где я рассказывал про aiogram.

Завершение

Надеюсь это статья вам поможет начать вам лучше писать ваших ботов, и спасибо всем кто прочитал до конца.

Полезные ресурсы

Курс Михаила на Stepik по созданию телеграм ботов для начинающих

Курс Михаила на Stepik по созданию телеграм ботов для продвинутых

Книга Груши(Groosha) по созданию телеграм ботов

Документация по aiogram_dialog

Реализация взаимодействия с БД через Middleware в Telegram-ботах