如何在Discord Bot中用SQLAlchemy初始化PostgreSQL并适配Cogs

原代码实现 #

我当前已编写的基础代码如下：

engine = create_engine(settings.database_url,
                   pool_size=40, max_overflow=80)

Session = sessionmaker(autocommit=False, autoflush=False, 
                                                  bind=engine)


def get_db():
    db = Session()
    try:
        yield db
    finally:
        db.close()


class Bot(commands.Bot):
    def __init__(self):
        self.description = "Bot - A Tipbot"

        super().__init__(
            command_prefix={"!"},
            owner_ids={settings.discord_bot_id},
            intents=discord.Intents.all(),
            help_command=PrettyHelp(),
            description=self.description,
            case_insensitive=True,
        )

    async def on_ready(self):
        logger.info(f'{self.user} is ready!')


client = Bot()


async def load_cogs():
    for filename in os.listdir("./app/cogs"):
        if filename.endswith(".py"):
            await client.load_extension(f"cogs.{filename[:-3]}")


async def main():
    await load_cogs()
    await client.start(settings.discord_bot_token)


asyncio.run(main())

我的尝试与疑问 #

我不清楚如何在Cogs中使用数据库连接，于是尝试在Bot类的__init__方法中初始化Session，以便在Cogs中通过self.session调用，修改后的Bot类代码如下：

class Bot(commands.Bot):
def __init__(self):
    self.description = "Bot - A Tipbot"

    super().__init__(
        command_prefix={"!"},
        owner_ids={settings.discord_bot_id},
        intents=discord.Intents.all(),
        help_command=PrettyHelp(),
        description=self.description,
        case_insensitive=True,
    )
self.session = sessionmaker(autocommit=False, autoflush=False, 
                                              bind=engine)
# 此处不确定如何调用get_db()

同时尝试在Cogs中初始化连接执行查询：

class WithdrawCommand(commands.Cog):
def __init__(self, client):
    self.client = client
    self.Session = self.Session.get_db()

@commands.Cog.listener()
async def on_ready(self):
    logger.info(f'{self} ON!')

@commands.command
async def withdraw(self, ctx):
    res = self.Session.query()
    await ctx.send(res)

但不确定这种实现方式是否正确，希望得到技术指导。

正确实现方案 #

核心问题分析 #

你当前的尝试存在几个关键问题：

1. 重复创建sessionmaker工厂：sessionmaker是全局工厂类，只需初始化一次，重复创建会浪费资源。
2. 错误的会话持有方式：在Cog的__init__中保存会话实例会导致连接长期占用，容易引发连接池耗尽、数据一致性问题。
3. 同步SQLAlchemy与异步事件循环的冲突：直接在Discord.py的异步上下文执行同步数据库操作会阻塞事件循环，影响Bot响应速度。

修正后的实现步骤 #

1. 保留全局的数据库配置

原代码中全局的engine和Session工厂无需修改，保持单例初始化即可：

engine = create_engine(settings.database_url, pool_size=40, max_overflow=80)
Session = sessionmaker(autocommit=False, autoflush=False, bind=engine)

2. Bot类无需额外修改

不需要在Bot类中重复创建sessionmaker，保持原有Bot类实现即可。

3. Cogs中安全使用数据库会话

每个命令/事件处理时，创建独立的会话，用完立即关闭，避免资源占用。同时用asyncio.to_thread包裹同步数据库操作，避免阻塞事件循环。

示例Cog代码：

import asyncio
from sqlalchemy.exc import SQLAlchemyError
import logger

class WithdrawCommand(commands.Cog):
    def __init__(self, client):
        self.client = client

    @commands.Cog.listener()
    async def on_ready(self):
        logger.info(f'{self} ON!')

    @commands.command()
    async def withdraw(self, ctx):
        # 定义同步数据库查询函数
        def perform_query():
            db = Session()
            try:
                # 替换为你的实际查询逻辑
                # res = db.query(YourModel).filter(...).first()
                res = "模拟查询结果"
                return res
            except SQLAlchemyError as e:
                logger.error(f"数据库查询错误: {e}")
                return None
            finally:
                db.close()
        
        # 在线程中执行同步操作，避免阻塞事件循环
        query_result = await asyncio.to_thread(perform_query)
        
        if query_result:
            await ctx.send(f"查询结果：{query_result}")
        else:
            await ctx.send("查询失败，请稍后重试")

4. 可选：复用get_db生成器

如果想复用原有的get_db生成器，可以在Cog中这样使用：

def get_db():
    db = Session()
    try:
        yield db
    finally:
        db.close()

class WithdrawCommand(commands.Cog):
    # ... __init__和on_ready方法不变 ...

    @commands.command()
    async def withdraw(self, ctx):
        async def query_db():
            for db in get_db():
                try:
                    res = db.query(YourModel).all()
                    return res
                except SQLAlchemyError as e:
                    logger.error(f"查询错误: {e}")
                    return None
        
        result = await asyncio.to_thread(query_db)
        await ctx.send(f"共查询到 {len(result) if result else 0} 条数据")

关键注意事项 #

• 会话隔离：每个命令/事件使用独立的会话，避免不同请求之间的数据干扰。
• 连接释放：务必在finally块中关闭会话，确保连接返回连接池。
• 异步适配：同步SQLAlchemy必须放在线程中执行，避免阻塞Discord.py的事件循环；如果追求更好的异步支持，可以改用SQLAlchemy 1.4+的异步模式（AsyncEngine和AsyncSession）。

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