Home > Back-end >  How to initialize postgres with sqlalchemy to a discord bot
How to initialize postgres with sqlalchemy to a discord bot

Time:01-04

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.

  • Related