From 020fe91267f24edbf780dca5862edbbccae29bd1 Mon Sep 17 00:00:00 2001
From: pablohashescobar <nikhilschacko@gmail.com>
Date: Fri, 19 Jan 2024 12:35:43 +0530
Subject: [PATCH] dev: create safe migration script so that the migration
 doesn't fail when running in replicas

---
 apiserver/bin/takeoff                         |  2 +-
 apiserver/bin/takeoff.local                   |  2 +-
 .../db/management/commands/safe_migrate.py    | 48 +++++++++++++++++++
 3 files changed, 50 insertions(+), 2 deletions(-)
 create mode 100644 apiserver/plane/db/management/commands/safe_migrate.py

diff --git a/apiserver/bin/takeoff b/apiserver/bin/takeoff
index 0ec2e495c..e7b681f69 100755
--- a/apiserver/bin/takeoff
+++ b/apiserver/bin/takeoff
@@ -1,7 +1,7 @@
 #!/bin/bash
 set -e
 python manage.py wait_for_db
-python manage.py migrate
+python manage.py safe_migrate
 
 # Create the default bucket
 #!/bin/bash
diff --git a/apiserver/bin/takeoff.local b/apiserver/bin/takeoff.local
index b89c20874..a3745eeea 100755
--- a/apiserver/bin/takeoff.local
+++ b/apiserver/bin/takeoff.local
@@ -1,7 +1,7 @@
 #!/bin/bash
 set -e
 python manage.py wait_for_db
-python manage.py migrate
+python manage.py safe_migrate
 
 # Create the default bucket
 #!/bin/bash
diff --git a/apiserver/plane/db/management/commands/safe_migrate.py b/apiserver/plane/db/management/commands/safe_migrate.py
new file mode 100644
index 000000000..fbe026670
--- /dev/null
+++ b/apiserver/plane/db/management/commands/safe_migrate.py
@@ -0,0 +1,48 @@
+# Python imports
+import uuid
+
+# Django imports
+from django.core.management.base import BaseCommand
+from django.core.management import call_command
+from django.db.migrations.executor import MigrationExecutor
+from django.db import connections, DEFAULT_DB_ALIAS
+
+# Module imports
+from plane.settings.redis import redis_instance
+
+
+class Command(BaseCommand):
+    help = 'Run migrations with Redis distributed locking if there are pending migrations'
+
+    def handle(self, *args, **kwargs):
+        # Check for pending migrations
+        if not self._pending_migrations():
+            self.stdout.write("No pending migrations.")
+            return
+
+        # Proceed with acquiring the lock and running migrations
+        lock_key = 'django_migration_lock'
+        lock_value = str(uuid.uuid4())  # Unique value for the lock
+        client = redis_instance()
+
+        # Try acquiring the lock
+        if client.set(lock_key, lock_value, nx=True, ex=300):  # 5 minutes expiry
+            try:
+                self.stdout.write("Acquired migration lock, running migrations...")
+                call_command('migrate')
+            finally:
+                # Release the lock if it belongs to this process
+                stored_value = client.get(lock_key)
+                if stored_value and stored_value.decode() == lock_value:
+                    client.delete(lock_key)
+                    self.stdout.write("Released migration lock.")
+                else:
+                    self.stdout.write("Lock was not released, as it was held by another instance.")
+        else:
+            self.stdout.write("Migration lock is held by another instance.")
+
+    def _pending_migrations(self):
+        connection = connections[DEFAULT_DB_ALIAS]
+        executor = MigrationExecutor(connection)
+        targets = executor.loader.graph.leaf_nodes()
+        return bool(executor.migration_plan(targets))