FastAPI 連接 SQL 資料庫

FastAPI 本身不限制你使用任何特定的資料庫。你可以使用 PostgreSQL, MySQL, SQLite, Oracle 等任何關聯式資料庫。

在 Python 生態系中,最流行的 ORM (Object-Relational Mapping) 工具是 SQLAlchemy。 而 SQLModel 是由 FastAPI 作者開發的一個基於 SQLAlchemy 和 Pydantic 的函式庫,它讓 FastAPI 與資料庫的整合變得異常簡單。

本章將示範如何使用 SQLModel 來連接 SQLite 資料庫。

安裝 SQLModel

pip install sqlmodel

定義模型 (Model)

SQLModel 的強大之處在於,同一個類別可以同時是 Pydantic Model (用於驗證 API 資料) 和 SQLAlchemy Model (用於定義資料表)。

from typing import Optional
from sqlmodel import Field, SQLModel

class Hero(SQLModel, table=True):
    id: int | None = Field(default=None, primary_key=True)
    name: str
    secret_name: str
    age: int | None = None
  • table=True:告訴 SQLModel 這是一個資料表。
  • Field(primary_key=True):設定主鍵。

建立資料庫連線

from sqlmodel import create_engine, SQLModel

sqlite_file_name = "database.db"
sqlite_url = f"sqlite:///{sqlite_file_name}"

# `connect_args={"check_same_thread": False}` 是 SQLite 專用的設定
engine = create_engine(sqlite_url, echo=True, connect_args={"check_same_thread": False})

def create_db_and_tables():
    SQLModel.metadata.create_all(engine)

依賴注入:取得 Session

我們需要一個 Dependency 來為每個 Request 提供一個資料庫 Session。

from sqlmodel import Session
from fastapi import Depends, FastAPI

def get_session():
    with Session(engine) as session:
        yield session

實作 CRUD API

現在我們可以開始寫 API 了。

app = FastAPI()

@app.on_event("startup")
def on_startup():
    create_db_and_tables()

from typing import Annotated

@app.post("/heroes/", response_model=Hero)
async def create_hero(hero: Hero, session: Annotated[Session, Depends(get_session)]):
    session.add(hero)
    session.commit()
    session.refresh(hero)
    return hero

@app.get("/heroes/", response_model=list[Hero])
async def read_heroes(session: Annotated[Session, Depends(get_session)]):
    heroes = session.exec(select(Hero)).all()
    return heroes

分離 Pydantic 與 Table 模型

在真實專案中,與 User 互動的模型 (Request Body) 和儲存到資料庫的模型 (Table) 往往不同 (例如:不想讓使用者直接設定 id)。

我們可以利用繼承來解決:

class HeroBase(SQLModel):
    name: str
    secret_name: str
    age: int | None = None

class Hero(HeroBase, table=True):
    id: Optional[int] = Field(default=None, primary_key=True)

class HeroCreate(HeroBase):
    pass

class HeroRead(HeroBase):
    id: int

@app.post("/heroes/", response_model=HeroRead)
async def create_hero(hero: HeroCreate, session: Annotated[Session, Depends(get_session)]):
    # Pydantic v2 / SQLModel 的現代寫法
    db_hero = Hero.model_validate(hero)
    session.add(db_hero)
    session.commit()
    session.refresh(db_hero)
    return db_hero

總結

  1. SQLModel 結合了 SQLAlchemy 和 Pydantic 的優點,非常適合 FastAPI。
  2. 使用 depends(get_session) 來管理資料庫連線的生命週期。
  3. 利用繼承來區分 API 模型 (Create/Read) 與 資料庫模型 (Table),保持程式碼整潔與安全。