1
\$\begingroup\$

I have a prototype pattern for leasing sessions to a PostgreSQL database and would like a review. I took inspiration from this project on github.

from functools import wraps
from typing import Callable, Any, Sequence

from enum import Enum
from typing import AsyncGenerator

from contextlib import asynccontextmanager
from contextvars import ContextVar, Token
from sqlalchemy import text, select
from sqlalchemy.ext.asyncio import (
    AsyncSession,
    async_scoped_session,
    async_sessionmaker,
    create_async_engine,
)
from sqlalchemy.orm import DeclarativeBase, Session, Mapped, mapped_column
from sqlalchemy.sql.expression import Delete, Insert, Update

from core.config import config


SESSION_CONTEXT: ContextVar[str] = ContextVar("session_context")
POOL_RECYCLE = 3600


class EngineType(Enum):
    WRITER = "writer"
    READER = "reader"


ENGINES = {
    EngineType.WRITER: create_async_engine(config.WRITER_DB_URL, pool_recycle=POOL_RECYCLE),
    EngineType.READER: create_async_engine(config.READER_DB_URL, pool_recycle=POOL_RECYCLE),
}


def get_session_context() -> str:
    return SESSION_CONTEXT.get()


def set_session_context(session_id: str) -> Token:
    return SESSION_CONTEXT.set(session_id)


def reset_session_context(context: Token) -> None:
    SESSION_CONTEXT.reset(context)


class RoutingSession(Session):
    def get_bind(self, mapper = None, clause = None, **kwargs):
        if self._flushing or isinstance(clause, (Update, Delete, Insert)):
            return ENGINES[EngineType.WRITER].sync_engine
        else:
            return ENGINES[EngineType.READER].sync_engine


_ASYNC_SESSION_FACTORY = async_sessionmaker(
    class_=AsyncSession,
    sync_session_class=RoutingSession,
    expire_on_commit=False,
)
session = async_scoped_session(
    session_factory=_ASYNC_SESSION_FACTORY,
    scopefunc=get_session_context,
)


class Base(DeclarativeBase):
    ...


@asynccontextmanager
async def session_factory() -> AsyncGenerator[AsyncSession, None]:
    _session = async_sessionmaker(
        class_=AsyncSession,
        sync_session_class=RoutingSession,
        expire_on_commit=False,
    )()
    try:
        yield _session
    finally:
        await _session.close()


class Transactional:

    def __init__(self, is_readonly: bool):
        self.is_readonly = is_readonly

    def __call__(self, func: Callable[..., Any]) -> Callable[..., Any]:
        @wraps(func)
        async def _transactional(*args, **kwargs):
            async with session_factory() as session:
                try:
                    result = await func(*args, **kwargs, session=session)
                    if not self.is_readonly:
                        await session.commit()
                    return result
                except Exception as e:
                    await session.rollback()
                    raise e

        return _transactional


class Test(Base):

    __tablename__ = 'test'

    id: Mapped[str] = mapped_column(primary_key=True)


class TestRepository:

    @Transactional(is_readonly=False)
    async def create(self, *, record: Test, session: AsyncSession):
        session.add(record)

    @Transactional(is_readonly=True)
    async def get(
        self,
        *,
        test_id,
        session: AsyncSession,
    ) -> Test | None:
        query = select(Test).where(Test.id == test_id)
        result = await session.execute(query)
        data = result.scalars().first()
        return data


# For demo purposes...
async def find_record():
    # Get record
    repo = TestRepository()
    result = await repo.get(test_id="ss")
    return result

if __name__ == "__main__":
    data = find_record()
    print(data)
\$\endgroup\$

1 Answer 1

1
\$\begingroup\$

lock duration

I'm skeptical we want do this conditionally:

                    if not self.is_readonly:
                        await session.commit()

For example I have seen this happen in production:

  • Popular web application query does SELECT a,b,c and reads some of those rows.
  • Eventually that TCP connection to the postgres server gets closed, but by then new readers have cropped up.
  • DBA attempts ALTER TABLE to rename column x to y.

So the SELECT is holding a reader lock. Since the DBA needs an exclusive lock, the column rename never succeeds.

Changing the app to consume all rows up to the specified LIMIT immediately fixed up the observed behavior, as that allows closing the no-longer-active connection. Equivalently we found that explicit COMMIT is also observed to release the reader lock.

Strictly speaking this is still racy, improving the chance of winning the race. An app might choose to add support for a DBA to briefly lock out readers in order to acquire an exclusive lock with minimal impact on uptime. Changing the DB connect password and then changing it back can produce similar effect.

annotations

In recent years these have been a WIP. At some point the typing module will deprecate portions and will become much smaller than it is now. It doesn't hurt to pick up good annotating habits now.

from typing import Callable, Any, Sequence

For the first and last, prefer this:

from collections.abc import Callable, Sequence
\$\endgroup\$
1
  • \$\begingroup\$ Thanks these suggestions have been noted. However, given the size of the bounty, I am holding out for intellectual annihilation :) \$\endgroup\$
    – Bob
    Commented Jun 25 at 0:53

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service and acknowledge you have read our privacy policy.

Not the answer you're looking for? Browse other questions tagged or ask your own question.