2024-08-27 14:46

여피 프로젝트 FastAPI, SQLAlchemy 쿼리 최적화 - 1. 중복 쿼리 없애기

특정 케이스에 대해 유저의 상태를 계산하는 API가 몇 초 단위로 걸려 이를 최적화 하려 한다. 먼저 중복 쿼리가 어떤게 있는지 살펴보자.

( /v1/api/qual/870, 1주택 비과세 및 조특 50% 감면 임대주택 / 2862ms)

official_price 조회 쿼리 - 17번
complex 조회 쿼리 - 12번
property 조회 쿼리 - 4번
조정지역 조회 쿼리 - 3번
공동주택 유형 조회 쿼리 - 3번
user_relation 조회 쿼리 - 2번 (id로 조회 2번, 거주지 같은조건으로 조회 2번, 세대주 같은 조건으로 조회 1번)

( /v1/api/qual/8703, 일시적 2주택자 / 2771ms)

property_ownership_relation 조회 6번
조정지역 조회 쿼리 - 6번
complex 조회 쿼리 - 5번
user_relation 조회 쿼리 - 3번 (id로 조회 3번, 거주지 같은조건으로 조회 3번, 세대주 같은 조건으로 조회 3번, id 여러 개로 조회 3번)

중복 쿼리를 메모리에 캐싱하거나, 애플리케이션 코드 상에서 중복되게 호출하는 부분을 줄여서 개선할 수 있을 것 같다.

그런데 로그를 보면 cached since … 이 보이는데 알아서 캐싱이 됐나?

https://docs.sqlalchemy.org/en/20/core/connections.html#sql-caching

SQL 이 생성되는 부분에서 캐싱이 발생하는 것이지 쿼리 결과를 캐싱해주지는 않는다.
그러므로 캐싱하는 것은 의미가 있다. identity map 을 통한 내부 캐싱도 가능하겠지만 쿼리를보면 PK가 아닌 다른 조건으로 쿼리를 요청하는 경우도 많아(특히 많이 조회하는 official_price) 완전한 최적화는 아니다.

위에서 많이 중복되는 쿼리를 prefetch 하고 재사용하도록 코드를 수정하면 좋을 것 같다.

중복 쿼리 캐싱 결과

2862ms → 1318ms (user_relation_id: 870)

기존 코드에서 다음과 같이 필요한 정보를 미리 DB에서 불러와 메모리에 올려놓고 사용하는 방식으로 최적화함.

houses_with_related_info = await property_ownership_service.get_houses_with_related_info(
    user_relation_id, calculation_date, ignore_property_ownership_ids
)

이제 중복 쿼리가 없을까?

여전히 중복쿼리가 발생함.

prefetch 하는 메서드 내에서도 동일한 데이터를 여러 번 가져오려 하기 때문

해결방법

서비스 메서드에서 PK를 받아 DB에서 조회하는 방식이 아니라, 엔티티를 파라미터로 받아 사용하는 방식으로 수정한다. (애플리케이션 코드 수정)
PK 로 조회할 경우 sqlalchemy identity map 을 사용하여 메모리에서 캐싱된 상태로 가져올 수 있도록 한다. (메모리 캐싱)

1번 방법은 간단하고 무조건 쿼리를 줄일 수 있는 방법이지만, 코드 수정해야할 양이 많고, 코드 가독성을 해치면서 성능을 올리는 방법이다.

2번 방법을 사용해보자.

sqlalchemy identity map 사용하기

identity map은 PK 와 엔티티를 세션 메모리에 저장해두고 해당 PK 로 get 메서드를 통해 조회 시 쿼리를 날리지 않고 세션 메모리에서 엔티티를 반환하는 캐싱 기능이다.

기존 base_repository 코드를 쿼리를 날리는 것에서 get 메서드를 통해 PK로 조회해오는 것으로 수정하자.

async def get_one_by_id(self, id: int | str, type: Type | None = None) -> MODEL:
    model_type = type or self.model_type

    if type and isinstance(id, str):
        id = int(id[1:])

    model = await session.get(model_type, id)

    if not model:
        raise NotFoundException(f"Cannot find {self.model_type.__name__} with id: {id}")

    return model

기존에 4번 발생하던 property 조회 쿼리가 2번으로 줄어들었다.

근데 1번이 아니라 2번이다. 이유가 뭘까?

직접 identity map 의 상태를 살펴보기 위해 로그를 찍어보자.


    async def get_one_by_id(self, id: int | str, type: Type | None = None) -> MODEL:
        model_type = type or self.model_type

        if type and isinstance(id, str):
            id = int(id[1:])

        for key, obj in session.identity_map.items():
            logging.warning('one map' + str(key) + str(obj) + str(inspect(obj).key) + get_session_context())

        model = await session.get(model_type, id)

        if not model:
            raise NotFoundException(f"Cannot find {self.model_type.__name__} with id: {id}")

        for key, obj in session.identity_map.items():
            logging.warning('one map' + str(key) + str(obj) + str(inspect(obj).key) + " " + str(model == obj) + get_session_context())

        return model

쿼리가 나가는 부분

[root] [WARNING] one map<app.domains.property_ownership.models.property_ownership.PropertyOwnership object at 0x10dcaf0e0>(<class 'app.domains.property_ownership.models.property_ownership.PropertyOwnership'>, (717,), None)a4815dd2-6b73-4bdc-8671-725f70e58ac0
[sqlalchemy.engine.Engine] [INFO] SELECT property.id AS property_id, property.complex_id AS property_complex_id, property.type AS property_type, property.dong AS property_dong, property.ho AS property_ho, property.net_leasable_area AS property_net_leasable_area, property.apt_code AS property_apt_code, property.updated_at AS property_updated_at 
FROM property 
WHERE property.id = %s
[sqlalchemy.engine.Engine] [INFO] [generated in 0.00012s] (2128526,)
[sqlalchemy.engine.Engine] [DEBUG] Col ('property_id', 'property_complex_id', 'property_type', 'property_dong', 'property_ho', 'property_net_leasable_area', 'property_apt_code', 'property_updated_at')
[sqlalchemy.engine.Engine] [DEBUG] Row (2128526, 'A86322', 'A', '101', '204', 84.98, '20345416_0002_0007', datetime.datetime(2024, 3, 19, 10, 42, 43))
[root] [WARNING] one map<app.domains.property_ownership.models.property_ownership.PropertyOwnership object at 0x10dcaf0e0>(<class 'app.domains.property_ownership.models.property_ownership.PropertyOwnership'>, (717,), None) Falsea4815dd2-6b73-4bdc-8671-725f70e58ac0
[root] [WARNING] one map<app.domains.property.models.property.Property object at 0x10dc86930>(<class 'app.domains.property.models.property.Property'>, (2128526,), None) Truea4815dd2-6b73-4bdc-8671-725f70e58ac0

쿼리가 나가지 않고 identity map 에서 조회하는 부분

[root] [WARNING] one map<app.domains.property_ownership.models.property_ownership.PropertyOwnership object at 0x10dcaf0e0>(<class 'app.domains.property_ownership.models.property_ownership.PropertyOwnership'>, (717,), None)a4815dd2-6b73-4bdc-8671-725f70e58ac0
[root] [WARNING] one map<app.domains.complex.models.apartment.Apartment object at 0x10dcafe00>(<class 'app.domains.complex.models.apartment.Apartment'>, (86322,), None)a4815dd2-6b73-4bdc-8671-725f70e58ac0
[root] [WARNING] one map<app.domains.property.models.property.Property object at 0x10dcafef0>(<class 'app.domains.property.models.property.Property'>, (2128526,), None)a4815dd2-6b73-4bdc-8671-725f70e58ac0
[root] [WARNING] one map<app.domains.property_ownership.models.property_ownership.PropertyOwnership object at 0x10dcaf0e0>(<class 'app.domains.property_ownership.models.property_ownership.PropertyOwnership'>, (717,), None) Falsea4815dd2-6b73-4bdc-8671-725f70e58ac0
[root] [WARNING] one map<app.domains.complex.models.apartment.Apartment object at 0x10dcafe00>(<class 'app.domains.complex.models.apartment.Apartment'>, (86322,), None) Falsea4815dd2-6b73-4bdc-8671-725f70e58ac0
[root] [WARNING] one map<app.domains.property.models.property.Property object at 0x10dcafef0>(<class 'app.domains.property.models.property.Property'>, (2128526,), None) Truea4815dd2-6b73-4bdc-8671-725f70e58ac0

쿼리가 나가지 않는 부분에서는 정상적으로 identity map 에서 가져오지만, 쿼리가 나가는 부분에서는 identity map 자체에 엔티티가 저장되어 있지 않아 발생하는 문제이다.

identity map 에 저장하는 로직이 어떻게 되어 있는지는 공식문서에도 나와 있지 않으나, 백그라운드 스레드로 실행하는지 어쨋든 identity map 에 저장되기 전에 쿼리가 나가서 발생하는 문제인 것 같다.

sqlalchemy identity map 사용하기 - read only Transactional

기존에는 항상 쿼리로 조회해올 때, 항상 새로운 세션을 할당받아 쿼리를 수행했는데 조회만 일어나는 경우에는 같은 세션을 사용하여 cache hit 가 더 자주 발생하도록 하고, 불필요한 세션 생성을 줄이면 시간을 절약할 수 있을 것 같다.

from functools import wraps

from app.core.db.session import (
    session,
    get_session_context,
    add_transactional_session,
    remove_transactional_session,
    is_in_transactional_session,
)


class Transactional:
    def __init__(self, read_only=False):
        self.read_only = read_only

    def __call__(self, func):
        @wraps(func)
        async def _transactional(*args, **kwargs):
            context = get_session_context()

            if not is_in_transactional_session(context):
                add_transactional_session(context)

            try:
                result = await func(*args, **kwargs)

                if self.read_only:
                    await session.rollback()
                elif context.startswith("test-session"):
                    await session.flush()
                else:
                    await session.commit()

            except Exception as e:
                await session.rollback()
                raise e
            finally:
                remove_transactional_session(context)

            return result

        return _transactional

read_only 세션도 일반 트랜잭션과 동일하게 별도의 트랜잭션 세션으로 관리하여 동일한 조회 세션을 사용할 수 있도록 한다.
트랜잭션이 끝나면 롤백시킨다.

@asynccontextmanager
async def session_factory() -> AsyncGenerator[AsyncSession, None]:
    context = get_session_context()

    if context in transactional_sessions or context.startswith("test-session"):
        yield session
        return

    _session = async_sessionmaker(
        bind=async_engine,
        class_=AsyncSession,
        expire_on_commit=False,
    )()
    try:
        yield _session
    finally:
        await _session.close()

context in transactional_sessions 부분을 추가하여 read_only_transaction 세션일 경우 전역 세션을 반환하도록 한다.

최적화 결과

2862ms → 1318ms → 433ms (user_relation_id: 870)

2771ms → 386ms (user_relation_id: 693)

SQLAlchemy

FastAPI