Add DatabaseMiddleware

mybot/middlewares/__init__.py

@@ -1,8 +1,10 @@
 from telebot import TeleBot
 
-from .arguments import ArgumentsMiddleware
 from ..i18n import I18N
+from .arguments import ArgumentsMiddleware
+from .database import DatabaseMiddleware
 
 
-def setup_middlewares(bot: TeleBot, i18n: I18N):
+def setup_middlewares(bot: TeleBot, i18n: I18N, engine):
     bot.setup_middleware(ArgumentsMiddleware(i18n))
+    bot.setup_middleware(DatabaseMiddleware(engine))

mybot/middlewares/database.py

@@ -0,0 +1,29 @@
+from telebot.handler_backends import BaseMiddleware
+from telebot.types import Message, CallbackQuery
+
+from sqlalchemy.orm import Session
+
+from ..database 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()