From 5865cc450f8cde04e65425793ca55ad1c2be3a99 Mon Sep 17 00:00:00 2001 From: Allan Eising Date: Mon, 19 May 2025 09:22:02 +0200 Subject: [PATCH] Implement async db access in admin --- .../src/sshecret_admin/auth/__init__.py | 2 + .../src/sshecret_admin/auth/authentication.py | 11 +++++ .../src/sshecret_admin/core/db.py | 46 ++++++++++++++++++- .../src/sshecret_admin/core/dependencies.py | 1 + .../src/sshecret_admin/core/settings.py | 5 ++ .../sshecret_admin/frontend/dependencies.py | 7 ++- .../src/sshecret_admin/frontend/router.py | 12 ++++- .../src/sshecret_admin/frontend/views/auth.py | 8 ++-- 8 files changed, 85 insertions(+), 7 deletions(-) diff --git a/packages/sshecret-admin/src/sshecret_admin/auth/__init__.py b/packages/sshecret-admin/src/sshecret_admin/auth/__init__.py index e83c7f5..da34765 100644 --- a/packages/sshecret-admin/src/sshecret_admin/auth/__init__.py +++ b/packages/sshecret-admin/src/sshecret_admin/auth/__init__.py @@ -2,6 +2,7 @@ from .authentication import ( authenticate_user, + authenticate_user_async, create_access_token, create_refresh_token, check_password, @@ -16,6 +17,7 @@ __all__ = [ "Token", "User", "authenticate_user", + "authenticate_user_async", "check_password", "create_access_token", "create_refresh_token", diff --git a/packages/sshecret-admin/src/sshecret_admin/auth/authentication.py b/packages/sshecret-admin/src/sshecret_admin/auth/authentication.py index 7f0ad57..8329d4b 100644 --- a/packages/sshecret-admin/src/sshecret_admin/auth/authentication.py +++ b/packages/sshecret-admin/src/sshecret_admin/auth/authentication.py @@ -8,6 +8,7 @@ import bcrypt import jwt from sqlalchemy import select +from sqlalchemy.ext.asyncio import AsyncSession from sqlalchemy.orm import Session from sshecret_admin.core.settings import AdminServerSettings @@ -72,6 +73,16 @@ def check_password(plain_password: str, hashed_password: str) -> None: raise AuthenticationFailedError() +async def authenticate_user_async(session: AsyncSession, username: str, password: str) -> User | None: + """Authenticate user async.""" + user = (await session.scalars(select(User).where(User.username == username))).first() + if not user: + return None + if not verify_password(password, user.hashed_password): + return None + return user + + def authenticate_user(session: Session, username: str, password: str) -> User | None: """Authenticate user.""" user = session.scalars(select(User).where(User.username == username)).first() diff --git a/packages/sshecret-admin/src/sshecret_admin/core/db.py b/packages/sshecret-admin/src/sshecret_admin/core/db.py index 550ece3..07d8e82 100644 --- a/packages/sshecret-admin/src/sshecret_admin/core/db.py +++ b/packages/sshecret-admin/src/sshecret_admin/core/db.py @@ -1,11 +1,15 @@ """Database setup.""" -from collections.abc import Generator, Callable +from contextlib import asynccontextmanager + +from collections.abc import AsyncIterator, Generator, Callable from sqlalchemy.orm import Session from sqlalchemy.engine import URL from sqlalchemy import create_engine, Engine +from sqlalchemy.ext.asyncio import AsyncConnection, create_async_engine, AsyncEngine, AsyncSession, async_sessionmaker + def setup_database( db_url: URL | str, @@ -20,3 +24,43 @@ def setup_database( yield session return engine, get_db_session + + +class DatabaseSessionManager: + def __init__(self, host: URL | str, **engine_kwargs: str): + self._engine: AsyncEngine | None = create_async_engine(host, **engine_kwargs) + self._sessionmaker: async_sessionmaker[AsyncSession] | None = async_sessionmaker(autocommit=False, bind=self._engine, expire_on_commit=False) + + async def close(self): + if self._engine is None: + raise Exception("DatabaseSessionManager is not initialized") + await self._engine.dispose() + + self._engine = None + self._sessionmaker = None + + @asynccontextmanager + async def connect(self) -> AsyncIterator[AsyncConnection]: + if self._engine is None: + raise Exception("DatabaseSessionManager is not initialized") + + async with self._engine.begin() as connection: + try: + yield connection + except Exception: + await connection.rollback() + raise + + @asynccontextmanager + async def session(self) -> AsyncIterator[AsyncSession]: + if self._sessionmaker is None: + raise Exception("DatabaseSessionManager is not initialized") + + session = self._sessionmaker() + try: + yield session + except Exception: + await session.rollback() + raise + finally: + await session.close() diff --git a/packages/sshecret-admin/src/sshecret_admin/core/dependencies.py b/packages/sshecret-admin/src/sshecret_admin/core/dependencies.py index 358cc7a..176f180 100644 --- a/packages/sshecret-admin/src/sshecret_admin/core/dependencies.py +++ b/packages/sshecret-admin/src/sshecret_admin/core/dependencies.py @@ -4,6 +4,7 @@ from collections.abc import AsyncGenerator, Awaitable, Callable, Generator from dataclasses import dataclass from typing import Self +from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker from sqlalchemy.orm import Session from sshecret_admin.auth import User from sshecret_admin.services import AdminBackend diff --git a/packages/sshecret-admin/src/sshecret_admin/core/settings.py b/packages/sshecret-admin/src/sshecret_admin/core/settings.py index 5ef2e48..e4b3321 100644 --- a/packages/sshecret-admin/src/sshecret_admin/core/settings.py +++ b/packages/sshecret-admin/src/sshecret_admin/core/settings.py @@ -31,3 +31,8 @@ class AdminServerSettings(BaseSettings): def admin_db(self) -> URL: """Construct database url.""" return URL.create(drivername="sqlite", database=self.database) + + @property + def async_db_url(self) -> URL: + """Construct database url with sync handling.""" + return URL.create(drivername="sqlite+aiosqlite", database=self.database) diff --git a/packages/sshecret-admin/src/sshecret_admin/frontend/dependencies.py b/packages/sshecret-admin/src/sshecret_admin/frontend/dependencies.py index 24ca4e3..e4516db 100644 --- a/packages/sshecret-admin/src/sshecret_admin/frontend/dependencies.py +++ b/packages/sshecret-admin/src/sshecret_admin/frontend/dependencies.py @@ -1,10 +1,11 @@ """Frontend dependencies.""" from dataclasses import dataclass -from collections.abc import Callable, Awaitable +from collections.abc import AsyncGenerator, Callable, Awaitable from typing import Self from sqlalchemy.orm import Session +from sqlalchemy.ext.asyncio import AsyncSession from jinja2_fragments.fastapi import Jinja2Blocks from fastapi import Request @@ -14,6 +15,7 @@ from sshecret_admin.auth.models import User UserTokenDep = Callable[[Request, Session], Awaitable[User]] UserLoginDep = Callable[[Request, Session], Awaitable[bool]] +AsyncSessionDep = Callable[[], AsyncGenerator[AsyncSession, None]] @dataclass @@ -25,6 +27,7 @@ class FrontendDependencies(BaseDependencies): get_user_from_access_token: UserTokenDep get_user_from_refresh_token: UserTokenDep get_login_status: UserLoginDep + get_async_session: AsyncSessionDep @classmethod def create( @@ -35,6 +38,7 @@ class FrontendDependencies(BaseDependencies): get_user_from_access_token: UserTokenDep, get_user_from_refresh_token: UserTokenDep, get_login_status: UserLoginDep, + get_async_session: AsyncSessionDep ) -> Self: """Create from base dependencies.""" return cls( @@ -45,4 +49,5 @@ class FrontendDependencies(BaseDependencies): get_user_from_access_token=get_user_from_access_token, get_user_from_refresh_token=get_user_from_refresh_token, get_login_status=get_login_status, + get_async_session=get_async_session, ) diff --git a/packages/sshecret-admin/src/sshecret_admin/frontend/router.py b/packages/sshecret-admin/src/sshecret_admin/frontend/router.py index e12cef2..db3e870 100644 --- a/packages/sshecret-admin/src/sshecret_admin/frontend/router.py +++ b/packages/sshecret-admin/src/sshecret_admin/frontend/router.py @@ -19,6 +19,7 @@ from starlette.datastructures import URL from sshecret_admin.auth import PasswordDB, User, decode_token from sshecret_admin.core.dependencies import BaseDependencies from sshecret_admin.services.admin_backend import AdminBackend +from sshecret_admin.core.db import DatabaseSessionManager from .dependencies import FrontendDependencies from .exceptions import RedirectException @@ -47,7 +48,9 @@ def create_router(dependencies: BaseDependencies) -> APIRouter: session: Annotated[Session, Depends(dependencies.get_db_session)] ): """Get admin backend API.""" - password_db = session.scalars(select(PasswordDB).where(PasswordDB.id == 1)).first() + password_db = session.scalars( + select(PasswordDB).where(PasswordDB.id == 1) + ).first() if not password_db: raise HTTPException( 500, detail="Error: The password manager has not yet been set up." @@ -116,6 +119,12 @@ def create_router(dependencies: BaseDependencies) -> APIRouter: return False return True + async def get_async_session(): + """Get async session.""" + sessionmanager = DatabaseSessionManager(dependencies.settings.async_db_url) + async with sessionmanager.session() as session: + yield session + view_dependencies = FrontendDependencies.create( dependencies, get_admin_backend, @@ -123,6 +132,7 @@ def create_router(dependencies: BaseDependencies) -> APIRouter: get_user_from_access_token, get_user_from_refresh_token, get_login_status, + get_async_session, ) app.include_router(audit.create_router(view_dependencies)) diff --git a/packages/sshecret-admin/src/sshecret_admin/frontend/views/auth.py b/packages/sshecret-admin/src/sshecret_admin/frontend/views/auth.py index c9359b0..175d0d6 100644 --- a/packages/sshecret-admin/src/sshecret_admin/frontend/views/auth.py +++ b/packages/sshecret-admin/src/sshecret_admin/frontend/views/auth.py @@ -8,13 +8,13 @@ from fastapi import APIRouter, Depends, Query, Request, Response, status from fastapi.responses import RedirectResponse from fastapi.security import OAuth2PasswordRequestForm -from sqlalchemy.orm import Session +from sqlalchemy.ext.asyncio import AsyncSession from sshecret_admin.services import AdminBackend from starlette.datastructures import URL from sshecret_admin.auth import ( User, - authenticate_user, + authenticate_user_async, create_access_token, create_refresh_token, ) @@ -80,7 +80,7 @@ def create_router(dependencies: FrontendDependencies) -> APIRouter: async def login_user( request: Request, response: Response, - session: Annotated[Session, Depends(dependencies.get_db_session)], + session: Annotated[AsyncSession, Depends(dependencies.get_async_session)], admin: Annotated[AdminBackend, Depends(dependencies.get_admin_backend)], form_data: Annotated[OAuth2PasswordRequestForm, Depends()], next: Annotated[str, Query()] = "/dashboard", @@ -100,7 +100,7 @@ def create_router(dependencies: FrontendDependencies) -> APIRouter: }, ) - user = authenticate_user(session, form_data.username, form_data.password) + user = await authenticate_user_async(session, form_data.username, form_data.password) login_failed = RedirectException( to=URL("/login").include_query_params( error_title="Login Error", error_message="Invalid username or password"