Updated fastapi-users sqlalchemy plugin

There were a large number of follow-on effects that had to be resolved

Author: sandalwoodbox

docs/coding/backend.md

@@ -20,11 +20,10 @@ The backend uses [FastAPI](fastapi.tiangolo.com/), an async [Python](https://www
 2. Install [homebrew](https://brew.sh/), a package manager for Mac OS.
 3. In a terminal, run:
    ```
-   brew install pyenv libpq openssl
+   brew install pyenv
    ```
-4. Follow the printed instructions to make sure [pyenv](https://github.com/pyenv/pyenv) and libpq are initialized when your terminal starts up. It should be something like:
+4. Follow the printed instructions to make sure [pyenv](https://github.com/pyenv/pyenv) is initialized when your terminal starts up. It should be something like:
    ```
-   echo 'export PATH="/usr/local/opt/libpq/bin:$PATH"' >> ~/.zshrc
    echo 'eval "$(pyenv init --path)"' >> ~/.zprofile
    echo 'eval "$(pyenv init -)"' >> ~/.zshrc
    ```

fanviddb/api_keys/db.py

@@ -8,64 +8,73 @@
 from passlib import pwd  # type: ignore
 from passlib.context import CryptContext  # type: ignore
 from sqlalchemy import select
+from sqlalchemy.exc import IntegrityError
+from sqlalchemy.ext.asyncio import AsyncSession
 from starlette.exceptions import HTTPException
 from starlette.status import HTTP_401_UNAUTHORIZED
 
-from fanviddb.db import database
+from fanviddb.db import get_async_session
 
-from . import db
+from . import models
 
 X_API_KEY = APIKeyHeader(name="X-API-Key", auto_error=False)
 api_key_context = CryptContext(schemes=["bcrypt"], deprecated="auto")
 
 
-async def generate() -> str:
+async def generate(session: AsyncSession) -> str:
     api_key = None
     attempts = 0
     while api_key is None:
         pk = pwd.genword(length=8, charset="ascii_50")
         secret = pwd.genword(entropy="secure", charset="ascii_50")
         api_key = f"{pk}_{secret}"
 
-        query = db.api_keys.insert().values(
+        instance = models.ApiKeyTable(
             pk=pk,
             hashed_api_key=api_key_context.hash(api_key),
             created_timestamp=datetime.datetime.utcnow(),
             state="active",
         )
-        transaction = await database.transaction()
-        try:
-            await database.execute(query)
-        except UniqueViolationError:
-            api_key = None
-            await transaction.rollback()
-        else:
-            await transaction.commit()
+        async with session.begin_nested() as nested:
+            try:
+                session.add(instance)
+                await nested.commit()
+            except IntegrityError as exc:
+                await nested.rollback()
+                # This is necessary because of how sqlalchemy currently
+                # nests errors.
+                if exc.orig.sqlstate != UniqueViolationError.sqlstate:
+                    raise
+                api_key = None
         attempts += 1
         if attempts > 10:
             raise ValueError(_("Too many collisions"))
 
     return api_key
 
 
-async def verify(api_key: str):
+async def verify(session: AsyncSession, api_key: str):
     try:
         pk, _ = api_key.split("_")
     except ValueError:
         return False
-    query = select([db.api_keys]).where(db.api_keys.c.pk == pk)
-    result = await database.fetch_one(query)
-    if result is None:
+    query = select([models.api_keys]).where(models.api_keys.c.pk == pk)
+    result = await session.execute(query)
+    row = result.first()
+    if row is None:
         return False
-    return api_key_context.verify(api_key, result["hashed_api_key"])
+    return api_key_context.verify(api_key, row.hashed_api_key)
 
 
-async def check_api_key_header(api_key: Optional[str] = Depends(X_API_KEY)):
+async def check_api_key_header(
+    session: AsyncSession = Depends(get_async_session),
+    api_key: Optional[str] = Depends(X_API_KEY),
+):
     # If no API key is present, that's fine; however, an invalid or revoked api key
     # is always an error.
     if api_key is None:
         return False
-    is_valid = await verify(api_key)
+    is_valid = await verify(session, api_key)
     if not is_valid:
         raise HTTPException(
             status_code=HTTP_401_UNAUTHORIZED, detail=_("Invalid api key")

fanviddb/api_keys/helpers.py

@@ -1,5 +1,24 @@
-from pydantic import BaseModel
+from sqlalchemy import Column
+from sqlalchemy import DateTime
+from sqlalchemy import String
 
+from fanviddb.db import Base
 
-class ApiKey(BaseModel):
-    api_key: str
+
+class ApiKeyTable(Base):
+
+    __tablename__ = "api_keys"
+
+    # pk functions as a public "username" so that we can find the
+    # correct hashed secret to check.
+    pk = Column(String(), primary_key=True)
+    hashed_api_key = Column(String(), nullable=False)
+
+    # Admin-only
+    state = Column(String())
+
+    # Internal
+    created_timestamp = Column(DateTime())
+
+
+api_keys = ApiKeyTable.__table__

fanviddb/api_keys/models.py

@@ -1,12 +1,16 @@
 from fastapi import APIRouter
+from fastapi import Depends
+from sqlalchemy.ext.asyncio import AsyncSession
+
+from fanviddb.db import get_async_session
 
 from .helpers import generate
-from .models import ApiKey
+from .schema import ApiKey
 
 api_key_router = APIRouter()
 
 
 @api_key_router.post("", response_model=ApiKey)
-async def create_api_key():
-    api_key = await generate()
+async def create_api_key(session: AsyncSession = Depends(get_async_session)):
+    api_key = await generate(session)
     return {"api_key": api_key}

fanviddb/api_keys/router.py

@@ -0,0 +1,5 @@
+from pydantic import BaseModel
+
+
+class ApiKey(BaseModel):
+    api_key: str

fanviddb/api_keys/schema.py

@@ -8,7 +8,6 @@
 from .api_keys.router import api_key_router
 from .auth.routers import auth_router
 from .auth.routers import users_router
-from .db import database
 from .email import EmailSendFailed
 from .fanvids.router import router as fanvid_router
 
@@ -26,42 +25,42 @@ async def check_config(self) -> None:
 
 main_app = Starlette()
 
-api = FastAPI(docs_url=None)
+api_app = FastAPI(docs_url=None)
 
-api.include_router(
+api_app.include_router(
     fanvid_router,
     prefix="/fanvids",
     tags=["Fanvids"],
 )
-api.include_router(
+api_app.include_router(
     auth_router,
     prefix="/auth",
     tags=["Auth"],
 )
-api.include_router(
+api_app.include_router(
     users_router,
     prefix="/users",
     tags=["Users"],
 )
-api.include_router(
+api_app.include_router(
     api_key_router,
     prefix="/api_keys",
     tags=["API Keys"],
 )
 
 
-frontend = Starlette()
+frontend_app = Starlette()
 
 
-@frontend.middleware("http")
+@frontend_app.middleware("http")
 async def default_response(request, call_next):
     response = await call_next(request)
     if response.status_code == 404:
         return FileResponse("frontend/build/index.html")
     return response
 
 
-frontend.mount(
+frontend_app.mount(
     "/", HelpfulStaticFiles(directory="frontend/build/", html=True, check_dir=False)
 )
 main_app.mount(
@@ -78,27 +77,17 @@ async def default_response(request, call_next):
 )
 main_app.mount(
     "/api",
-    api,
+    api_app,
     name="api",
 )
 main_app.mount(
     "/",
-    frontend,
+    frontend_app,
     name="frontend",
 )
 
 
-@main_app.on_event("startup")
-async def startup():
-    await database.connect()
-
-
-@main_app.on_event("shutdown")
-async def shutdown():
-    await database.disconnect()
-
-
-@api.exception_handler(EmailSendFailed)
+@api_app.exception_handler(EmailSendFailed)
 async def email_send_failed_handler(__: Request, exc: EmailSendFailed):
     return JSONResponse(
         status_code=503,

fanviddb/app.py

@@ -12,22 +12,23 @@
 from fastapi_users.authentication import AuthenticationBackend
 from fastapi_users.authentication import CookieTransport
 from fastapi_users.authentication import JWTStrategy
+from fastapi_users.db import SQLAlchemyUserDatabase
+from sqlalchemy.ext.asyncio import AsyncSession
 from sqlalchemy.sql import exists
 from sqlalchemy.sql import select
 from zxcvbn import zxcvbn  # type: ignore
 
 from fanviddb import conf
-from fanviddb.db import database
+from fanviddb.db import get_async_session
 from fanviddb.email import send_email
 from fanviddb.i18n.utils import get_fluent
 from fanviddb.i18n.utils import get_request_locales
 
-from .db import get_user_db
-from .db import users
-from .models import User
-from .models import UserCreate
-from .models import UserDB
-from .models import UserUpdate
+from .models import UserTable
+from .schema import User
+from .schema import UserCreate
+from .schema import UserDB
+from .schema import UserUpdate
 
 AUTH_LIFETIME = 60 * 60 * 24 * 14
 cookie_transport = CookieTransport(
@@ -53,6 +54,7 @@ class UserManager(BaseUserManager[UserCreate, UserDB]):
     reset_password_token_lifetime_seconds = 60 * 5
     verification_token_secret = conf.EMAIL_TOKEN_SECRET_KEY
     verification_token_lifetime_seconds = 60 * 5
+    user_db: SQLAlchemyUserDatabase
 
     async def on_after_forgot_password(
         self, user: UserDB, token: str, request: Optional[Request] = None
@@ -130,16 +132,21 @@ async def create(
         request: Optional[Request] = None,
     ) -> UserDB:
         user = cast(UserCreate, user)
-        query = select([exists().where(users.c.username == user.username)])
-        result = await database.fetch_one(query)
-        if result and result.get("anon_1"):
+        query = select([exists().where(UserTable.username == user.username)])
+        result = await self.user_db.session.execute(query)
+        row = result.first()
+        if row and row._mapping.get("anon_1"):
             raise HTTPException(
                 status_code=status.HTTP_400_BAD_REQUEST,
                 detail="REGISTER_USERNAME_ALREADY_EXISTS",
             )
         return await super().create(user, safe, request)
 
 
+async def get_user_db(session: AsyncSession = Depends(get_async_session)):
+    yield SQLAlchemyUserDatabase(UserDB, session, UserTable)
+
+
 async def get_user_manager(user_db=Depends(get_user_db)):
     yield UserManager(user_db)
 

fanviddb/auth/db.py

@@ -1,19 +1,12 @@
-from typing import Optional
+from fastapi_users.db import SQLAlchemyBaseUserTable
+from sqlalchemy import Column
+from sqlalchemy import String
 
-from fastapi_users import models
+from fanviddb.db import Base
 
 
-class User(models.BaseUser):
-    username: Optional[str]
+class UserTable(Base, SQLAlchemyBaseUserTable):
+    username = Column(String(length=40), nullable=False, unique=True)
 
 
-class UserCreate(models.BaseUserCreate):
-    username: str
-
-
-class UserUpdate(models.BaseUserUpdate):
-    username: Optional[str]
-
-
-class UserDB(User, models.BaseUserDB):
-    username: str
+users = UserTable.__table__  # type: ignore