FastAPI + SQLModel Postgres DB연동

 

개요

FastAPI + SQLModel Postgres DB연동 검색해도 잘 안나와서 포스팅 합니다.

 

다음 공식 사이트는 sqlite로 되어 있습니다.

https://sqlmodel.tiangolo.com/

SQLModel

 

Postgres연동

postgres연동으로 교체한 후 실행

from typing import Optional

from sqlmodel import Field, SQLModel, create_engine, Session, select


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


hero_1 = Hero(name="Deadpond", secret_name="Dive Wilson")
hero_2 = Hero(name="Spider-Boy", secret_name="Pedro Parqueador")
hero_3 = Hero(name="Rusty-Man", secret_name="Tommy Sharp", age=48)

# DB 접속 정보
db_host = "your_db_host" # ex) ec2-123-123-123-123.ap-northeast-2.compute.amazonaws.com
user = "your_db_user"
password = "your_db_password"

engine = create_engine(f'postgresql://{user}:{password}@{db_host}/your_db_name')


SQLModel.metadata.create_all(engine)

with Session(engine) as session:
    session.add(hero_1)
    session.add(hero_2)
    session.add(hero_3)
    session.commit()

with Session(engine) as session:
    statement = select(Hero).where(Hero.name == "Spider-Boy")
    hero = session.exec(statement).first()
    print(hero)

 

접속정보 env처리

보안상 소스코드에 DB접속정보는 넣지 않습니다. 다음과 같이 .env처리 할 수 있습니다.

# DB 접속 정보
db_host = os.getenv('DB_HOST')
user = os.getenv('DB_USERNAME')
password = os.getenv('DB_PASSWORD')

 

 

결과

DB에 생성

 

 

콘솔 출력

 

 

DB Connection분리

 

session.py

import os

from sqlmodel import create_engine, Session, SQLModel

db_host = os.getenv('DB_HOST')
user = os.getenv('DB_USERNAME')
password = os.getenv('DB_PASSWORD')

engine = create_engine(f'postgresql://{user}:{password}@{db_host}')


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


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

 

test_db_integration.py

위 session.py에서 yield session으로 리턴 했기 때문에 next(get_session())을 써서 session 객체를 받아줍니다.

from sqlmodel import select

from db.session import get_session, create_db_and_tables
from models.user_model import User


create_db_and_tables()

hero_1 = User(name="Kyeongrok", secret_name="Dive Wilson", age=38)
hero_2 = User(name="Taehwan", secret_name="Pedro Parqueador")
hero_3 = User(name="Junbin", secret_name="Tommy Sharp")

session = next(get_session())  # Generator를 썼기 때문에 next()를 사용해야 합니다.
session.add(hero_1)
session.add(hero_2)
session.add(hero_3)
session.commit()

statement = select(User).where(User.name == "Deadpond")
hero = session.exec(statement).first()
print(hero)

 

Router적용

main.py에 너무 많은 내용이 들어있기 때문에 router를 이용해 endpoint를 분리 해줍니다.

main.py

from fastapi import FastAPI

from api.v1.api import api_router

app = FastAPI(
    title="Flow View API",
    openapi_url="/api/v1/openapi.json"
)

app.include_router(api_router)

 

api.v1/api.py

from fastapi import APIRouter

from api.v1.endpoints import (root, user)

api_router = APIRouter()
api_router.include_router(root.router, tags=["root"])
api_router.include_router(user.router, prefix="/users", tags=["users"])

 

 

api.v1.endpoints/root.py

from fastapi import APIRouter

router = APIRouter()


@router.get("/")
async def root():
    return {"message": "Hello World"}


@router.get("/hello/{name}")
async def say_hello(name: str):
    # return apple string
    return {"message": f"Hello {name}"}