[DB] SQLModel로 데이터베이스 다루기 <1>

게시 2024/09/16

By Yu.S.H 6 분읽는 시간

공식문서 참조

1. SQLModel이란?

백엔드를 구성하다 보면 파이썬의 타입 검증 라이브러리 Pydantic과 파이썬 ORM sqlalchemy를 동시에 써야 할 경우가 굉장히 많다.

이로 인해 db에 쓰일 스키마와, 백엔드 코드에 쓸 스키마를 따로 구성해야 하기 때문에 코드의 양이 늘어나고, 실수확률이 올라가 불편하다.

이때, 이 두 가지 라이브러리를 한번에 합친 SQLModel을 사용한다면 문제를 해결 할 수 있다.

심지어 SQL Model 클래스 자체가 Pydantic의 Basemodel을 상속한다고 하니 정말 개꿀이다.

1.1 설치

pip install sqlmodel

한 줄이면 설치 가능하다. 편하다.

2. SQLModel로 DB와 연결

순서는 다음과 같다.

Engine 가동.
Engine 에서 Connection 연결.
Session 요청을 Connection으로 처리

세션과 커넥션의 개념은 여기 에 정리되어 있다.

2.1 Engine

Engine 은 데이터베이스와의 Connection을 관리하기 위한 객체이다. 엔진은 데이터베이스 서버와의 통신 채널같은 역할을 한다.

SQLModel 에선 자체적으로 db와의 엔진을 만드는 기능을 제공하지만, 비동기적 엔진을 만들기 위해선 SQLAlchemy의 비동기 기능을 사용해야 한다.

  
from sqlalchemy.ext.asyncio import create_async_engine

DB_URL = f"mysql+aiomysql://{Config.DB_USER}:{Config.DB_PASS}@{Config.DB_HOST}:{Config.DB_PORT}/{Config.DB_NAME}"

async_engine = create_async_engine(url=DB_URL, echo = False)

일반적으로 엔진은 애플리케이션 생성 시에 한 번 시작되며, 애플리케이션 가동 중에는 계속 유지된다. 엔진을 종료할때는 다음과 같다.

  
async_engine.dispose()

2.2 Connection

비동기적 엔진을 만들었다면 엔진을 통해 connection을 관리할 수 있게 된다.

엔진에서는 두 가지 방법으로 커넥션을 생성할 수 있다.

.connect() : SQL 쿼리를 시작하기 위한 기본적 커넥션 제공, 자동 트랜잭션X, 수동 커밋, 롤백
.begin() : SQL 쿼리를 시작하기 위한 커넥션 제공, 자동 트랜잭션 O, 정상적으로 트랜잭션이 종료되면 자동 커밋, 오류가 발생한다면 자동 롤백.

  
async def create_table(async_engine: AsyncEngine):
    async with async_engine.begin() as conn:
        await conn.run_sync(SQLModel.metadata.create_all)
    

비동기적 엔진을 사용했으므로 async with 를 달아서 사용한다.

2.3 Session

Session을 요청하게 된다면, Engine은 Connection을 Session에 할당해준다. 이 때 Session을 관리하는 주체는 Engine이 된다. 역시 SQLAlchemy의 비동기적 기능인 async_sessionmaker를 사용하게 된다.

  
from core.database.engine import async_engine
from sqlalchemy.ext.asyncio import async_sessionmaker
from sqlmodel.ext.asyncio.session import AsyncSession

async def get_async_session():
    async_session = async_sessionmaker(
        bind = async_engine,
        class_=AsyncSession,
    )

    async with async_session() as session:
        yield session

bind 에 데이터베이스 엔진을 지정해준다.
async with 와 yield를 통해 session을 열고 닫는 것을 자동으로 실행해준다. 이것으로 세션의 생명주기를 안전하게 관리 가능하다.

3. DB에 Table 생성

SQLModel에선 SQLModel 클래스를 상속한 클래스를 만드는 것으로 테이블을 손쉽게 만들 수 있다.

  
from sqlmodel import Field, SQLModel
from sqlalchemy.ext.asyncio import create_async_engine, AsyncEngine
import asyncio
from config import Config

DB_URL = f"mysql+aiomysql://{Config.DB_USER}:{Config.DB_PASS}@{Config.DB_HOST}:{Config.DB_PORT}/{Config.DB_NAME}"

engine = create_async_engine(url=DB_URL, echo=True)

class Hero(SQLModel, table=True):
    id: int | None = Field(default=None, primary_key=True)
    name: str
    secret_name: str
    age: int | None = None

async def create_db_and_tables():
    async with engine.begin() as conn:
        await conn.run_sync(SQLModel.metadata.create_all)

asyncio.run(create_db_and_tables())

이때 id 는 자동증가를 통한 값이 들어갈 것이므로, 객체 생성시 입력하지 않아도 되게 하기 위해 Optional으로 지정해준다.

Field + 타입 지정을 통해 기존 쿼리문에서 하던 각종 설정들을 해줄 수 있다.

미리 생성해둔 db에 테이블이 정상적으로 들어간 것을 볼 수 있다.

3.1 FastAPI에서 예시

FastAPI를 사용한다면 lifespan 관리를 통해 애플리케이션 시작 시에 테이블이 먼저 생성되도록 해줄 수 있다.

테이블 생성 함수

  
from config import Config
from sqlalchemy.ext.asyncio import AsyncEngine, create_async_engine
from sqlmodel import SQLModel

DB_URL = f"mysql+aiomysql://{Config.DB_USER}:{Config.DB_PASS}@{Config.DB_HOST}:{Config.DB_PORT}/{Config.DB_NAME}"

async_engine = create_async_engine(url=DB_URL, echo = False)

async def create_table(async_engine: AsyncEngine):
    async with async_engine.begin() as conn:
        await conn.run_sync(SQLModel.metadata.create_all)

FastAPI Configs 에서 lifespan 관리

  
class FastAPIConfigs:
    def __new__(cls, app: FastAPI):
        ExceptionHandelers(app)

    @classmethod
    @asynccontextmanager
    async def lifespan(cls, app: FastAPI):
        await create_table(async_engine)
        await insert_init_data()
        yield
        await async_engine.dispose()

lifespan 종료 시에 엔진을 닫아주도록 yield 뒤에 위치시킨다.

Main.py에서 컨트롤

  
app = FastAPI(
    lifespan=FastAPIConfigs.lifespan
)

FastAPIConfigs(app)

이렇게 작성하면 애플리케이션 시작 시에 테이블이 생성된다.

python sqlmodel db mariadb