diff --git a/src/backend/langflow/services/database/service.py b/src/backend/langflow/services/database/service.py index f41a2484d..3ecbfd971 100644 --- a/src/backend/langflow/services/database/service.py +++ b/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