refactor: use one engine for database migrations and bot (#688)

Author: onerandomusername

monty/__main__.py

@@ -8,6 +8,7 @@
 import disnake
 import redis
 import redis.asyncio
+from sqlalchemy.ext.asyncio import create_async_engine
 
 from monty import constants, monkey_patches
 from monty.bot import Monty
@@ -56,19 +57,24 @@ async def main() -> None:
         constants.Client.config_prefix, session=redis_session, prefix=constants.Redis.prefix
     )
 
+    database_engine = create_async_engine(constants.Database.postgres_bind)
     # run alembic migrations
     if constants.Database.run_migrations:
         log.info(f"Running database migrations to target {constants.Database.migration_target}")
-        await run_alembic()
+        await run_alembic(database_engine)
     else:
         log.info("Skipping database migrations per environment settings.")
+        # we still need to connect to the database to verify connection info is correct
+        async with database_engine.connect():
+            pass
 
     # ping redis
     await redis_session.ping()
     log.debug("Successfully pinged redis server.")
 
     bot = Monty(
         redis_session=redis_session,
+        database_engine=database_engine,
         command_prefix=constants.Client.default_command_prefix,
         activity=constants.Client.activity,
         allowed_mentions=constants.Client.allowed_mentions,

monty/bot.py

@@ -10,7 +10,7 @@
 import redis.asyncio
 import sqlalchemy as sa
 from disnake.ext import commands
-from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker, create_async_engine
+from sqlalchemy.ext.asyncio import AsyncEngine, AsyncSession, async_sessionmaker
 from sqlalchemy.orm import selectinload
 
 from monty import constants
@@ -50,7 +50,13 @@ class Monty(commands.Bot):
 
     name = constants.Client.name
 
-    def __init__(self, redis_session: redis.asyncio.Redis, proxy: str | None = None, **kwargs) -> None:
+    def __init__(
+        self,
+        redis_session: redis.asyncio.Redis,
+        database_engine: AsyncEngine,
+        proxy: str | None = None,
+        **kwargs,
+    ) -> None:
         if TEST_GUILDS:
             kwargs["test_guilds"] = TEST_GUILDS
             log.warning("registering as test_guilds")
@@ -62,8 +68,8 @@ def __init__(self, redis_session: redis.asyncio.Redis, proxy: str | None = None,
 
         self.create_http_session(proxy=proxy)
 
-        self.db_engine = engine = create_async_engine(constants.Database.postgres_bind)
-        self.db_session = async_sessionmaker(engine, expire_on_commit=False, class_=AsyncSession)
+        self.db_engine = database_engine
+        self.db_session = async_sessionmaker(database_engine, expire_on_commit=False, class_=AsyncSession)
 
         self.guild_configs: dict[int, GuildConfig] = {}
         self.guild_db: dict[int, Guild] = {}
@@ -161,8 +167,6 @@ async def get_prefix(self, message: disnake.Message) -> list[str] | str | None:
 
     async def _create_features(self) -> None:
         """Update the database with all features defined immediately upon launch. No more lazy creation."""
-        await self.wait_until_first_connect()
-
         async with self._feature_db_lock, self.db.begin() as session:
             stmt = sa.select(Feature).options(selectinload(Feature.rollout))
             result = await session.scalars(stmt)

monty/migrations.py

@@ -23,7 +23,7 @@ async def run_async_upgrade(engine: AsyncEngine) -> None:
         await conn.run_sync(run_upgrade, alembic_cfg)
 
 
-async def run_alembic() -> None:
+async def run_alembic(engine: AsyncEngine | None = None) -> None:
     """Run alembic migrations."""
-    engine = create_async_engine(constants.Database.postgres_bind)
+    engine = engine or create_async_engine(constants.Database.postgres_bind)
     await run_async_upgrade(engine)