o2o-site-AEO/solution/backend/common/database/db_session_manager.py

from asyncio import current_task

from sqlalchemy.orm import sessionmaker
from sqlalchemy.exc import IntegrityError
from sqlalchemy.ext.asyncio import AsyncSession, create_async_engine, async_scoped_session
from sqlalchemy.util._collections import immutabledict

from common.database.model.models import MAIN_BASE
from common.enums import DBType, DBWRType, ErrorType
from common.logger import LOG
from common.singleton import Singleton
from config.server_configs import main_db_config


class DBSessionManager(Singleton):
    """DB 세션/엔진 관리자 (싱글톤).

    핵심 패턴
    - DBType(논리 DB) x DBWRType(Read/Write) 조합마다 별도 async 엔진을 둔다.
      => 조회는 Read 복제본, 변경은 Write 주 DB 로 자연스럽게 분리된다.
    - 비즈니스 로직(service)은 직접 세션을 열지 않고 "람다"를 넘긴다.
      execute_lambda        : 단일 쿼리 (주로 조회)
      execute_lambda_run    : 동일 DB 의 여러 변경 쿼리를 한 트랜잭션으로 commit
      세션 open/close 와 commit/rollback 은 매니저가 책임진다.
    """

    def __init__(self):
        if DBSessionManager.is_init():
            LOG.e_no_callstack("already init DBSessionManager")
            return
        DBSessionManager.set_init()

        self.__DB_URL_MAP = {"postgresql": "postgresql+asyncpg"}
        # 종료 시 dispose 하기 위해 생성한 엔진을 모아둔다.
        self.__engines = []
        # 논리 DB -> config. DB 가 늘어나면 여기에 추가만 하면 된다.
        self.__db_type_map = {
            DBType.MAIN.value: main_db_config,
        }

        # Write 엔진 맵
        self.__write_session = {
            DBType.MAIN.value: self.create_engine(DBType.MAIN.value, DBWRType.DB_WRITE.value),
        }
        # Read 엔진 맵
        self.__read_session = {
            DBType.MAIN.value: self.create_engine(DBType.MAIN.value, DBWRType.DB_READ.value),
        }

    def create_engine(self, db_type: int, db_wr_type: int):
        db_config = self.__db_type_map.get(db_type)
        if not db_config:
            raise ValueError("Invalid database type")

        if db_wr_type == DBWRType.DB_READ.value:
            pw = (":" + db_config.read_pw) if len(db_config.read_pw) > 0 else ""
            db_url = f"{self.__DB_URL_MAP[db_config.db_type]}://{db_config.read_id}{pw}@{db_config.read_host}:{db_config.read_port}/{db_config.name}"
            LOG.i(f"Read DB create engine url : {db_url}")
        else:
            pw = (":" + db_config.write_pw) if len(db_config.write_pw) > 0 else ""
            db_url = f"{self.__DB_URL_MAP[db_config.db_type]}://{db_config.write_id}{pw}@{db_config.write_host}:{db_config.write_port}/{db_config.name}"
            LOG.i(f"Write DB create engine url : {db_url}")

        # SSL/TLS: 관리형 DB(RDS/Aurora/Azure)는 보통 TLS 필수. sslmode 가 설정되면 asyncpg 에 전달.
        connect_args = {}
        sslmode = (getattr(db_config, "sslmode", "") or "").lower()
        if sslmode and sslmode != "disable":
            connect_args["ssl"] = sslmode

        engine = create_async_engine(
            db_url,
            echo=db_config.show_log,
            pool_size=db_config.pool_size,
            max_overflow=db_config.max_overflow,
            pool_pre_ping=True,
            pool_recycle=600,
            connect_args=connect_args,
        )
        self.__engines.append(engine)
        scoped_session = async_scoped_session(
            sessionmaker(engine, class_=AsyncSession, expire_on_commit=False, autocommit=False, autoflush=False),
            scopefunc=current_task,
        )
        return scoped_session

    def is_registered(self, db_type: int) -> bool:
        """해당 논리 DB 가 등록돼 있는지 — 선택 연동(부속 DB)의 배치가 실행 전 확인하는 용도."""
        return db_type in self.__db_type_map

    async def dispose_all(self):
        """모든 엔진의 커넥션 풀을 정리한다. 앱 종료/테스트 종료 시 호출한다.
        호출하지 않으면 풀 커넥션이 이벤트 루프 종료 후 GC 되며 경고를 남긴다.
        """
        for engine in self.__engines:
            await engine.dispose()

    # ---- 세션 lifecycle -------------------------------------------------
    async def start_session(self, db_type: int, db_wr_type: int) -> AsyncSession:
        if db_wr_type == DBWRType.DB_WRITE.value:
            return self.__write_session[db_type]()
        return self.__read_session[db_type]()

    async def end_session(self, db_type: int, db_wr_type: int):
        if db_wr_type == DBWRType.DB_WRITE.value:
            await self.__write_session[db_type].remove()
        else:
            await self.__read_session[db_type].remove()

    # ---- 저수준 DB 연산 (crud 에서 호출) --------------------------------
    async def run(self, db: AsyncSession, err_msg="DB Run Failed", raise_error=True) -> ErrorType:
        try:
            await db.commit()
            return ErrorType.SUCCESS
        except IntegrityError as ex:
            await db.rollback()
            LOG.e_no_callstack(f"duplicated. {ex}")
            return ErrorType.DB_ALREADY_SAME_KEY
        except Exception as ex:
            await db.rollback()
            err_type = ErrorType.DB_RUN_FAILED
            LOG.e_no_callstack(f"[{err_type.name}] {err_msg=}, {ex=}")
            if raise_error:
                raise RuntimeError(err_type.name, err_msg)
            return err_type

    async def insert(self, db: AsyncSession, obj, err_msg="DB Failed", raise_error=True) -> ErrorType:
        try:
            if isinstance(obj, MAIN_BASE):
                db.add(obj)
            elif isinstance(obj, list):
                db.add_all(obj)
            else:
                raise RuntimeError("DO NOT USE QUERY IN DBJOB")
            return ErrorType.SUCCESS
        except Exception as ex:
            await db.rollback()
            err_type = ErrorType.DB_RUN_FAILED
            LOG.e_no_callstack(f"[{err_type.name}] {err_msg=}, {ex=}")
            if raise_error:
                raise RuntimeError(err_type.name, err_msg)
            return err_type

    async def add(self, db: AsyncSession, query, err_msg="DB Operation Failed", raise_error=True) -> ErrorType:
        """update/delete 등 비-select 쿼리 실행."""
        try:
            if hasattr(query, "column_descriptions"):
                raise RuntimeError("DO NOT USE SELECT QUERY IN DBJOB")
            await db.execute(query, execution_options=immutabledict({"synchronize_session": "fetch"}))
            return ErrorType.SUCCESS
        except IntegrityError as ex:
            await db.rollback()
            err_type = ErrorType.DB_ALREADY_SAME_KEY
            LOG.e_no_callstack(f"[{err_type.name}] {err_msg=}, {ex=}")
            return err_type
        except Exception as ex:
            await db.rollback()
            err_type = ErrorType.DB_RUN_FAILED
            LOG.e_no_callstack(f"[{err_type.name}] {err_msg=}, {ex=}")
            if raise_error:
                raise RuntimeError(err_type.name, err_msg)
            return err_type

    async def add_with_rowcount(self, db: AsyncSession, query, err_msg="DB Operation Failed") -> tuple[ErrorType, int]:
        """update/delete 등 비-select 쿼리 실행 후 (ErrorType, 영향행수) 반환.
        조건부 갱신(WHERE 로 상태를 거른 UPDATE)이 실제로 적용됐는지 판별하는 동시처리 가드용."""
        try:
            if hasattr(query, "column_descriptions"):
                raise RuntimeError("DO NOT USE SELECT QUERY IN DBJOB")
            res = await db.execute(query, execution_options=immutabledict({"synchronize_session": "fetch"}))
            return ErrorType.SUCCESS, res.rowcount
        except IntegrityError as ex:
            await db.rollback()
            err_type = ErrorType.DB_ALREADY_SAME_KEY
            LOG.e_no_callstack(f"[{err_type.name}] {err_msg=}, {ex=}")
            return err_type, 0
        except Exception as ex:
            await db.rollback()
            err_type = ErrorType.DB_RUN_FAILED
            LOG.e_no_callstack(f"[{err_type.name}] {err_msg=}, {ex=}")
            return err_type, 0

    async def execute(self, db: AsyncSession, query, err_msg="DB Query Execution Failed", raise_error=True) -> tuple[ErrorType, list]:
        """select 쿼리 실행 후 결과 리스트 반환."""
        try:
            if not hasattr(query, "column_descriptions"):
                raise RuntimeError("DO NOT USE NON-SELECT QUERY IN DBJOB")
            res = await db.execute(query, execution_options=immutabledict({"synchronize_session": "fetch"}))
            return ErrorType.SUCCESS, res.scalars().fetchall() if 1 == len(query.column_descriptions) else res.all()
        except Exception as ex:
            err_type = ErrorType.DB_RUN_FAILED
            LOG.e_no_callstack(f"[{err_type.name}] {err_msg=}, {ex=}")
            if raise_error:
                raise RuntimeError(err_type.name, err_msg)
            return err_type, []

    # ---- 람다 실행 진입점 (service 에서 호출) ---------------------------
    async def execute_lambda(self, db_type: int, db_wr_type: int, func):
        """단일 쿼리 호출. func(session) 한 개를 실행하고 결과를 그대로 반환."""
        s = await self.start_session(db_type, db_wr_type)
        try:
            return await func(s)
        finally:
            await self.end_session(db_type, db_wr_type)

    async def execute_lambda_run(self, db_type_list: list[int], func_list: list):
        """동일 DB 의 변경 쿼리 여러 개를 한 트랜잭션으로 실행 후 commit.
        하나라도 SUCCESS 가 아니면 즉시 중단(rollback)된다.
        """
        temp_list = list(set(db_type_list))
        if len(temp_list) != 1:
            return ErrorType.DB_INVALID_TYPE

        db_type = temp_list[0]
        s = await self.start_session(db_type, DBWRType.DB_WRITE.value)
        try:
            for func in func_list:
                err_type = await func(s)
                if err_type != ErrorType.SUCCESS:
                    return err_type
            return await self.run(s)
        except Exception as ex:
            LOG.e_no_callstack(ex)
            return ErrorType.DB_RUN_FAILED
        finally:
            await self.end_session(db_type, DBWRType.DB_WRITE.value)

    async def execute_lambda_write(self, db_type: int, func):
        """Write 세션에서 func(session) 을 실행하고 commit 한 뒤 **func 의 반환값을 그대로** 돌려준다.

        execute_lambda_run 은 ErrorType 만, execute_lambda_claim 은 (ErrorType, 적용행수) 만 돌려준다.
        작업 큐처럼 "변경하면서 값을 받아와야" 하는 경우(RETURNING 절)를 위한 진입점이다 —
        원자적 claim(FOR UPDATE SKIP LOCKED + UPDATE + RETURNING)은 조회/변경을 나눌 수 없다.
        예외는 rollback 후 그대로 전파한다(호출측이 잡 실패로 처리)."""
        s = await self.start_session(db_type, DBWRType.DB_WRITE.value)
        try:
            result = await func(s)
            await s.commit()
            return result
        except Exception:
            await s.rollback()
            raise
        finally:
            await self.end_session(db_type, DBWRType.DB_WRITE.value)

    async def execute_lambda_claim(self, db_type: int, func) -> tuple[ErrorType, int]:
        """조건부 변경 쿼리 1건을 한 트랜잭션으로 실행/commit 하고 (ErrorType, 적용행수) 반환.
        동시처리 가드용 — func(session) -> (ErrorType, rowcount). 적용행수 0 이면 다른 호출자가 이미 처리한 것.
        (Postgres READ COMMITTED 에서 같은 행 UPDATE 는 행 잠금으로 직렬화되어, 진 호출자는 0 을 받는다.)"""
        s = await self.start_session(db_type, DBWRType.DB_WRITE.value)
        try:
            err_type, rowcount = await func(s)
            if err_type != ErrorType.SUCCESS:
                return err_type, 0
            commit_err = await self.run(s)
            return commit_err, rowcount
        except Exception as ex:
            LOG.e_no_callstack(ex)
            return ErrorType.DB_RUN_FAILED, 0
        finally:
            await self.end_session(db_type, DBWRType.DB_WRITE.value)


DB_SESSION_MNG = DBSessionManager()