FastAPI Database Integration with SQLAlchemy

from sqlalchemy.ext.asyncio import create_async_engine, async_sessionmaker, AsyncSession
from sqlalchemy.orm import DeclarativeBase
import os

# Production-grade PostgreSQL async connection string
DATABASE_URL = os.getenv("DATABASE_URL", "postgresql+asyncpg://user:pass@localhost/forge_db")

# create_async_engine prevents the event loop from blocking during I/O
engine = create_async_engine(DATABASE_URL, echo=False, future=True)

# Factory for creating session objects
AsyncSessionLocal = async_sessionmaker(
    bind=engine, 
    class_=AsyncSession, 
    expire_on_commit=False
)

class Base(DeclarativeBase):
    """Base class for all SQLAlchemy models"""
    pass

# Dependency: The 'Session-per-Request' provider
async def get_db():
    async with AsyncSessionLocal() as session:
        try:
            yield session
            # Session is automatically committed/closed if using 'async with' context
        except Exception:
            await session.rollback()
            raise
        finally:
            await session.close()

from sqlalchemy import String, Integer, select
from sqlalchemy.orm import Mapped, mapped_column
from fastapi import FastAPI, Depends, HTTPException, status
from io.thecodeforge.db.session import Base, get_db, AsyncSession
from pydantic import BaseModel, EmailStr

# 1. SQLAlchemy Model (Database Identity)
class User(Base):
    __tablename__ = 'users'
    id: Mapped[int] = mapped_column(primary_key=True, index=True)
    username: Mapped[str] = mapped_column(String(50), unique=True)
    email: Mapped[str] = mapped_column(unique=True)

# 2. Pydantic Schema (API Contract)
class UserCreate(BaseModel):
    username: str
    email: EmailStr

app = FastAPI()

@app.post('/forge/users', status_code=status.HTTP_201_CREATED)
async def create_user(payload: UserCreate, db: AsyncSession = Depends(get_db)):
    new_user = User(username=payload.username, email=payload.email)
    db.add(new_user)
    await db.commit()
    await db.refresh(new_user) # Reloads auto-generated ID from DB
    return new_user

@app.get('/forge/users/{user_id}')
async def read_user(user_id: int, db: AsyncSession = Depends(get_db)):
    result = await db.execute(select(User).where(User.id == user_id))
    user = result.scalar_one_or_none()
    if not user:
        raise HTTPException(status_code=404, detail="User not found")
    return user

from io.thecodeforge.db.session import AsyncSession
from io.thecodeforge.db.models import User
from sqlalchemy import select

class UserRepository:
    def __init__(self, db: AsyncSession):
        self.db = db

    async def get_by_id(self, user_id: int) -> User | None:
        result = await self.db.execute(select(User).where(User.id == user_id))
        return result.scalar_one_or_none()

    async def create(self, username: str, email: str) -> User:
        user = User(username=username, email=email)
        self.db.add(user)
        await self.db.commit()
        await self.db.refresh(user)
        return user

# In endpoint:
@app.get('/forge/users/{user_id}')
async def read_user(user_id: int, db: AsyncSession = Depends(get_db)):
    repo = UserRepository(db)
    user = await repo.get_by_id(user_id)
    if not user:
        raise HTTPException(404)
    return user

from logging.config import fileConfig
from sqlalchemy import pool
from sqlalchemy.ext.asyncio import create_async_engine
from alembic import context
from io.thecodeforge.db.session import Base
from io.thecodeforge.db.models import User  # noqa: F401

config = context.config
fileConfig(config.config_file_name)

target_metadata = Base.metadata

def run_migrations_offline():
    url = config.get_main_option("sqlalchemy.url")
    context.configure(url=url, target_metadata=target_metadata, literal_binds=True)
    with context.begin_transaction():
        context.run_migrations()

def do_run_migrations(connection):
    context.configure(connection=connection, target_metadata=target_metadata)
    with context.begin_transaction():
        context.run_migrations()

async def run_async_migrations():
    connectable = create_async_engine(config.get_main_option("sqlalchemy.url"), poolclass=pool.NullPool)
    async with connectable.connect() as connection:
        await connection.run_sync(do_run_migrations)
    await connectable.dispose()

def run_migrations_online():
    import asyncio
    asyncio.run(run_async_migrations())

context.run_migrations = run_migrations_online if context.is_offline_mode() else run_migrations_offline

import pytest
from httpx import AsyncClient, ASGITransport
from sqlalchemy.ext.asyncio import create_async_engine, async_sessionmaker, AsyncSession
from io.thecodeforge.db.session import Base, get_db
from io.thecodeforge.main import app

# Use a separate test database
TEST_DATABASE_URL = "postgresql+asyncpg://user:pass@localhost/test_forge_db"

@pytest.fixture
async def async_engine():
    engine = create_async_engine(TEST_DATABASE_URL, echo=False)
    async with engine.begin() as conn:
        await conn.run_sync(Base.metadata.create_all)
    yield engine
    await engine.dispose()

@pytest.fixture
async def db_session(async_engine) -> AsyncSession:
    session = async_sessionmaker(bind=async_engine, class_=AsyncSession, expire_on_commit=False)()
    try:
        yield session
    finally:
        await session.close()

@pytest.fixture
async def client(db_session):
    async def override_get_db():
        yield db_session
    app.dependency_overrides[get_db] = override_get_db
    async with AsyncClient(transport=ASGITransport(app=app), base_url="http://testserver") as ac:
        yield ac

@pytest.mark.asyncio
async def test_create_user(client: AsyncClient):
    response = await client.post("/forge/users", json={"username": "tester", "email": "tester@test.com"})
    assert response.status_code == 201
    data = response.json()
    assert data["username"] == "tester"

async def create_order(db: AsyncSession, order_data: OrderCreate) -> Order:
    async with db.begin():
        # Automatically commits or rolls back on exception
        user = await db.execute(select(User).where(User.id == order_data.user_id))
        user = user.scalar_one_or_none()
        if not user:
            raise HTTPException(status_code=404, detail="User not found")
        order = Order(user_id=user.id, total=order_data.total)
        db.add(order)
        # Use nested transaction for partial rollback
        async with db.begin_nested():
            for item in order_data.items:
                inventory_item = await db.execute(select(Inventory).where(Inventory.product_id == item.product_id))
                inventory_item = inventory_item.scalar_one_or_none()
                if not inventory_item or inventory_item.quantity < item.quantity:
                    raise HTTPException(status_code=400, detail="Insufficient stock")
                inventory_item.quantity -= item.quantity
    await db.refresh(order)
    return order