SQLAlchemy 1.4.46提交触发IllegalStateChangeError的原因及解决

升级集成asyncio的SQLAlchemy至1.4.46版本后，执行commit操作时触发sqlalchemy.exc.IllegalStateChangeError异常，错误提示：

Method 'commit()' can't be called here; method '_connection_for_bind()' is already in progress and this would cause an unexpected state change to <SessionTransactionState.CLOSED: 5>

此前版本运行正常，相关实现代码及调用方式如下：

# -*- coding:utf-8 -*-

from sqlalchemy import exc, event, text
from sqlalchemy.ext.asyncio import create_async_engine, AsyncSession as SQLAlchemyAsyncSession
from sqlalchemy.orm import sessionmaker, Session
from sqlalchemy.ext.asyncio import AsyncEngine
from sqlalchemy.pool import NullPool, Pool
from contextvars import ContextVar
from sanic import Sanic
import asyncio


class EngineNotInitialisedError(Exception):
    pass


class DBSessionContext:
    def __init__(self, session: Session, commit_on_exit: bool = True) -> None:
        self.session = session
        self._query = None
        self.commit_on_exit = commit_on_exit
        self.token = None

    async def close(self, exc_type=None, exc_value=None, traceback=None):
        if self._query:
            if exc_value and getattr(exc_value, 'status_code', 500) > 300:
                await self._query.rollback()
                self._post_processing.clear()
            else:
                await self._query.commit()
                await self.run_post_processing()

            await self._query.close()

        if self._post_processing:
            await self.run_post_processing()

    def set_token(self, token):
        self.token = token

    @property
    def query(self) -> Session:
        if not self._query:
            self._query = self.session()

        return self._query


class AsyncSession(SQLAlchemyAsyncSession):
    async def execute(self, statement, **parameters):
        try:
            if isinstance(statement, str):
                # We wrap around the `text()` method automatically
                statement = text(statement)
            return await super().execute(statement, parameters)
        except exc.OperationalError as e:
            if e.orig.args[0] == 1205:
                # Lock wait timeout exceeded
                await self.rollback()
                return await super().execute(statement, parameters)

            raise e


class DBSession:
    def __init__(self):
        self.engine = None
        self.session = None
        self._session = None
        self.context = ContextVar("context", default=None)

    def init_app(self, app: Sanic, url: str, commit_on_exit: bool = True) -> None:
        self.commit_on_exit = commit_on_exit

        engine_args = {
            'echo': app.config.get('DATABASE_ECHO', cast=bool, default=False),
            'echo_pool': app.config.get('DATABASE_ECHO_POOL', cast=bool, default=False),
            'poolclass': NullPool,  # will be used to create a connection pool instance using the connection parameters given in the URL
            # if pool_class is not NullPool:

            # the number of connections to allow in connection pool “overflow”
            # 'max_overflow': app.config.get('DATABASE_MAX_OVERFLOW', cast=int, default=10),
            # if True will enable the connection pool “pre-ping” feature that tests connections for liveness upon each checkout
            # 'pool_pre_ping': app.config.get('DATABASE_POOL_PRE_PING', cast=bool, default=True),
            # the number of connections to keep open inside the connection pool
            # 'pool_size': app.config.get('DATABASE_POOL_SIZE', cast=int, default=5),
            # this setting causes the pool to recycle connections after the given number of seconds has passed
            # 'pool_recycle': app.config.get('DATABASE_POOL_RECYCLE', cast=int, default=-1),
            # number of seconds to wait before giving up on getting a connection from the pool
            # 'pool_timeout': app.config.get('DATABASE_POOL_TIMEOUT', cast=int, default=3600),
        }

        self.engine = create_async_engine(
            url,
            **engine_args
        )

        self.session = sessionmaker(
            bind=self.engine,
            expire_on_commit=False,
            class_=AsyncSession,
            autoflush=False
        )

    async def __aenter__(self):
        if not isinstance(self.engine, AsyncEngine):
            raise EngineNotInitialisedError

        session_ctx = DBSessionContext(self.session, self.commit_on_exit)
        session_ctx.set_token(self.context.set(session_ctx))

        return session_ctx

    async def __aexit__(self, exc_type, exc_value, traceback):
        session_ctx = self.context.get()
        await asyncio.shield(session_ctx.close(exc_type, exc_value, traceback))

        self.context.reset(session_ctx.token)

    @property
    def query(self) -> Session:
        return self.context.get().query


@event.listens_for(Pool, "checkout")
def check_connection(dbapi_con, con_record, con_proxy):
    '''Listener for Pool checkout events that pings every connection before using.
    Implements pessimistic disconnect handling strategy. See also:
    http://docs.sqlalchemy.org/en/rel_0_8/core/pooling.html#disconnect-handling-pessimistic'''

    cursor = dbapi_con.cursor()
    try:
        cursor.execute("SELECT 1")
    except exc.OperationalError as ex:
        if ex.args[0] in (2006,   # MySQL server has gone away
                          2013,   # Lost connection to MySQL server during query
                          2055):  # Lost connection to MySQL server at '%s', system error: %d
            raise exc.DisconnectionError()  # caught by pool, which will retry with a new connection
        else:
            raise

    cursor.close()


db = DBSession()

调用方式：

async with db:
    await db.query.execute('INSERT INTO ...')

1. 异步会话初始化不规范：DBSessionContext的query属性中，直接通过self.session()同步创建异步会话，但SQLAlchemy 1.4+的异步会话工厂（由sessionmaker(class_=AsyncSession)生成）返回的是异步对象，必须通过await完成初始化。同步创建会导致会话内部的连接建立流程未完成，后续调用commit()时，会话仍处于连接操作的中间状态，触发严格的状态校验错误。
2. 新版本状态校验强化：SQLAlchemy 1.4.46对异步会话的状态流转逻辑做了更严格的检查，禁止在连接操作（_connection_for_bind()）执行期间调用commit()/rollback()这类会改变会话状态的方法，旧版本的宽松校验掩盖了代码中的异步调用问题。
3. 代码逻辑缺失：DBSessionContext中未定义_post_processing属性和run_post_processing方法，会导致close方法执行时触发额外异常，进一步干扰会话状态流转。

针对上述问题，对代码做如下修改：

1. 修正异步会话创建逻辑 #

将DBSessionContext的query属性改为异步属性，通过await正确初始化异步会话：

class DBSessionContext:
    def __init__(self, session: Session, commit_on_exit: bool = True) -> None:
        self.session = session
        self._query = None
        self.commit_on_exit = commit_on_exit
        self.token = None
        # 初始化_post_processing列表
        self._post_processing = []

    async def get_query(self) -> SQLAlchemyAsyncSession:
        if not self._query:
            # 异步会话必须通过await初始化
            self._query = await self.session()
        return self._query

    # 改为异步属性
    @property
    async def query(self) -> SQLAlchemyAsyncSession:
        return await self.get_query()

    # 补充run_post_processing方法实现
    async def run_post_processing(self):
        for task in self._post_processing:
            await task
        self._post_processing.clear()

2. 修复close方法的异常处理 #

为close方法添加try/finally确保会话始终被关闭，同时处理异常场景：

async def close(self, exc_type=None, exc_value=None, traceback=None):
    if self._query:
        try:
            if exc_value and getattr(exc_value, 'status_code', 500) > 300:
                await self._query.rollback()
                self._post_processing.clear()
            else:
                await self._query.commit()
                await self.run_post_processing()
        finally:
            # 无论是否异常，确保会话关闭
            await self._query.close()

    if self._post_processing:
        await self.run_post_processing()

3. 调整DBSession的query访问方式 #

由于query变为异步属性，修改DBSession提供异步获取接口：

class DBSession:
    # ... 其他代码保持不变 ...
    async def get_query(self) -> SQLAlchemyAsyncSession:
        session_ctx = self.context.get()
        return await session_ctx.query

对应的调用方式改为：

async with db:
    query = await db.get_query()
    await query.execute('INSERT INTO ...')

4. 移除不必要的asyncio.shield #

__aexit__中的asyncio.shield会掩盖会话关闭时的异常，且当前场景下无必要，直接移除：

async def __aexit__(self, exc_type, exc_value, traceback):
    session_ctx = self.context.get()
    await session_ctx.close(exc_type, exc_value, traceback)
    self.context.reset(session_ctx.token)

内容的提问来源于stack exchange，提问作者Cyril N.