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