Fix circular imports

mybot/__init__.py

@@ -19,7 +19,7 @@ def create_bot(config: Config, i18n: I18N, engine):
                   use_class_middlewares=True,
                   state_storage=state_storage)
     register_handlers(bot)
-    setup_middlewares(bot, i18n)
+    setup_middlewares(bot, i18n, engine)
     add_custom_filters(bot, config)
     return bot
 

mybot/config.py

@@ -54,10 +54,14 @@ class StateStorageConfig:
 @dataclass
 class DatabaseConfig:
     url: str
+    pool_recycle: int
+    pool_pre_ping: bool
 
     @classmethod
     def from_env(cls):
-        return cls(os.getenv("DATABASE_URL", "sqlite:///bot.db"))
+        return cls(os.getenv("DB_URL", "sqlite:///bot.db"),
+                   int(os.getenv("DB_POOL_RECYCLE", 3600)),
+                   bool(int(os.getenv("DB_POOL_PRE_PING", True))))
 
 
 @dataclass

mybot/database/__init__.py

@@ -6,8 +6,8 @@ from ..config import DatabaseConfig
 
 def get_engine(config: DatabaseConfig):
     engine = create_engine(config.url,
-                           pool_recycle=3600,
-                           pool_pre_ping=True)
+                           pool_recycle=config.pool_recycle,
+                           pool_pre_ping=config.pool_pre_ping)
     return engine
 
 

mybot/database/models.py

@@ -8,5 +8,8 @@ class User (Base):
     __tablename__ = "user"
 
     id: Mapped[int] = mapped_column(BIGINT, primary_key=True, unique=True, autoincrement=False)
-    username: Mapped[int] = mapped_column(String(32), unique=True, nullable=True)
+    username: Mapped[str] = mapped_column(String(32), unique=True, nullable=True)
     # additional fields go here
+
+    def __init__(self, id: int, username: str):
+        super().__init__(id=id, username=username)

mybot/middlewares/__init__.py

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

mybot/middlewares/arguments.py

@@ -2,7 +2,7 @@ from telebot.handler_backends import BaseMiddleware
 from telebot.types import Message, CallbackQuery
 
 
-class ExtraArguments(BaseMiddleware):
+class ArgumentsMiddleware (BaseMiddleware):
     def __init__(self, i18n):
         super().__init__()
         self.i18n = i18n

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.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()