30 lines
946 B
Python
30 lines
946 B
Python
from telebot.handler_backends import BaseMiddleware
|
|
from telebot.types import Message, CallbackQuery
|
|
|
|
from sqlalchemy.orm import Session
|
|
|
|
from ..database.models import User
|
|
|
|
|
|
class DatabaseMiddleware (BaseMiddleware):
|
|
def __init__(self, engine):
|
|
super().__init__()
|
|
self.engine = engine
|
|
self.update_types = ["message", "callback_query"]
|
|
|
|
def pre_process(self, obj: [Message, CallbackQuery], data: dict):
|
|
session = Session(self.engine)
|
|
user = session.get(User, obj.from_user.id)
|
|
if user is None:
|
|
user = User(id=obj.from_user.id, username=obj.from_user.username)
|
|
session.add(user)
|
|
session.commit()
|
|
data["db"] = session
|
|
data["user"] = user
|
|
|
|
def post_process(self, message, data: dict, exception: BaseException):
|
|
if "db" in data:
|
|
session: Session = data["db"]
|
|
session.rollback()
|
|
session.close()
|