Fix database initialization and migration

This commit is contained in:
Gabriel Luiz Freitas Almeida 2023-12-01 16:30:26 -03:00
commit 163ef6973b

View file

@ -1,17 +1,18 @@
from contextlib import contextmanager
from dataclasses import dataclass from dataclasses import dataclass
from typing import TYPE_CHECKING from typing import TYPE_CHECKING
from loguru import logger
from contextlib import contextmanager
from alembic.util.exc import CommandError from alembic.util.exc import CommandError
from loguru import logger
from sqlmodel import Session from sqlmodel import Session
if TYPE_CHECKING: if TYPE_CHECKING:
from langflow.services.database.service import DatabaseService from langflow.services.database.service import DatabaseService
def initialize_database(): def initialize_database(fix_migration: bool = False):
logger.debug("Initializing database") 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) database_service: "DatabaseService" = service_manager.get(ServiceType.DATABASE_SERVICE)
try: try:
@ -28,7 +29,7 @@ def initialize_database():
logger.error(f"Error checking schema health: {exc}") logger.error(f"Error checking schema health: {exc}")
raise RuntimeError("Error checking schema health") from exc raise RuntimeError("Error checking schema health") from exc
try: try:
database_service.run_migrations() database_service.run_migrations(fix=fix_migration)
except CommandError as exc: except CommandError as exc:
# if "overlaps with other requested revisions" or "Can't locate revision identified by" # if "overlaps with other requested revisions" or "Can't locate revision identified by"
# are not in the exception, we can't handle it # are not in the exception, we can't handle it