From 81eef39bfdd643453c6392fe3f15653315f69094 Mon Sep 17 00:00:00 2001
From: Gianluca Ficarelli <gianluca.ficarelli@epfl.ch>
Date: Wed, 9 Oct 2024 09:58:05 +0200
Subject: [PATCH] Acquire exclusive lock during db migration

---
 alembic/env.py | 22 +++++++++++++---------
 1 file changed, 13 insertions(+), 9 deletions(-)

diff --git a/alembic/env.py b/alembic/env.py
index 25d85c5..d80114d 100644
--- a/alembic/env.py
+++ b/alembic/env.py
@@ -4,7 +4,7 @@
 from logging.config import fileConfig
 
 import alembic_postgresql_enum  # noqa: F401
-from sqlalchemy import pool
+from sqlalchemy import pool, text
 from sqlalchemy.engine import Connection
 from sqlalchemy.ext.asyncio import async_engine_from_config
 
@@ -16,6 +16,15 @@
 
 L = logging.getLogger("alembic.env")
 
+# server settings to reduce possible service disruptions while running the migration
+SERVER_SETTINGS = {
+    # Abort any statement that takes more than the specified amount of time
+    "statement_timeout": "6000",
+    # Abort any statement that waits longer than the specified amount of time while
+    # attempting to acquire a lock on a table, index, row, or other database object
+    "lock_timeout": "4000",
+}
+
 # this is the Alembic Config object, which provides
 # access to the values within the .ini file in use.
 config = context.config
@@ -87,6 +96,8 @@ def do_run_migrations(connection: Connection) -> None:
     )
 
     with context.begin_transaction():
+        # obtain an exclusive transaction-level advisory lock, waiting if necessary
+        connection.execute(text("SELECT pg_advisory_xact_lock(12345)"))
         context.run_migrations()
 
 
@@ -94,18 +105,11 @@ async def run_async_migrations() -> None:
     """In this scenario we need to create an Engine
     and associate a connection with the context.
     """
-    server_settings = {
-        # Abort any statement that takes more than the specified amount of time
-        "statement_timeout": "6000",
-        # Abort any statement that waits longer than the specified amount of time while
-        # attempting to acquire a lock on a table, index, row, or other database object
-        "lock_timeout": "4000",
-    }
     connectable = async_engine_from_config(
         config.get_section(config.config_ini_section, {}),
         prefix="sqlalchemy.",
         poolclass=pool.NullPool,
-        connect_args={"server_settings": server_settings},
+        connect_args={"server_settings": SERVER_SETTINGS},
     )
 
     async with connectable.connect() as connection: