From 163ef6973bf3996475815354f498718937d5c7d6 Mon Sep 17 00:00:00 2001 From: Gabriel Luiz Freitas Almeida Date: Fri, 1 Dec 2023 16:30:26 -0300 Subject: [PATCH] Fix database initialization and migration --- src/backend/langflow/services/database/utils.py | 11 ++++++----- 1 file changed, 6 insertions(+), 5 deletions(-) diff --git a/src/backend/langflow/services/database/utils.py b/src/backend/langflow/services/database/utils.py index aae567e91..826b8d3a5 100644 --- a/src/backend/langflow/services/database/utils.py +++ b/src/backend/langflow/services/database/utils.py @@ -1,17 +1,18 @@ +from contextlib import contextmanager from dataclasses import dataclass from typing import TYPE_CHECKING -from loguru import logger -from contextlib import contextmanager + from alembic.util.exc import CommandError +from loguru import logger from sqlmodel import Session if TYPE_CHECKING: from langflow.services.database.service import DatabaseService -def initialize_database(): +def initialize_database(fix_migration: bool = False): logger.debug("Initializing database") - from langflow.services import service_manager, ServiceType + from langflow.services import ServiceType, service_manager database_service: "DatabaseService" = service_manager.get(ServiceType.DATABASE_SERVICE) try: @@ -28,7 +29,7 @@ def initialize_database(): logger.error(f"Error checking schema health: {exc}") raise RuntimeError("Error checking schema health") from exc try: - database_service.run_migrations() + database_service.run_migrations(fix=fix_migration) except CommandError as exc: # if "overlaps with other requested revisions" or "Can't locate revision identified by" # are not in the exception, we can't handle it