【Python】sqlmodel: Python 資料庫管理ORM 的終極形態？

ORM

大家都知道ORM（Object Relational Mapping）是一種將物件和關聯式資料庫中的表進行對映的技術，它可以讓開發者更加方便地運算元據庫，而不用直接使用SQL語句。

直接使用SQL語句運算元據庫，雖然可以讓開發者直接與資料庫打交道，但手動編寫SQL語句，容易出錯，而且靈活性上比較欠缺。相比之下，使用ORM（以SQLAlchemy為例）有更加易於使用、更加靈活、能防止 SQL 注入攻擊、更加易於測試的優勢。

點選檢視優勢說明

更加易於使用: 可以使用 Python 物件來表示資料庫中的表和行，而不是直接使用 SQL 語句。這樣可以使程式碼更加易於編寫和維護。
更加靈活: SQLAlchemy 提供了靈活的查詢語言，可以通過鏈式呼叫的方式構建複雜的查詢語句。同時，SQLAlchemy 支援多種資料庫，可以在不同的資料庫之間進行切換，而不需要修改程式碼。
防止 SQL 注入攻擊: SQLAlchemy 提供了引數化查詢的方式，可以有效地防止 SQL 注入攻擊。使用引數化查詢可以將使用者輸入的資料轉換為引數，從而避免了 SQL 注入攻擊。
更加易於測試: 使用 SQLAlchemy 可以將業務邏輯和資料庫操作分離，從而使得程式碼更加易於測試。可以通過 Mock 物件模擬資料庫操作，從而進行單元測試和整合測試。
...

當然，使用 SQLAlchemy 也會增加程式碼的複雜度，需要學習額外的知識和 API。因此，在實際應用中需要根據具體情況進行選擇。

那麼有沒有一種技術或者框架 既不用增加太多的應用成本，又兼具以SQLAlchemy為代表的ORM 框架的優勢 呢？答案是肯定的，那就是我們今天介紹的主角 sqlmodel.

我們就以 Fastapi 開發建立使用者和查詢使用者 兩個功能的介面來對比一下，SQLAlchemy 和 sqlmodel， sqlmodel 和只使用 SQL的差異。

使用SQLAlchemy

安裝

pip install sqlalchemy

範例程式碼

from fastapi import FastAPI, Depends, HTTPException
from sqlalchemy import create_engine, Column, Integer, String
from sqlalchemy.orm import Session, declarative_base, sessionmaker

SQLALCHEMY_DATABASE_URL = "mysql://user:password@host:port/database"

engine = create_engine(SQLALCHEMY_DATABASE_URL)
SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)
Base = declarative_base()

app = FastAPI()

class User(Base):
    __tablename__ = "users"

    id = Column(Integer, primary_key=True, index=True)
    name = Column(String(50))
    age = Column(Integer)

class UserIn(Base):
    name: str
    age: int

class UserOut(Base):
    id: int
    name: str
    age: int

class UserUpdate(Base):
    name: Optional[str] = None
    age: Optional[int] = None

Base.metadata.create_all(bind=engine)

def get_db():
    db = None
    try:
        db = SessionLocal()
        yield db
    finally:
        db.close()

def create_user(db: Session, user: UserIn):
    db_user = User(name=user.name, age=user.age)
    db.add(db_user)
    db.commit()
    db.refresh(db_user)
    return db_user

def read_user(db: Session, user_id: int):
    db_user = db.query(User).filter(User.id == user_id).first()
    if not db_user:
        raise HTTPException(status_code=404, detail="User not found")
    return db_user

def read_all_user(db: Session, ):
    db_user = db.query(User).all()
    if not db_user:
        raise HTTPException(status_code=404, detail="User not found")
    return db_user

@app.post("/users/", response_model=UserOut)
async def create_user_view(user: UserIn, db: Session = Depends(get_db)):
    return create_user(db, user)

@app.get("/users/{user_id}", response_model=UserOut)
async def read_user_view(user_id: int, db: Session = Depends(get_db)):
    return read_user(db, user_id)

@app.get("/users/", response_model=UserOut)
async def read_all_user_view(db: Session = Depends(get_db)):
    return read_all_user(db)

程式碼解釋

User 是資料模型類的名稱，id、name、age 是表中的列名。UserIn 是建立使用者的請求引數模型，UserOut 是查詢使用者的響應資料模型，UserUpdate 是更新使用者的請求引數模型。

使用 create_engine 函數建立一個資料庫連線引擎，使用 sessionmaker 函數建立一個資料庫對談工廠，使用 declarative_base 函數建立一個基礎類別。在建立表時，使用 Base.metadata.create_all 函數建立表。

使用 get_db 函數獲取資料庫對談物件，使用 create_user 和 read_user 函數進行資料庫操作。在檢視函數中，只需要呼叫這些函數即可完成相應的業務邏輯。

上面的程式碼已經非常簡潔直觀，但是還是有有一定的學習成本，下面我們來看下使用我們今天的主角 -- sqlmodel 需要怎樣來實現上面的介面。

使用sqlmodel

安裝 sqlmodel

pip install sqlmodel

範例程式碼

點選檢視完整程式碼

# -*- coding: utf-8 -*-
"""
@File   :dda.py
@Date   :2023-06-05
@user   :bingoHe
"""
from typing import Optional

from fastapi import FastAPI, Depends, HTTPException
from sqlalchemy import create_engine
from sqlalchemy.orm import Session
from sqlmodel import SQLModel, Field, create_all, Session as SQLModelSession

SQLALCHEMY_DATABASE_URL = "mysql://user:password@host:port/database"

engine = create_engine(SQLALCHEMY_DATABASE_URL)

app = FastAPI()

class UserBase(SQLModel):
    name: Optional[str] = None
    age: Optional[int] = None

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

class UserIn(UserBase):
    pass

class UserOut(UserBase):
    id: int

class UserUpdate(UserBase):
    pass


create_all(engine)

def get_db():
    """獲取資料庫對談物件"""
    db = None
    try:
        db = SQLModelSession(engine)
        yield db
    finally:
        db.close()

def create_user(db: SQLModelSession, user: UserIn):
    """建立使用者"""
    db_user = User.from_orm(user)
    db.add(db_user)
    db.commit()
    db.refresh(db_user)
    return db_user

def read_user(db: SQLModelSession, user_id: int):
    """查詢使用者"""
    db_user = db.get(User, user_id)
    if not db_user:
        raise HTTPException(status_code=404, detail="User not found")
    return db_user

@app.post("/users/", response_model=UserOut)
async def create_user(user: UserIn, db: SQLModelSession = Depends(get_db)):
    """建立使用者"""
    return create_user(db, user)

@app.get("/users/{user_id}", response_model=UserOut)
async def read_user(user_id: int, db: SQLModelSession = Depends(get_db)):
    """查詢使用者"""
    return read_user(db, user_id)

和SQLAlchemy的主要使用差異在引數的定義上，使用多處繼承，而不是各自定義的方法：

# Code above omitted