diff --git a/src/backend/langflow/services/database/utils.py b/src/backend/langflow/services/database/utils.py index 20b2bbbb4..94bcd6651 100644 --- a/src/backend/langflow/services/database/utils.py +++ b/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")