From 37dd5ba3e868f166f909b910326664adb90d0d61 Mon Sep 17 00:00:00 2001
From: stijndcl <declercq.stijn@outlook.com>
Date: Sat, 11 Jun 2022 01:26:19 +0200
Subject: [PATCH] Check current migration on startup

---
 database/migrations.py | 26 ++++++++++++++++++++++++++
 main.py                | 10 ++++++++--
 2 files changed, 34 insertions(+), 2 deletions(-)
 create mode 100644 database/migrations.py

diff --git a/database/migrations.py b/database/migrations.py
new file mode 100644
index 0000000..464c811
--- /dev/null
+++ b/database/migrations.py
@@ -0,0 +1,26 @@
+import logging
+
+from alembic import config, script
+from alembic.runtime import migration
+
+from database.engine import engine
+
+
+async def ensure_latest_migration():
+    """Make sure we are currently on the latest revision, otherwise raise an exception"""
+    alembic_config = config.Config("alembic.ini")
+    alembic_script = script.ScriptDirectory.from_config(alembic_config)
+
+    async with engine.begin() as connection:
+        current_revision = await connection.run_sync(
+            lambda sync_connection: migration.MigrationContext.configure(sync_connection).get_current_revision()
+        )
+
+        alembic_head = alembic_script.get_current_head()
+
+        if current_revision != alembic_head:
+            error_message = (
+                f"Pending migrations (current revision is {current_revision}, while head is at {alembic_head})"
+            )
+            logging.error(error_message)
+            raise RuntimeError(error_message)
diff --git a/main.py b/main.py
index 58f9e1b..cb4aa8e 100644
--- a/main.py
+++ b/main.py
@@ -4,6 +4,7 @@ from logging.handlers import RotatingFileHandler
 import asyncio
 
 import settings
+from database.migrations import ensure_latest_migration
 from didier import Didier
 
 
@@ -28,7 +29,12 @@ def setup_logging():
     logging.getLogger("discord").setLevel(logging.ERROR)
 
 
-if __name__ == "__main__":
+async def main():
+    """Do some setup & checks, and then run the bot"""
     setup_logging()
+    await ensure_latest_migration()
+    await run_bot()
 
-    asyncio.run(run_bot())
+
+if __name__ == "__main__":
+    asyncio.run(main())