SQLAlchemy и SQLite: подключение и хранение данных в Telegram-боте

В предыдущей части мы написали шахматный бот, и добавили в него интерактивную клавиатуру.

Однако, если перезапустить бота в середине игры, игра пропадет, поскольку все доски хранятся в Python-словаре во внутренней памяти программы.

Одно из возможных решений — хранить состояние игры в базе данных.

В этой статье мы добавим поддержку базы данных SQLite с использованием SQLAlchemy. Это позволит сохранять состояние игры между перезапусками.

Почему SQLAlchemy и SQLite?

SQLite:

  • Легкость установки и использования: не требует установки дополнительных программ.
  • Подходит для небольших и тестовых проектов, таких как наш Telegram-бот.

SQLAlchemy:

  • Позволяет работать с базой данных через ORM (объектно-реляционное отображение), что делает код более читаемым и поддерживаемым.
  • Поддерживает асинхронные операции, необходимые для работы с асинхронным Telegram-ботом.

Установка пакетов

Для добавления функциональности установим нужные зависимости:

pip install sqlalchemy aiosqlite

Создание дополнительных файлов

Для удобства работы создадим два новых файла:

db_models.py — содержит описание базы данных и её конфигурацию. board.py — реализует интерфейс для взаимодействия с базой данных.

Создание модели

Файл db_models.py:

from sqlalchemy import BigInteger, String
from sqlalchemy.orm import Mapped, mapped_column
from sqlalchemy.orm import DeclarativeBase
from sqlalchemy.ext.asyncio import AsyncAttrs, async_sessionmaker, create_async_engine, AsyncSession


engine = create_async_engine(url='sqlite+aiosqlite:///db.sqlite3')
async_session = async_sessionmaker(engine, class_=AsyncSession)


class Base(AsyncAttrs, DeclarativeBase):
    pass


class ChessBoard(Base):
    __tablename__ = 'boards'

    telegram_user_id: Mapped[int] = mapped_column(BigInteger, primary_key=True)
    board: Mapped[str] = mapped_column(String, nullable=False)
  • create_async_engine — создает асинхронное подключение к SQLite.
  • async_sessionmaker — предоставляет фабрику сессий для взаимодействия с базой.
  • Класс Base — базовый класс для всех моделей. Мы наследуемся от AsyncAttrs, поскольку в асинхронном боте мы хотим асинхронный sqlalchemy.
  • ChessBoard — модель для хранения состояния доски.
    • __tablename__ — определяет имя таблицы в базе данных
    • telegram_user_id — идентификатор пользователя Telegram, тип BigInteger, отображается в тип Python int.
    • board — строка FEN, описывающая текущее состояние доски.

Создание таблиц

Создадим необходимую таблицу в БД вручную. Для этого допишем в db_models.py:

if __name__ == "__main__":
    import asyncio

    async def create_tables():
        async with engine.begin() as conn:
            await conn.run_sync(Base.metadata.create_all)

    asyncio.run(create_tables())

Этот кусок кода выполнится только, если запустить db_models.py как самостоятельную программу. Запустим db_models.py один раз, чтобы создать таблицы. После запуска в папке с проектом у вас появится файл db.sqlite3. Это база данных нашего Telegram-бота.

Интерфейс доступа к SQLAlchemy

Напишем отдельный класс для доступа к базе данных и преобразования в объект chess.Board. Файл board.py:

from typing import Optional

import sqlalchemy as sql
import chess

from db_models import ChessBoard, async_session


class UserBoard:
    def __init__(self, telegram_user_id: int):
        self.telegram_user_id = telegram_user_id

    async def create(self) -> chess.Board:
        board = chess.Board()
        async with async_session() as session:
            board_obj = ChessBoard(
                telegram_user_id=self.telegram_user_id,
                board=board.fen(),
            )
            session.add(board_obj)
            await session.commit()

        return board

    async def get(self) -> Optional[chess.Board]:
        async with async_session() as session:
            board_obj = await session.scalar(
                sql.select(ChessBoard).filter_by(telegram_user_id=self.telegram_user_id)
            )
        if board_obj is None:
            return None
        return chess.Board(board_obj.board)

    async def get_or_create(self) -> chess.Board:
        async with async_session() as session:
            board_obj = await session.scalar(
                sql.select(ChessBoard).filter_by(telegram_user_id=self.telegram_user_id)
            )
            if board_obj is not None:
                return chess.Board(board_obj.board)

            board = chess.Board()
            board_obj = ChessBoard(
                telegram_user_id=self.telegram_user_id,
                board=board.fen(),
            )
            session.add(board_obj)
            await session.commit()
            return board


    async def update(self, new_board: chess.Board) -> chess.Board:
        async with async_session() as session:
            await session.execute(
                sql.update(ChessBoard).filter_by(
                    telegram_user_id=self.telegram_user_id
                ).values(
                    board=new_board.fen()
                )
            )
            await session.commit()

        return new_board

    async def delete(self) -> None:
        async with async_session() as session:
            await session.execute(
                sql.delete(ChessBoard).filter_by(
                    telegram_user_id=self.telegram_user_id
                )
            )
            await session.commit()

Так выглядит простой интерфейс create read update delete для sqlalchemy.

Осталось модифицировать файл chessbot.py:

import asyncio
import logging
import sys

from aiogram import Bot, Dispatcher, F
from aiogram.client.default import DefaultBotProperties
from aiogram.enums import ParseMode
from aiogram.filters import CommandStart, Command, StateFilter
from aiogram.types import Message, BufferedInputFile, CallbackQuery
from aiogram.fsm.context import FSMContext
from aiogram.fsm.state import StatesGroup, State

from my_credentials import API_TOKEN
from chess_utils import generate_board_image, first_tap_keyboard, second_tap_keyboard, make_move, make_bot_move
from board import UserBoard


dp = Dispatcher()


class ChessMatchStatesGroup(StatesGroup):
    has_chosen_piece = State()


@dp.message(CommandStart())
async def send_welcome(message: Message) -> None:
    await message.reply("Привет! Я ваш первый Telegram-бот.")

@dp.message(Command("newgame"))
async def start_game(message: Message) -> None:
    user_board = UserBoard(message.from_user.id)
    initial_board = await user_board.get_or_create()

    await message.answer_photo(
        photo=BufferedInputFile(generate_board_image(initial_board), "game.png"),
        caption="Игра началась! Делай ход. Нажми на клавиатуру: первым нажатием выбери фигуру, вторым - куда она пойдет",
        reply_markup=first_tap_keyboard(initial_board),
    )

@dp.callback_query(StateFilter(None), F.data.startswith("turn_"))
async def process_from_move(message: CallbackQuery, state: FSMContext) -> None:
    user_board = UserBoard(message.from_user.id)
    board = await user_board.get()
    if board is None:
        return

    square_name = message.data.removeprefix("turn_")
    await state.set_data({"square": square_name})
    await state.set_state(ChessMatchStatesGroup.has_chosen_piece)
    await message.message.edit_reply_markup(
        reply_markup=second_tap_keyboard(board, square_name),
    )


@dp.callback_query(ChessMatchStatesGroup.has_chosen_piece, F.data == "turn2_cancel")
async def process_to_cancel(message: Message, state: FSMContext) -> None:
    user_board = UserBoard(message.from_user.id)
    board = await user_board.get()
    if board is None:
        return

    await state.clear()
    await message.message.edit_reply_markup(
        reply_markup=first_tap_keyboard(board),
    )


@dp.callback_query(ChessMatchStatesGroup.has_chosen_piece, F.data.startswith("turn2_"))
async def process_to_move(message: Message, state: FSMContext) -> None:
    user_board = UserBoard(message.from_user.id)
    board = await user_board.get()
    if board is None:
        return

    async def send_end_of_game():
        await message.message.answer_photo(
            photo=BufferedInputFile(generate_board_image(board), "game.png"),
            caption=f"Игра закончена! Результат - {outcome.result()}. Новая игра - /newgame"
        )
        await message.message.delete()
        await state.clear()
        await user_board.delete()

    square_name = message.data.removeprefix("turn2_")
    data = await state.get_data()

    outcome = make_move(board, f"{data.get('square')}{square_name}")
    if outcome is not None:
        await send_end_of_game()
        return

    outcome = make_bot_move(board)
    if outcome is not None:
        await send_end_of_game()
        return

    await user_board.update(board)
    await message.message.answer_photo(
        photo=BufferedInputFile(generate_board_image(board), "game.png"),
        caption=f"Делайте следующий ход!",
        reply_markup=first_tap_keyboard(board),
    )
    await message.message.delete()
    await state.clear()


async def main() -> None:
    bot = Bot(token=API_TOKEN, default=DefaultBotProperties(parse_mode=ParseMode.HTML))
    await dp.start_polling(bot)


if __name__ == "__main__":
    logging.basicConfig(level=logging.INFO, stream=sys.stdout)
    asyncio.run(main())

Теперь, даже если в середине игры перезапустить бота, то командой /newgame можно продолжить игру.

Однако, если бот перезапустится после первого нажатия на ход, но до второго нажатия на клавиатуру, то вторая клавиатура работать не будет. Это происходит потому, что, помимо шахматной доски, бот также сохраняет состояния в FSM (в предыдущей статье это подробно описано). А состояния по умолчанию сохраняются в памяти и теряются при перезапуске программы.

Решить эту проблему можно несколькими способами:

  1. Включив другое хранилище FSM (Redis, или MongoDB). Требует установки соответствующего сервиса на компьютер.
  2. Написав собственное хранилище, например, в той же базе SQLite.
  3. Избавившись от FSM, или же избавившись полностью от состояний (это возможно, поскольку оно у нас одно).

Какой из способов выбрать - решать Вам. Мы рассмотрим эту проблему в одной из следующих статей.

Архив с кодом Telegram-бота (без файла с токеном) можно скачать здесь.

Популярные статьи

Объект range

Изучение объекта range в Python. Создание последовательностей чисел, использование в циклах и примеры применения.

PEP 257 - соглашения для строк документации (docstrings)

Целью данного PEP является стандартизация структуры строк документации: что они должны содержать и что должны объяснять.

Бесплатные курсы Python

Обзор бесплатных курсов, обучающих видео по языку программирования Python