🐛 fix(utils.py): handle CommandError when running migrations and delete alembic_version table if wrong revision in DB

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

@@ -1,7 +1,7 @@
 from typing import TYPE_CHECKING
 from langflow.utils.logger import logger
 from contextlib import contextmanager
-
+from alembic.util.exc import CommandError
 from sqlmodel import Session
 
 if TYPE_CHECKING:
@@ -13,7 +13,23 @@ def initialize_database():
     from langflow.services import service_manager, ServiceType
 
     database_manager = service_manager.get(ServiceType.DATABASE_MANAGER)
-    database_manager.run_migrations()
+    try:
+        database_manager.run_migrations()
+    except CommandError as exc:
+        if "Can't locate revision identified by" not in str(exc):
+            raise exc
+        # This means there's wrong revision in the DB
+        # We need to delete the alembic_version table
+        # and run the migrations again
+        logger.warning(
+            "Wrong revision in DB, deleting alembic_version table and running migrations again"
+        )
+        with session_getter(database_manager) as session:
+            session.execute("DROP TABLE alembic_version")
+        database_manager.run_migrations()
+    except Exception as exc:
+        logger.error(f"Error running migrations: {exc}")
+        raise RuntimeError("Error running migrations") from exc
     database_manager.create_db_and_tables()
     logger.debug("Database initialized")