refactor: move db and file watcher setup to use async context manager

This commit is contained in:
2026-03-09 14:25:24 -04:00
parent 2de0aac23a
commit 2e1d16ef9e

View File

@@ -1,5 +1,6 @@
import asyncio import asyncio
from typing import Any from contextlib import asynccontextmanager
from typing import Any, AsyncGenerator
from chitai.services.book import BookService from chitai.services.book import BookService
from chitai.services.consume import ConsumeDirectoryWatcher from chitai.services.consume import ConsumeDirectoryWatcher
@@ -63,14 +64,15 @@ oauth2_auth = OAuth2PasswordBearerAuth[User](
"/access/signup", "/access/signup",
"/opds", "/opds",
"/schema", "/schema",
"/syncs",
"/users/auth",
], ],
) )
watcher_task: asyncio.Task watcher_task: asyncio.Task
async def startup(): @asynccontextmanager
"""Run setup.""" async def setup_db_connection(app: Litestar) -> AsyncGenerator[None, None]:
# Setup databse # Setup databse
async with settings.alchemy_config.get_session() as db_session: async with settings.alchemy_config.get_session() as db_session:
# Create default library if none exist # Create default library if none exist
@@ -86,21 +88,30 @@ async def startup():
) )
await db_session.commit() await db_session.commit()
# book_service = BookService(session=db_session) try:
yield
finally:
await db_session.aclose()
@asynccontextmanager
async def setup_directory_watcher(app: Litestar) -> AsyncGenerator[None, None]:
# Create book covers directory if it does not exist # Create book covers directory if it does not exist
await create_directory(settings.book_cover_path) await create_directory(settings.book_cover_path)
# Create consume directory # Create consume directory
await create_directory(settings.consume_path) await create_directory(settings.consume_path)
async with settings.alchemy_config.get_session() as db_session:
book_service = BookService(session=db_session)
library_service = LibraryService(session=db_session)
# file_watcher = ConsumeDirectoryWatcher(settings.consume_path, library_service, book_service) file_watcher = ConsumeDirectoryWatcher(settings.consume_path, library_service, book_service)
# watcher_task = asyncio.create_task(file_watcher.init_watcher()) watcher_task = asyncio.create_task(file_watcher.init_watcher())
async def shutdown():
""" Run shutdown tasks. """
try:
yield
finally:
watcher_task.cancel() watcher_task.cancel()
def create_app() -> Litestar: def create_app() -> Litestar:
@@ -114,12 +125,14 @@ def create_app() -> Litestar:
c.PublisherController, c.PublisherController,
c.TagController, c.TagController,
c.OpdsController, c.OpdsController,
c.DeviceController,
c.KosyncController,
create_static_files_router(path="/covers", directories=["./covers"]), create_static_files_router(path="/covers", directories=["./covers"]),
index, index,
healthcheck, healthcheck,
], ],
exception_handlers=exception_handlers, exception_handlers=exception_handlers,
on_startup=[startup], lifespan=[setup_db_connection, setup_directory_watcher],
plugins=[alchemy], plugins=[alchemy],
on_app_init=[oauth2_auth.on_app_init], on_app_init=[oauth2_auth.on_app_init],
openapi_config=OpenAPIConfig( openapi_config=OpenAPIConfig(