[sqlalchemy] @contextmanger로 트랜젝션 관리

FastAPI + Sqlalchemy를 사용하여 개발하던 중 insert 실패 후 다른 데이터를 insert를 할때 아래 와 같은 애러가 발생하였다. 

 

sqlalchemy.exc.PendingRollbackError: This Session's transaction has been rolled back due to a previous exception during flush. To begin a new transaction with this Session, first issue Session.rollback(). Original exception was: (psycopg2.errors.UniqueViolation) duplicate key value violates unique constraint

 

 

위의 애러를 살펴보면 롤백이 필요하다고 해서 아래와 같이 코드를 변경하였더니 정상 동작이 되었다. 외래키 떄문에 insert가 실패가 났는데 하나의 세션을 계속 사용해서 그런 것 같다. DI를 위해 아래의 Repository에 세션을 주입하는 형식으로 진행했다. 

class WorkspaceRepository:
    def __init__(self, session: Session):
        self.session = session

class FileRepository:
    def __init__(self, session: Session):
        self.session = session

 

두 레파지토리에 Dependency-Injector 라이브러리를 사용해서 같은 세션을 주입하다보니 발생하는 애러로 생각된다.

    session = providers.Singleton(SessionLocal)

    # repository
    file_repository = providers.Factory(FileRepository, session=session)
    workspace_repository = providers.Factory(WorkspaceRepository, session=session)

 

실제 세션의 메모리를 살펴보니 같은 세션을 사용하고 있는 것을 확인할 수 있었다. 따라서 첫번째 쿼리에서 실패가 되면 rollback을 진행하고 다른 세션을 사용할 수 있도록 변경이 필요해 보였다.

<sqlalchemy.orm.session.Session object at 0x7fecf0824210>
<sqlalchemy.orm.session.Session object at 0x7fecf0824210>

 

 

이를 해결하기 위해 세션 주입 부분은 생략하고 우선 contextmanger를 사용해서 자원 관리를 진행해보도록 하였다.

@contextmanager
def session_manager():
    """
    example : 
      with session_manager() as session:
        result = session.query(User).all()
        print(result)
    """
    
    session = SessionLocal()
    try:
        yield session
    except Exception as e:
        print("rollback transaction")
        session.rollback()
        raise
    finally:
        print("closing connection")
        session.close()

 

with session_manager() as session:
	session.add(document)

 

이렇게 다른 세션을 사용함을 확인 하였고 롤백도 정상적으로 진행되는 것을 확인 할 수 있었다.

<sqlalchemy.orm.session.Session object at 0x7fa5b1c8f8d0>
<sqlalchemy.orm.session.Session object at 0x7fa5afc33150>