Fix migration issue with autogenerate diffs

src/backend/langflow/services/database/service.py

@@ -5,17 +5,16 @@ from typing import TYPE_CHECKING
 import sqlalchemy as sa
 from alembic import command, util
 from alembic.config import Config
-from loguru import logger
-from sqlalchemy import inspect
-from sqlalchemy.exc import OperationalError
-from sqlmodel import Session, SQLModel, create_engine
-
 from langflow.services.base import Service
 from langflow.services.database import models  # noqa
 from langflow.services.database.models.user.crud import get_user_by_username
 from langflow.services.database.utils import Result, TableResults
 from langflow.services.deps import get_settings_service
 from langflow.services.utils import teardown_superuser
+from loguru import logger
+from sqlalchemy import inspect
+from sqlalchemy.exc import OperationalError
+from sqlmodel import Session, SQLModel, create_engine
 
 if TYPE_CHECKING:
     from sqlalchemy.engine import Engine
@@ -118,9 +117,10 @@ class DatabaseService(Service):
         alembic_cfg.set_main_option("script_location", str(self.script_location))
         alembic_cfg.set_main_option("sqlalchemy.url", self.database_url)
         command.stamp(alembic_cfg, "head")
+        # command.upgrade(alembic_cfg, "head")
         logger.info("Alembic initialized")
 
-    def run_migrations(self):
+    def run_migrations(self, fix=False):
         # First we need to check if alembic has been initialized
         # If not, we need to initialize it
         # if not self.script_location.exists(): # this is not the correct way to check if alembic has been initialized
@@ -151,16 +151,32 @@ class DatabaseService(Service):
             if isinstance(exc, util.exc.CommandError) or isinstance(exc, util.exc.AutogenerateDiffsDetected):
                 command.upgrade(alembic_cfg, "head")
 
-        # We should check the schema health after running migrations
         try:
             command.check(alembic_cfg)
-        except util.exc.AutogenerateDiffsDetected:
-            # downgrade to base and upgrade again
-            logger.warning("Autogenerate diffs detected, downgrading and upgrading")
-            command.downgrade(alembic_cfg, "-1")
-            # wait for the database to be ready
-            time.sleep(5)
-            command.upgrade(alembic_cfg, "head")
+        except util.exc.AutogenerateDiffsDetected as exc:
+            logger.exception("AutogenerateDiffsDetected: {exc}")
+            logger.warning("Something went wrong running migrations. Please, run `langflow migration --fix`")
+
+        if fix:
+            self.try_downgrade_upgrade_until_success(alembic_cfg)
+
+    def try_downgrade_upgrade_until_success(self, alembic_cfg, retries=5):
+        # Try -1 then head, if it fails, try -2 then head, etc.
+        # until we reach the number of retries
+        for i in range(1, retries + 1):
+            try:
+                command.check(alembic_cfg)
+                break
+            except util.exc.AutogenerateDiffsDetected as exc:
+
+                # downgrade to base and upgrade again
+                logger.warning(f"AutogenerateDiffsDetected: {exc}")
+                command.downgrade(alembic_cfg, f"-{i}")
+                # wait for the database to be ready
+                time.sleep(3)
+                command.upgrade(alembic_cfg, "head")
+
+
 
     def run_migrations_test(self):
         # This method is used for testing purposes only