Connection.use()

api/kiwi_vpn_api/db/connection.py

@@ -6,6 +6,19 @@ from sqlalchemy.orm import Session, sessionmaker
 from .models import ORMBaseModel
 
 
+class SessionManager:
+    __session: Session
+
+    def __init__(self, session: Session) -> None:
+        self.__session = session
+
+    def __enter__(self) -> Session:
+        return self.__session
+
+    def __exit__(self, *args) -> None:
+        self.__session.close()
+
+
 class Connection:
     engine: Engine | None = None
     session_local: sessionmaker | None = None
@@ -18,6 +31,13 @@ class Connection:
         )
         ORMBaseModel.metadata.create_all(bind=engine)
 
+    @classmethod
+    def use(cls) -> SessionManager | None:
+        if cls.session_local is None:
+            return None
+
+        return SessionManager(cls.session_local())
+
     @classmethod
     async def get(cls) -> Generator[Session | None, None, None]:
         if cls.session_local is None:

api/kiwi_vpn_api/main.py

@@ -39,7 +39,7 @@ async def on_startup() -> None:
         Connection.connect(await current_config.db.db_engine)
 
         # some testing
-        async for db in Connection.get():
+        with Connection.use() as db:
             print(schemas.User.from_db(db, "admin"))
             print(schemas.User.from_db(db, "nonexistent"))
 

api/kiwi_vpn_api/routers/admin.py

@@ -25,7 +25,7 @@ async def install(
     await config.save()
     Connection.connect(await config.db.db_engine)
 
-    async for db in Connection.get():
+    with Connection.use() as db:
         admin_user = schemas.User.create(
             db=db,
             user=admin_user,