从同步SQLAlchemy迁移到异步时，asyncpg如何设置keepalives参数？

问题 #

正尝试从同步SQLAlchemy迁移至异步架构，使用asyncpg替代psycopg2。同步初始化代码可正常运行：

SQLALCHEMY_DATABASE_URL = f"postgresql://{DB_USER}:{DB_PASSWORD}@{DB_HOST}:{DB_PORT}/{DB_NAME}"

engine = create_engine(
    SQLALCHEMY_DATABASE_URL,
    connect_args={
        "keepalives": 1,
        "keepalives_idle": 10,
        "keepalives_interval": 5,
        "keepalives_count": 5
    },
    pool_size=50,
    max_overflow=100
)

但修改为异步代码后出现错误：TypeError: connect() got an unexpected keyword argument 'keepalives'：

SQLALCHEMY_DATABASE_URL = f"postgresql+asyncpg://{DB_USER}:{DB_PASSWORD}@{DB_HOST}:{DB_PORT}/{DB_NAME}"

engine = create_async_engine(
    SQLALCHEMY_DATABASE_URL,
    connect_args={
        "keepalives": 1,
        "keepalives_idle": 10,
        "keepalives_interval": 5,
        "keepalives_count": 5
    },
    pool_size=50,
    max_overflow=100
)

完整报错栈信息：

File "/home/vovan/PycharmProjects/proj/venv/lib/python3.10/site-packages/starlette/routing.py", line 621, in lifespan
    async with self.lifespan_context(app):
  File "/home/vovan/PycharmProjects/proj/venv/lib/python3.10/site-packages/starlette/routing.py", line 518, in __aenter__
    await self._router.startup()
  File "/home/vovan/PycharmProjects/proj/venv/lib/python3.10/site-packages/starlette/routing.py", line 598, in startup
    await handler()
  File "/home/vovan/PycharmProjects/proj/main.py", line 57, in on_startup
    await create_tables()
  File "/home/vovan/PycharmProjects/proj/core/models/models.py", line 78, in create_tables
    async with engine.begin() as conn:
  File "/home/vovan/PycharmProjects/proj/venv/lib/python3.10/site-packages/sqlalchemy/ext/asyncio/base.py", line 60, in __aenter__
    return await self.start(is_ctxmanager=True)
  File "/home/vovan/PycharmProjects/proj/venv/lib/python3.10/site-packages/sqlalchemy/ext/asyncio/engine.py", line 609, in start
    await self.conn.start(is_ctxmanager=is_ctxmanager)
  File "/home/vovan/PycharmProjects/proj/venv/lib/python3.10/site-packages/sqlalchemy/ext/asyncio/engine.py", line 154, in start
    await (greenlet_spawn(self.sync_engine.connect))
  File "/home/vovan/PycharmProjects/proj/venv/lib/python3.10/site-packages/sqlalchemy/util/_concurrency_py3k.py", line 123, in greenlet_spawn
    result = context.switch(*args, **kwargs)
  File "/home/vovan/PycharmProjects/proj/venv/lib/python3.10/site-packages/sqlalchemy/future/engine.py", line 406, in connect
    return super(Engine, self).connect()
  File "/home/vovan/PycharmProjects/proj/venv/lib/python3.10/site-packages/sqlalchemy/engine/base.py", line 3204, in connect
    return self._connection_cls(self, close_with_result=close_with_result)
  File "/home/vovan/PycharmProjects/proj/venv/lib/python3.10/site-packages/sqlalchemy/engine/base.py", line 96, in __init__
    else engine.raw_connection()
  File "/home/vovan/PycharmProjects/proj/venv/lib/python3.10/site-packages/sqlalchemy/engine/base.py", line 3283, in raw_connection
    return self._wrap_pool_connect(self.pool.connect, _connection)
  File "/home/vovan/PycharmProjects/proj/venv/lib/python3.10/site-packages/sqlalchemy/engine/base.py", line 3250, in _wrap_pool_connect
    return fn()
  File "/home/vovan/PycharmProjects/proj/venv/lib/python3.10/site-packages/sqlalchemy/pool/base.py", line 310, in connect
    return _ConnectionFairy._checkout(self)
  File "/home/vovan/PycharmProjects/proj/venv/lib/python3.10/site-packages/sqlalchemy/pool/base.py", line 868, in _checkout
    fairy = _ConnectionRecord.checkout(pool)
  File "/home/vovan/PycharmProjects/proj/venv/lib/python3.10/site-packages/sqlalchemy/pool/base.py", line 476, in checkout
    rec = pool._do_get()
  File "/home/vovan/PycharmProjects/proj/venv/lib/python3.10/site-packages/sqlalchemy/pool/impl.py", line 145, in _do_get
    with util.safe_reraise():
  File "/home/vovan/PycharmProjects/proj/venv/lib/python3.10/site-packages/sqlalchemy/util/langhelpers.py", line 70, in __exit__
    compat.raise_(
  File "/home/vovan/PycharmProjects/proj/venv/lib/python3.10/site-packages/sqlalchemy/util/compat.py", line 207, in raise_
    raise exception
  File "/home/vovan/PycharmProjects/proj/venv/lib/python3.10/site-packages/sqlalchemy/pool/impl.py", line 143, in _do_get
    return self._create_connection()
  File "/home/vovan/PycharmProjects/proj/venv/lib/python3.10/site-packages/sqlalchemy/pool/base.py", line 256, in _create_connection
    return _ConnectionRecord(self)
  File "/home/vovan/PycharmProjects/proj/venv/lib/python3.10/site-packages/sqlalchemy/pool/base.py", line 371, in __init__
    self.__connect()
  File "/home/vovan/PycharmProjects/proj/venv/lib/python3.10/site-packages/sqlalchemy/pool/base.py", line 665, in __connect
    with util.safe_reraise():
  File "/home/vovan/PycharmProjects/proj/venv/lib/python3.10/site-packages/sqlalchemy/util/langhelpers.py", line 70, in __exit__
    compat.raise_(
  File "/home/vovan/PycharmProjects/proj/venv/lib/python3.10/site-packages/sqlalchemy/util/compat.py", line 207, in raise_
    raise exception
  File "/home/vovan/PycharmProjects/proj/venv/lib/python3.10/site-packages/sqlalchemy/pool/base.py", line 661, in __connect
    self.dbapi_connection = connection = pool._invoke_creator(self)
  File "/home/vovan/PycharmProjects/proj/venv/lib/python3.10/site-packages/sqlalchemy/engine/create.py", line 590, in connect
    return dialect.connect(*cargs, **cparams)
  File "/home/vovan/PycharmProjects/proj/venv/lib/python3.10/site-packages/sqlalchemy/engine/default.py", line 597, in connect
    return self.dbapi.connect(*cargs, **cparams)
  File "/home/vovan/PycharmProjects/proj/venv/lib/python3.10/site-packages/sqlalchemy/dialects/postgresql/asyncpg.py", line 758, in connect
    await_only(self.asyncpg.connect(*arg, **kw)),
TypeError: connect() got an unexpected keyword argument 'keepalives'

解决方案 #

报错原因是psycopg2与asyncpg的连接参数命名规则不同，SQLAlchemy会将connect_args中的参数直接传递给底层数据库驱动，而asyncpg不支持psycopg2的keepalives系列参数名，需要替换为asyncpg对应的参数：

修改后的异步代码如下：

SQLALCHEMY_DATABASE_URL = f"postgresql+asyncpg://{DB_USER}:{DB_PASSWORD}@{DB_HOST}:{DB_PORT}/{DB_NAME}"

engine = create_async_engine(
    SQLALCHEMY_DATABASE_URL,
    connect_args={
        "tcp_keepalive": True,
        "tcp_keepidle": 10,
        "tcp_keepintvl": 5,
        "tcp_keepcnt": 5
    },
    pool_size=50,
    max_overflow=100
)

内容的提问来源于stack exchange，提问作者Ivan