Aiogram 3: PostgreSQL - интеграция

https://habr.com/ru/articles/822083/

В этой статье мы напишем:

1. Функцию для создания таблицы пользователей.
2. Функцию, которая проверяет наличие пользователя в базе данных и возвращает его полную информацию или None.
3. Функцию для добавления пользователя в базу данных.
4. Хендлер для команды /start, который проверяет наличие пользователя в базе данных и запускает соответствующий сценарий.
5. Логику для записи пользователя в базу данных после заполнения анкеты.
6. Блок профиля, который отображает всю необходимую информацию о пользователе по клику.

Настройка базы данных PostgreSQL

Для начала убедитесь, что у вас настроена база данных PostgreSQL. Это может быть как локальная машина, так и PostgreSQL на VPS сервере. Подробную инструкцию по настройке PostgreSQL вы найдете в статье: "Разворачиваем PostgreSQL на VPS сервере за 5 минут".

Установка asyncpg-lite

В качестве менеджера для взаимодействия с базой данных PostgreSQL я буду использовать библиотеку asyncpg-lite. Подробный обзор основных функций этой библиотеки я делал тут: "asynpg-lite: Простой асинхронный менеджер для PostgreSQL на Python"

Установим ее:

pip install --upgrade asyncpg-lite

Проверьте, чтобы после установки у вас была версия не ниже 0.22.2.1.

Подключение к базе данных

Для подключения к PostgreSQL через asyncpg-lite можно использовать два способа:

1. Через аргумент dsn с ссылкой вида: postgresql://USER:PASSWORD@HOST_IP:PORT/NAME_DB
2. Через передачу параметров подключения: хост, порт, название базы данных, имя пользователя и т.д.

Оба эти формата подключения к PostgreSQL обхватывает asyncpg-lite

Будем использовать 1 вариант со ссылкой, импортируя ее из .env файла. Также для взаимодействия с asyncpg-lite необходимо передать обязательный параметр deletion_password, который нужен для дополнительной защиты в критических операциях (очистка всех данных, удаление таблицы).

Импорт и создание объекта DatabaseManager

from asyncpg_lite import DatabaseManager
from decouple import config

pg_manager = DatabaseManager(dsn=config('PG_LINK'), deletion_password=config('ROOT_PASS'))

Этот объект будет взаимодействовать с базой данных, выполняя необходимые действия.

Написание функций

Теперь давайте начнем писать код. Нам необходимо будет написать 3 функции:

• Функцию для создания таблицы с пользователями.
• Функцию для получения данных о пользователе из таблицы.
• Функцию для добавления данных о пользователе в таблицу.

Функции пропишу в отдельный файл, а после импортирую их в хендлеры бота.

Выполняем импорты:

from create_bot import pg_manager

Пишем функцию для создания таблицы:

async def create_table_users(table_name='users_reg'):
    async with pg_manager:
        columns = ['user_id INT8 PRIMARY KEY', 'gender VARCHAR(50)', 'age INT',
                   'full_name VARCHAR(255)', 'user_login VARCHAR(255) UNIQUE',
                   'photo TEXT', 'about TEXT', 'date_reg TIMESTAMP DEFAULT CURRENT_TIMESTAMP']
        await pg_manager.create_table(table_name=table_name, columns=columns)

Обратите внимание на синтаксис:

• Взаимодействие с базой данных асинхронное и осуществляется через async with.
• Для генерации таблицы создаем питоновский список, где каждое значение из списка описывает колонку таблицы.
• Выполняем саму генерацию. Для этого передаем два обязательных параметра: название таблицы и список с колонками.

После выполнения этой функции мы создадим таблицу с именем users_reg и с такими полями:

• user_id – BIGINT (большое число, это будет телеграмм айди пользователя)
• gender – Пол, строка максимум в 50 знаков
• age – обычное число для хранения возраста
• full_name – Строка максимальной длиной в 255 символов
• user_login – Логин пользователя максимальной длины в 255 символов с меткой, что значение должно быть уникальным
• photo – тип данных текст, так как мы будем хранить айди фотографии, а они в длине могут превысить 255 знаков
• date_reg – В этой колонке будет фиксироваться дата и время регистрации пользователя. По умолчанию будут устанавливаться текущие данные, а это то, что нам нужно.

Импортируем asyncio, запускаем функцию и смотрим, что у нас получилось:

2024-06-15 22:31:45,817 - asyncpg_lite - INFO - Соединение с базой данных установлено.
2024-06-15 22:31:45,962 - asyncpg_lite - INFO - Таблица users_reg успешно создана.
2024-06-15 22:31:46,112 - asyncpg_lite - INFO - Соединение с базой данных закрыто.

Логи нам подсказывают, что таблица «users_reg» успешно создана. Давайте проверим.

Так выглядит созданная таблица.

Видим, что таблица действительно создана. Теперь давайте напишем функцию, которая будет доставать данные о пользователе из указанной таблицы:

async def get_user_data(user_id: int, table_name='users_reg'):
    async with pg_manager:
        user_info = await pg_manager.select_data(table_name=table_name, where_dict={'user_id': user_id}, one_dict=True)
        return user_info if user_info else None

Функция будет принимать 2 параметра:

• user_id (telegram_id пользователя)
• table_name (имя таблицы, откуда берем данные)

Конструкция await pg_manager.select_data(table_name=table_name, where_dict={'user_id': user_id}, one_dict=True) обеспечивает, что функция вернет или 1 словарь, или пустой список.

Если данные не будут получены, то функция вернет None. Это мы будем использовать и в отображении данных в профиле пользователя и в проверке, существует ли пользователь, чтобы бот понял, запускать сценарий анкетирования или отправлять пользователю клавиатуру.

И напишем функцию для добавления пользователя в базу данных:

async def insert_user(user_data: dict, table_name='users_reg'):
    async with pg_manager:
        await pg_manager.insert_data(table_name=table_name, records_data=user_data)

Функция принимает:

• словарь с данными о пользователе (его мы собираем через FSM сценарий)
• имя таблицы

Этих настроек достаточно для того, чтобы мы сделали интеграцию базы данных с ботом.

Переписываем хендлер start.py

from aiogram import Router, F
from aiogram.filters import CommandStart
from aiogram.fsm.context import FSMContext
from aiogram.types import Message
from aiogram.utils.chat_action import ChatActionSender
from create_bot import bot
from db_handler.db_funk import get_user_data
from keyboards.all_kb import main_kb, gender_kb
from .my_anketa import Form

start_router = Router()

@start_router.message(CommandStart())
async def cmd_start(message: Message, state: FSMContext):
    await state.clear()
    async with ChatActionSender.typing(bot=bot, chat_id=message.chat.id):
        user_info = await get_user_data(user_id=message.from_user.id)

    if user_info:
        await message.answer('Привет. Я вижу, что ты зарегистрирован, а значит тебе можно '
                             'посмотреть, как выглядит твой профиль.', reply_markup=main_kb(message.from_user.id))
    else:
        await message.answer('Привет. Для начала выбери свой пол:', reply_markup=gender_kb())
        await state.set_state(Form.gender)

Из новых импортов мы видим:

from .my_anketa import Form
from db_handler.db_funk import get_user_data

Этими строками мы импортировали функцию для получения информации о пользователе и класс FSM. Form.gender нужен нам для запуска сценария анкетирования.

Теперь давайте рассмотрим сам хендлер.

В самом начале мы на всякий случай очищаем состояние FSM. Далее мы запускаем имитацию набора текста, и в то время пока база данных возвращает какие-то значения, как бы бот печатает.

Такой конструкцией мы замаскировали обращение к базе данных за имитацией ввода (пользователь не будет понимать, что в этот момент ждет ответа от базы данных).

После получения результатов выполнения функции get_user_data мы или запускаем сценарий анкетирования, или отправляем сообщение с клавиатурой и возможностью войти в личный профиль.

Теперь нам остается импортировать функцию insert_user в роутер с анкетированием и записать данные, если пользователь нажал на кнопку «Все верно»:

# сохраняем данные
@questionnaire_router.callback_query(F.data == 'correct', Form.check_state)
async def start_questionnaire_process(call: CallbackQuery, state: FSMContext):
    await call.answer('Данные сохранены')
    user_data = await state.get_data()
    await insert_user(user_data)
    await call.message.edit_reply_markup(reply_markup=None)
    await call.message.answer('Благодарю за регистрацию. Ваши данные успешно сохранены!',
                              reply_markup=main_kb(call.from_user.id))
    await state.clear()

Обратите внимание. Мы добавили 2 строки кода:

# Забрали данные из хранилища FSM
user_data = await state.get_data()

# Передали эти данные в функцию insert_user, тем самым выполнив INSERT операцию по добавлению данных о пользователе в таблицу.
await insert_user(user_data)

Перезапускаем бота и смотрим:

Видим, что после выполнения команды /start бот запустил сценарий анкетирования. Давайте его пройдем до конца:

Видим, что анкета завершена. Давайте теперь попробуем вызвать команду /start.

Отлично! Видим, что бот теперь нас признает. Давайте посмотрим в базу данных, добавлены ли мы:

Видим, что мы появились в таблице.

Теперь давайте настроим профиль пользователя. Для этого нам нужно будет написать клавиатуру, которая будет отправлять кнопку с профилем (для разнообразия ещё добавлю кнопку с админкой, если пользователь есть в списках админов):

def main_kb(user_telegram_id: int):
    kb_list = [[KeyboardButton(text="👤 Профиль")]]
    if user_telegram_id in admins:
        kb_list.append([KeyboardButton(text="⚙️ Админ панель")])
    keyboard = ReplyKeyboardMarkup(keyboard=kb_list,
                                   resize_keyboard=True,
                                   one_time_keyboard=True,
                                   input_field_placeholder="Воспользуйтесь меню:")
    return keyboard

И давайте добавим обработчик, который будет реагировать на команду «Профиль»:

@start_router.message(F.text.contains('Профиль'))
async def start_profile(message: Message, state: FSMContext):
    async with ChatActionSender.typing(bot=bot, chat_id=message.chat.id):
        user_info = await get_user_data(user_id=message.from_user.id)
        profile_message = (
            f"<b>👤 Профиль пользователя:</b>\n"
            f"<b>🆔 ID:</b> {user_info['user_id']}\n"
            f"<b>💼 Логин:</b> @{user_info['user_login']}\n"
            f"<b>📛 Полное имя:</b> {user_info['full_name']}\n"
            f"<b>🧑‍🦰 Пол:</b> {user_info['gender']}\n"
            f"<b>🎂 Возраст:</b> {user_info['age']}\n"
            f"<b>📅 Дата регистрации:</b> {user_info['date_reg']}\n"
            f"<b>📝 О себе:</b> {user_info['about']}\n"
        )

        await message.answer_photo(photo=user_info.get('photo'), caption=profile_message)

Получили данные о пользователе, красиво отформатировали их и просто отправили их в ответ на клик на «Профиль».

Смотрим, что получилось:

Видим, что информация успешно отобразилась, а это значит, что интеграция была завершена успешно.

Ничего сложного во взаимодействии с базой данных PostgreSQL нет, особенно когда работа идет через asyncpg-lite.