Improve db migration process

2 files changed, 41 insertions, 12 deletions

diff --git a/apply_migrations.py b/apply_migrations.py
index 4ade81e..4b00dfe 100644
--- a/apply_migrations.py
+++ b/apply_migrations.py
@@ -9,6 +9,10 @@ migrations_dir = 'migrations'  # Relative directory path
 # Get all migration files
 migration_files = sorted(os.listdir(migrations_dir))
 
+db.ensure_migration_table_exists()
+
+num_applied = 0
+
 # Apply migrations in alphabetical order
 for migration_file in migration_files:
   migration_path = os.path.join(migrations_dir, migration_file)
@@ -18,8 +22,12 @@ for migration_file in migration_files:
     migration_sql = f.read()
 
   print(migration_sql)
+  number = int(migration_file.split('_')[0])
 
-  db.apply_migration(migration_sql)
-  print(f"* Migration applied: {migration_path}")
+  if db.ensure_migration_applied(number, migration_sql):
+    print(f"* Migration applied: {migration_path}\n")
+    num_applied += 1
+  else:
+    print(f"* Migration already applied: {migration_path}\n")
 
-print(f"* {len(migration_files)} migrations applied.")
+print(f"* {num_applied} migrations applied.")
diff --git a/db.py b/db.py
index 2ee4cc0..d91692a 100644
--- a/db.py
+++ b/db.py
@@ -219,15 +219,6 @@ def vote_poll(poll_id: str, voter_name: str, selections: dict[str, int]) -> None
       conn.rollback()
       raise e
 
-def apply_migration(migration_sql: str) -> None:
-  with db.cursor() as (conn, cur):
-    try:
-      cur.execute(migration_sql)
-      conn.commit()
-    except Exception as e:
-      conn.rollback()
-      raise e
-
 def get_poll_by_code(code: str) -> Poll | None:
   with db.cursor() as (conn, cur):
     try:
@@ -318,3 +309,33 @@ def delete_poll(code: str) -> None:
     except Exception as e:
       conn.rollback()
       raise e
+
+### Migrations
+
+def ensure_migration_table_exists() -> None:
+  with db.cursor() as (conn, cur):
+    try:
+      cur.execute("CREATE TABLE IF NOT EXISTS applied_migrations ("
+                  "number INTEGER PRIMARY KEY"
+                  ")")
+      conn.commit()
+    except Exception as e:
+      conn.rollback()
+      raise e
+
+def ensure_migration_applied(number: int, migration_sql: str) -> bool:
+  """Returns True if the migration was applied, False if it was already applied."""
+  with db.cursor() as (conn, cur):
+    try:
+      cur.execute("SELECT * FROM applied_migrations WHERE number = %s", (number,))
+      if cur.fetchone() is not None:
+        return False
+
+      cur.execute("INSERT INTO applied_migrations (number) VALUES (%s)", (number,))
+      cur.execute(migration_sql)
+      conn.commit()
+      return True
+
+    except Exception as e:
+      conn.rollback()
+      raise e