Aiogram Bot #2. Registering a user. Working with SQL Alchemy

Всех приветствую, снова. Сегодня мы продолжаем писать функционального бота на aiogram. В конце вы можете увидеть весь код, который загружен на GitHub, а также другие полезные ссылки, но сначала прочитайте статью :)

В данном уроке мы реализуем:
- создадим базу данных для бота
- подключим и настроим СУБД PostgreSQL
- сделаем регистрацию для пользователя
- создадим команду для показа профиля

Давайте начнём!

Для начала нам нужно скачать PostgreSQL. Установка простая, думаю каждый справиться. После чего у нас появляется ярлык pgAdmin 4 (если не появился, то ищем через поиск).

Открываем. При входе нам нужно придумать пароль от pgAdmin. Этот пароль вы будете вводить каждый раз при повторном входе, а также когда будете запускать базу данных на своём компьютере. В общем запомните пароль.

Далее нам открывается такая картина, где мы видим строчку Databases. Она нам не нужна. Вместо этого мы наводим на PostgreSQL 13 > Create > Database.



В появившимся окне вводим любое название базы данных (я назвал её lolz). После чего нажимаем Save.



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



Теперь давайте откроем наш проект. Кто не смотрел прошлую статью, рекомендуется ознакомиться.

Итак, для начала давайте подключим нашу базу данных к проекту. Мы будем использовать данный stack технологий:
- psycopg2
- SQL Alchemy

Но сначала нам надо установить данные библиотеки, поэтому в терминале проекта прописываем:

pip install psycopg2

а также

pip install sqlalchemy

Отлично! Теперь создадим подключение к БД, а также модель юзера в файле utils/db_api/schemas/user.py (файл также создадим).

import os

from sqlalchemy import create_engine, Column, Integer, Boolean, String
from dotenv import load_dotenv
from sqlalchemy.orm import scoped_session, declarative_base, sessionmaker

load_dotenv()

host = str(os.getenv("HOST"))
password = str(os.getenv("PASSWORD"))
database = str(os.getenv("DATABASE"))

engine = create_engine(f"postgresql+psycopg2://postgres:{password}@{host}/{database}")

session = scoped_session(sessionmaker(bind=engine))
Base = declarative_base()
Base.query = session.query_property()

Python

import os



from sqlalchemy import create_engine, Column, Integer, Boolean, String

from dotenv import load_dotenv

from sqlalchemy.orm import scoped_session, declarative_base, sessionmaker



load_dotenv()



host = str(os.getenv("HOST"))

password = str(os.getenv("PASSWORD"))

database = str(os.getenv("DATABASE"))



engine = create_engine(f"postgresql+psycopg2://postgres:{password}@{host}/{database}")



session = scoped_session(sessionmaker(bind=engine))

Base = declarative_base()

Base.query = session.query_property()

В начале я записываю в переменные host, password и database значения из переменных окружения (данный принцип мы разбирали в прошлой статье).

Далее мы создаём соединение с нашей postgres, где передаём password, host, а также database. После мы инициализируем хранилище уже созданных сессий с помощью scoped_session, где передаём Sessionmaker.

Sessionmaker – фабрика для создания экземпляров Session с заданными параметрами. Это просто штука, которая немного упрощает жизнь: вместо того, чтобы каждый раз указывать список аргументов у сессии, его достаточно один раз указать у фабрики, а дальше уже создавать сессии без указания аргументов.

Base = declarative_base()

Здесь мы наследуемся от декларативного базового класса, и в конце настраиваем запросы для нашей Base.

Теперь давайте создадим начальную модель пользователя:

class User(Base):
__tablename__ = 'users'

id = Column(Integer, primary_key=True)
username = Column(String)
name = Column(String)
admin = Column(Boolean, default=False)

Python

class User(Base):

    __tablename__ = 'users'



    id = Column(Integer, primary_key=True)

    username = Column(String)

    name = Column(String)

    admin = Column(Boolean, default=False)

Класс User наследуется от Base. В начале, строка

__tablename__ = 'users'

отвечает за название таблицы с пользователя (в данном случае users). После чего указываем колонки таблицы, а также их типы в скобках.

Атрибут primary_key отвечает за некую уникальность той или иной колонки. В данном случает id пользователя будет уникальным, то есть не будет повторяться. При повторе у нас возникнет исключение.

Атрибут default отвечает за стандартное значение колонки, в данном случае, когда пользователь зарегистрировался его admin = False.

В конце файла пишем


Base.metadata.create_all(bind=engine)


Данная конструкция отвечает за создания всех таблиц, если их нет.

Отлично! Мы с вами полностью настроили подключение к базе данных, а также описали модель пользователя. Теперь давайте перейдём к его непосредственному созданию, регистрации.

Создаём файл utils/db_api/db_quick_commands.py. В нём мы будем описывать все быстрые (и не только) взаимодействия с базой данных.

В этом файле мы создаём функцию регистрации юзера:

from sqlalchemy.exc import PendingRollbackError, IntegrityError

from bot.utils.db_api.schemas.user import session, User


def register_user(message):
username = message.from_user.username if message.from_user.username else None
user = User(id=int(message.from_user.id), username=username, name=message.from_user.full_name)

session.add(user)

try:
session.commit()
return True
except IntegrityError:
session.rollback() # откатываем session.add(user)
return False

Python

from sqlalchemy.exc import PendingRollbackError, IntegrityError



from bot.utils.db_api.schemas.user import session, User





def register_user(message):

    username = message.from_user.username if message.from_user.username else None

    user = User(id=int(message.from_user.id), username=username, name=message.from_user.full_name)



    session.add(user)



    try:

        session.commit()

        return True

    except IntegrityError:

        session.rollback()  # откатываем session.add(user)

        return False

Функция принимает message от юзера.

В переменную username мы записываем username пользователя, если он имеется, иначе записываем полное имя.
В переменную user мы записываем объект модели User (которую недавно создали), где приравниваем все данные для регистрации кроме admin, так как у него default=False, и записывать его тут - не обязательно.

После чего через текущую сессию добавляем user в базу данных через метод add()
Далее мы ставим обработчик ошибок и в блок try записываем:

try:
session.commit()
return True

Python

    try:

        session.commit()

        return True

С помощью

session.commit()

мы подтверждаем все изменения в БД (в данном случае добавления пользователя через add()), а затем возвращаем True что в будущем даст нам понять об успешной регистрации юзера. Иначе, у нас есть блок except:

except IntegrityError:
session.rollback() # откатываем session.add(user)
return False

Python

    except IntegrityError:

        session.rollback()  # откатываем session.add(user)

        return False

Где в случае ошибки, мы, с помощью метода rollback() откатываем предыдущее изменение в БД (создания юзера), после - возвращаем False.

Отлично! Теперь переходим в файл handlers/users/start.py и переписываем ранее нами созданную команду /start:

from aiogram import types

from bot.loader import dp
from bot.utils.db_api.db_quick_commands import register_user


@dp.message_handler(commands=['start'])
async def command_start(message: types.Message):
user = register_user(message)

if user:
await message.answer('Вы успешно зарегистрировались!')
else:
await message.answer('Вы уже зарегистрированы!')

Python

from aiogram import types



from bot.loader import dp

from bot.utils.db_api.db_quick_commands import register_user





@dp.message_handler(commands=['start'])

async def command_start(message: types.Message):

    user = register_user(message)



    if user:

        await message.answer('Вы успешно зарегистрировались!')

    else:

        await message.answer('Вы уже зарегистрированы!')

В переменную user мы записываем ранее нами созданную функцию register_user(), где передаём message. После чего проверяем переменную user. Если она вернула True, то пользователь успешно создан, иначе - юзер уже был зарегистрирован.

Есть! Мы написали с вами регистрацию пользователя. Теперь давайте сделаем вывод информации о пользователе из БД. Но чтобы выводить информацию о пользователе, нам сначала нужно её получать. Поэтому возвращаемся в файлик db_quick_commands.py и создаем новую функцию select_user():

def select_user(user_id):
user = session.query(User).filter(User.id == user_id).first()
return user

Python

def select_user(user_id):

    user = session.query(User).filter(User.id == user_id).first()

    return user

В аргументе функции мы принимаем user_id пользователя, которого хотим получить из БД. Далее создаём переменную user куда передаём объект пользователя. В запросе мы выбираем строку где User.id (id юзера) == user_id (id которое мы передали). В конце пишем first() - это значит что нам нужна только первая запись, которая нашлась. После чего возвращаем user.

Теперь мы можем получать пользователя из базы данных, поэтому в папке handlers/users создаём файл profile.py и прописываем данные строчки:

from aiogram import types

from bot.loader import dp
from bot.utils.db_api.db_quick_commands import select_user


@dp.message_handler(commands=['profile'])
async def show_profile(message: types.Message):
user = select_user(message.from_user.id)

await message.answer(f"Твой профиль\n"
f"Name: {user.name}\n"
f"Username: @{user.username}\n"
f"Admin: {'Да' if user.admin else 'Нет'}")

Python

from aiogram import types



from bot.loader import dp

from bot.utils.db_api.db_quick_commands import select_user





@dp.message_handler(commands=['profile'])

async def show_profile(message: types.Message):

    user = select_user(message.from_user.id)



    await message.answer(f"Твой профиль\n"

                         f"Name: {user.name}\n"

                         f"Username: @{user.username}\n"

                         f"Admin: {'Да' if user.admin else 'Нет'}")

Тут мы создаём переменную user и присваиваем ей объект найденного пользователя. А так как это объект класса user, то мы можем обращаться к его атрибутам и в данном случае выводить всю интересующую нас информацию юзеру.

В строке

f"Admin: {'Да' if user.admin else 'Нет'}")

Мы выводим "Да", если user.admin == True, иначе выводим "Нет".

Теперь мы должны проинициализировать этот файл в handlers/users/__init__.py


from .profile import dp


Также давайте добавим команду /profile в команды бота:

from aiogram import types


async def set_default_commands(dp):
await dp.bot.set_my_commands(
[
types.BotCommand('start', 'Запустить бота'),
types.BotCommand('profile', 'Посмотреть профиль'),
]
)

Python

from aiogram import types





async def set_default_commands(dp):

    await dp.bot.set_my_commands(

        [

            types.BotCommand('start', 'Запустить бота'),

            types.BotCommand('profile', 'Посмотреть профиль'),

        ]

    )

Отлично теперь всё протестируем!




И как видите, всё отлично работает!

А если вам понравился данный урок, то дайте об этом знать ;) Всем спасибо и скоро увидимся!

Полезные ссылки:
GitHub проекта - https://github.com/Kolini77/aiogram-lessons
Telegram чат - https://t.me/+q9HNGuHpkpA1Y2Y6​