refactor code to add typing and fix linting issues

Signed-off-by: Grant Ramsay <seapagan@gmail.com>

Author: seapagan

db.py

@@ -1,33 +1,57 @@
 """Set up the database connection and session.""" ""
-from sqlalchemy.ext.asyncio import async_sessionmaker, create_async_engine
-from sqlalchemy.orm import declarative_base
+from collections.abc import AsyncGenerator
+from typing import Any
+
+from sqlalchemy import MetaData
+from sqlalchemy.ext.asyncio import (
+    AsyncSession,
+    async_sessionmaker,
+    create_async_engine,
+)
+from sqlalchemy.orm import DeclarativeBase
 
 DATABASE_URL = "postgresql+asyncpg://postgres:postgres@localhost/postgres"
-# DATABASE_URL = "sqlite+aiosqlite:///./test.db"
+# DATABASE_URL = "sqlite+aiosqlite:///./test.db"  # noqa: ERA001
 # Note that (as far as I can tell from the docs and searching) there is no need
 # to add 'check_same_thread=False' to the sqlite connection string, as
 # SQLAlchemy version 1.4+ will automatically add it for you when using SQLite.
 
-engine = create_async_engine(DATABASE_URL, echo=False)
-Base = declarative_base()
-async_session = async_sessionmaker(engine, expire_on_commit=False)
+
+class Base(DeclarativeBase):
+    """Base class for SQLAlchemy models.
+
+    All other models should inherit from this class.
+    """
+
+    metadata = MetaData(
+        naming_convention={
+            "ix": "ix_%(column_0_label)s",
+            "uq": "uq_%(table_name)s_%(column_0_name)s",
+            "ck": "ck_%(table_name)s_%(constraint_name)s",
+            "fk": "fk_%(table_name)s_%(column_0_name)s_%(referred_table_name)s",
+            "pk": "pk_%(table_name)s",
+        }
+    )
+
+
+async_engine = create_async_engine(DATABASE_URL, echo=False)
+async_session = async_sessionmaker(async_engine, expire_on_commit=False)
 
 
-async def get_db():
+async def get_db() -> AsyncGenerator[AsyncSession, Any]:
     """Get a database session.
 
     To be used for dependency injection.
     """
-    async with async_session() as session:
-        async with session.begin():
-            yield session
+    async with async_session() as session, session.begin():
+        yield session
 
 
-async def init_models():
+async def init_models() -> None:
     """Create tables if they don't already exist.
 
     In a real-life example we would use Alembic to manage migrations.
     """
-    async with engine.begin() as conn:
-        # await conn.run_sync(Base.metadata.drop_all)
+    async with async_engine.begin() as conn:
+        # await conn.run_sync(Base.metadata.drop_all)  # noqa: ERA001
         await conn.run_sync(Base.metadata.create_all)

main.py

@@ -1,16 +1,18 @@
 """An example of using FastAPI with Async SQLAlchemy 2."""
+from collections.abc import AsyncGenerator, Sequence
 from contextlib import asynccontextmanager
+from typing import Any
 
 import uvicorn
-from fastapi import Depends, FastAPI
-from sqlalchemy import select
-
 from db import get_db, init_models
+from fastapi import Depends, FastAPI
 from models import User
+from sqlalchemy import select
+from sqlalchemy.ext.asyncio import AsyncSession
 
 
 @asynccontextmanager
-async def lifespan(app: FastAPI):
+async def lifespan(app: FastAPI) -> AsyncGenerator[Any, None]:  # noqa: ARG001
     """Run tasks before and after the server starts."""
     await init_models()
     yield
@@ -20,25 +22,26 @@ async def lifespan(app: FastAPI):
 
 
 @app.get("/")
-async def root():
+async def root() -> dict[str, str]:
     """Root endpoint."""
     return {"message": "Test API for FastAPI and Async SQLAlchemy ."}
 
 
 @app.post("/users/")
-async def create_user(name: str, email: str, session=Depends(get_db)):
+async def create_user(
+    name: str, email: str, session: AsyncSession = Depends(get_db)
+) -> User:
     """Add a user."""
     user = User(name=name, email=email)
     session.add(user)
     return user
 
 
 @app.get("/users/")
-async def get_users(session=Depends(get_db)):
+async def get_users(session: AsyncSession = Depends(get_db)) -> Sequence[User]:
     """Get all users."""
     result = await session.execute(select(User))
-    users = result.scalars().all()
-    return users
+    return result.scalars().all()
 
 
 if __name__ == "__main__":

models.py

@@ -1,7 +1,6 @@
 """Define Models used in this example."""
-from sqlalchemy import Column, Integer, String
-
 from db import Base
+from sqlalchemy import Column, Integer, String
 
 
 class User(Base):