SQLAlchemy에서 automap_base 기능을 이용하여 데이터베이스 테이블을 간편히 불러오기

 

 

sqlalchemy를 이용하다보면, 데이터베이스의 테이블 구조를 python의 class 객체 형태로 매핑해준 후, 데이터에 접근하여 각종 작업을 수행해야 하는 경우가 많습니다. 예를 들면, 'User'라는 이름의 테이블을 불러오는 경우 다음과 같이 정의하곤 합니다.

 

class User(Base): 
    __tablename__ = 'users'
    id = Column(Integer, primary_key=True) 
    name = Column(String) 
    fullname = Column(String) 
    password = Column(String)

 

 

하지만, 여러 테이블을 다루다보면 이러한 작업이 다소 복잡하게 느껴집니다. sqlalchemy에서는 이러한 작업 없이도 테이블을 간편히 python class로 매핑할 수 있도록 automap_base 라는 기능을 제공합니다. 본 포스트에서는 이 automap_base 기능을 사용하여 간단히 데이터베이스의 데이터에 접근하는 예제를 소개해보겠습니다.

 

먼저, db.py 파일을 다음과 같이 구성합니다. 일반적으로는 get_db 함수 정도만 정의하여 사용하지만, 이는 트랜잭션에만 이용하며 테이블 정보를 불러오는 get_table 함수를 별도로 구성해봅시다. 

 

 

# db.py
from fastapi import HTTPException
from sqlalchemy.ext.automap import automap_base
from sqlalchemy.orm import sessionmaker
from sqlalchemy import create_engine
from dotenv import load_dotenv
import os

load_dotenv()

def get_db():
    DATABASE_URL = os.getenv(f"DATABASE_URL")
    engine = create_engine(DATABASE_URL)
    SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)
    db = SessionLocal()
    try:
        yield db
    finally:
        db.close()

def get_table(table_name: str):
    DATABASE_URL = os.getenv(f"DATABASE_URL")
    engine = create_engine(DATABASE_URL)
    Base = automap_base()
    Base.prepare(engine, reflect=True)
    if hasattr(Base.classes, table_name):
        return getattr(Base.classes, table_name)
    else:
        raise HTTPException(status_code=404, detail=f"'{table_name}' 테이블 정보를 찾을 수 없습니다.")

 

 

automap_base() 기능은 데이베이스 엔진 객체와 결합하여 테이블을 python class 형태로 불러올 수 있도록 합니다. 데이터베이스 정보는 dotenv 파일(.env)에 저장해놓았다고 가정합니다. 이제, 간단한 fastapi 앱을 띄워 테스트를 해볼 것입니다. db.py와 같은 경로에 main.py 파일을 다음과 같이 구성합니다. main.py의 '/api/is-user-exist' api는 데이터베이스에 입력한 사용자 이메일이 존재하는지 확인합니다. 참고로, 샘플로 입력한 이메일은 데이터베이스의 'tb_user' 테이블에 다음과 같이 실제로 존재하고 있습니다.

 

 

 

# main.py
from fastapi import FastAPI
from fastapi.responses import JSONResponse
from pydantic import BaseModel, Field
from db import get_db, get_table

app = FastAPI(title="automap_base() practice page")

@app.get("/")
def root():
    return {"message": "success"}

class SampleInputForm(BaseModel):
    email: str = Field(..., example="hsh@test.com")

@app.post("/api/is-user-exist/")
def is_user_exist(form: SampleInputForm):
    db = next(get_db())

    tb_user = get_table(table_name="tb_user")
    db_user = db.query(tb_user).filter(tb_user.email == form.email).first()

    if db_user:
        return JSONResponse({"detail": "사용자가 존재합니다."})
    else:
        return JSONResponse({"detail": "사용자가 존재하지 않습니다."})

if __name__ == "__main__":
    import uvicorn
    uvicorn.run('main:app', host="0.0.0.0", port=3000, reload=True)

 

 

 

get_db 함수로 불러온 db 객체는 generator 객체이기 때문에 next()로 감싸 불러와야 오류가 발생하지 않습니다. get_table 함수는 테이블 이름을 입력으로 받아, db_user 객체로 tb_user 테이블 객체를 매핑합니다. 입력 이메일의 존재 여부에 따라, 반환 메시지가 달라질 것입니다. 이제 host와 post를 적절히 설정한 후 main.py를 배포해봅시다. 

 

python main.py

 

 

배포된 주소에서, /docs를 추가로 입력하면 fastapi 문서로 진입합니다.

 

 

 

설정한 api 이름을 클릭하여 펼친 후, Try it out을 클릭하여 테스트해봅시다.

 

 

 

다음과 같이 성공적으로 데이터를 조회화여 메시지가 반환된 것을 볼 수 있습니다.

 

 

 

automap_base 기능은 간편하긴 하지만, 테이블 구조가 복잡하거나 테이블 변경 사항이 까다로울 경우에는 적합하지 않을 수 있습니다. 환경을 고려하여 automap_base 기능을 적절히 사용하면, 데이터베이스와의 상호작용 구조를 보다 직관적으로 개선할 수 있을 것입니다.