refactor: move db and file watcher setup to use async context manager
This commit is contained in:
@@ -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(
|
||||
|
||||
Reference in New Issue
Block a user