How to initialize postgres with sqlalchemy to a discord bot-CodePudding

Currently I have:

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())

But I don't know how to use it with my cogs, should I initialize the session in the class Bot, so inside the cogs I can use self.session? Like:

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() ....

So in each cogs I can simply initialize my db connection to make querys:

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)

Or something like that I was thinking, but I'm not sure if this is the right way

CodePudding user response：

The right way is to indeed add an attribute (or property, to make your life easier) on your Bot class & access that within your cogs. However, you should always create a new session whenever you need it & destroy it when you don't need it anymore. Don't keep one as an attribute in your Cog forever.

Session = sessionmaker(..., class_=AsyncSession)

class MyBot(commands.Bot):
    ...
    @property
    def session(self) -> AsyncSession:
        # Let the sessionmaker create a new session
        return Session()

and then you can use it as

class SomeCog(commands.Cog):
    bot: MyBot

    def __init__(self, bot: MyBot):
        self.bot = bot
        # Don't create a session in here

    ...
    async def some_command(self, ctx):
        # Using the session as a context manager automatically closes it afterwards
        async with self.bot.session as session:
            await do_whatever_with_your(session)

Note that you should use asynchronous SQLAlchemy to avoid freezing your bot while running database queries. There's a lot of info about it in the docs: https://docs.sqlalchemy.org/en/14/orm/extensions/asyncio.html#synopsis-orm

Discord.py docs about blocking: https://discordpy.readthedocs.io/en/stable/faq.html#what-does-blocking-mean

SQLAlchemy & asyncpg both already implement a connection pool for you to allow multiple simultaneous connections, so you don't have to worry about that part.