import logging
|
import time
|
|
from django.conf import settings
|
from django.core.management.commands.migrate import Command as MigrateCommand
|
from django.db import connections, transaction
|
|
logger = logging.getLogger(__name__)
|
|
DEFAULT_LOCK_ID = getattr(settings, 'MIGRATE_LOCK_ID', 1000)
|
|
LOCKED_MIGRATE_CMD_CONNECTION_ALIAS = 'locked_migrate_cmd_connection'
|
RETRY_INTERVAL = 5 # Time to wait between retries in seconds
|
MAX_WAIT_TIME = 300 # Maximum time to wait for the lock in seconds (5 minutes)
|
|
|
class Command(MigrateCommand):
|
help = 'Run Django migrations safely, using a lock'
|
|
def add_arguments(self, parser):
|
MigrateCommand.add_arguments(self, parser)
|
parser.add_argument(
|
'--migrate-lock-id',
|
default=DEFAULT_LOCK_ID,
|
type=int,
|
help='The id of the advisory lock to use',
|
)
|
|
def handle(self, *args, **options):
|
lock_id = options.pop('migrate_lock_id')
|
|
# Create a separate database connection to hold the lock
|
separate_lock_connection = connections.create_connection('default')
|
connections[LOCKED_MIGRATE_CMD_CONNECTION_ALIAS] = separate_lock_connection
|
try:
|
# Use a transaction to hold the lock for the duration of the migration
|
with transaction.atomic(using=LOCKED_MIGRATE_CMD_CONNECTION_ALIAS):
|
# Attempt to acquire the lock with retries
|
self.acquire_lock_with_retry(separate_lock_connection, lock_id)
|
# Run the standard Django migration once lock is acquired
|
super().handle(*args, **options)
|
logger.info('Migration complete, the migration lock has now been released.')
|
finally:
|
# Ensure the lock connection is closed to free resources
|
separate_lock_connection.close()
|
|
def acquire_lock_with_retry(self, lock_connection, lock_id):
|
start_time = time.time()
|
|
while True:
|
with lock_connection.cursor() as cursor:
|
logger.info(f'Attempting to acquire the postgres advisory transaction lock with id: {lock_id}.')
|
|
# Attempt to acquire the transaction-level lock without blocking
|
cursor.execute(f'SELECT pg_try_advisory_xact_lock({lock_id})')
|
lock_acquired = cursor.fetchone()[0]
|
|
if lock_acquired:
|
logger.info('Acquired the transaction lock, proceeding with migration.')
|
return # Exit the function if the lock is acquired
|
|
# Check if the maximum wait time has been reached
|
elapsed_time = time.time() - start_time
|
if elapsed_time >= MAX_WAIT_TIME:
|
logger.info('Could not acquire the transaction lock within the timeout period.')
|
raise TimeoutError('Failed to acquire PostgreSQL advisory transaction lock within 5 minutes.')
|
|
# Wait before retrying
|
logger.info(f'Lock not acquired. Retrying in {RETRY_INTERVAL} seconds...')
|
time.sleep(RETRY_INTERVAL)
|
