SQLAlchemyのセッションについて

SQLAlchemyのSessionをよくわからないままなんとなく使っていたので整理してみようと思います。

公式ドキュメントでscoped_sessionについて説明している記事があるのでこちらの内容を参考にしました。 docs.sqlalchemy.org

sessionmaker

sessionはsqlalchemyの機能を使ってDBにアクセスするために必要なオブジェクトで、sessionmakerを使うことでセッションを生成するファクトリを取得できます。

from sqlalchemy import create_engine
from sqlalchemy.orm import sessionmaker, scoped_session

engine = create_engine(DATABASE_URL)
Session = sessionmaker(bind=engine)

セッションファクトリーはそのまま呼び出すことでグローバルなセッションを取得することができます。

# 実際に使用する場合
with Session() as db:
    # DBアクセスの処理

また、セッションファクトリはセッションとして振る舞うことができます。

Session.add(user)
Session.commit()

FastAPIなどでDIとして注入する場合は次のような関数を定義して使うことが多いと思います。

async def get_db():
    with Session() as db:
        try:
            yield db
            db.commit()
        except: Exception as e:
            logger.error(e)

注意

他の記事でも触れられていましたが、sessionオブジェクトはクローズdb.close()しても、現在のセッションが破棄されるだけで、オブジェクト自体は有効のようです。つまり、session.close()したあとでもsession.queryなどを呼び出せてしまいます(内部で再びセッションが張られる)。例えば、FastAPIで非同期タスクを実行する際に、Dependsで注入したセッションオブジェクトをタスク側に渡してしまうと、API側でsession.closeを呼び出してもセッションオブジェクトがタスク内で生き続けてしまう可能性があります。このため、タスク内部では別にセッションを取得する必要がありそうです。 zenn.dev

@route.post("/some_task)
async some_task(background_tasks: BackgroundTasks, db: Session = Depends(get_db)):
    background_tasks.add_task(task, db)
    return { "message":  "created task" }

scoped_session

scoped_sessionはスレッドセーフなセッションファクトリを返します。scoped_sessionにより取得したセッションファクトリは、同一スレッド内においてセッションが閉じるまで常に同一のセッションを返し(最初に生成したセッションをキャッシュする)、セッションが閉じた後は新しいセッションが生成されます。

# スコープ化されたセッションを生成するファクトリーオブジェクトを返す。
SessionLocal = scoped_session(sessionmaker(engine=engine))
some_session = SessionLocal()
some_other_session = SessionLocal()

some_session is some_other_session # True

# セッションを破棄する
some_session.remove()

# 新しいセッションが生成される
new_session = Session()
new_session is some_session # False

scoped_sessionは内部でthreading.local(スレッドごとに分離したインメモリなストレージオブジェクトを返す)を使っていて、commitやrollbackなどで明示的にセッションが閉じなかった場合は各スレッドが終わるタイミングでGCによりセッションが解放されるので、この点についても、安全にマルチスレッド環境でセッションを利用できるようです。

まとめ

簡単にsessionmakerとscoped_sessionについて整理してみました。FastAPIのような非同期処理をサポートするWebAPI等でSQLAlchemyを利用する場合は、とりあえずscoped_sessionでセッションを生成しておけば問題なさそうです。