From 62325e000c3f79242192ca6ed14cf75613792075 Mon Sep 17 00:00:00 2001 From: Gabriel Luiz Freitas Almeida Date: Thu, 23 Nov 2023 16:03:19 -0300 Subject: [PATCH] Refactor database teardown process --- .../langflow/services/database/manager.py | 26 ++++++++----------- 1 file changed, 11 insertions(+), 15 deletions(-) diff --git a/src/backend/langflow/services/database/manager.py b/src/backend/langflow/services/database/manager.py index 7465c804f..6513ca280 100644 --- a/src/backend/langflow/services/database/manager.py +++ b/src/backend/langflow/services/database/manager.py @@ -1,17 +1,19 @@ from pathlib import Path from typing import TYPE_CHECKING + +import sqlalchemy as sa +from alembic import command +from alembic.config import Config 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.getters import get_settings_service -from sqlalchemy import inspect -import sqlalchemy as sa -from sqlalchemy.exc import OperationalError -from sqlmodel import SQLModel, Session, create_engine +from langflow.services.utils import teardown_superuser from loguru import logger -from alembic.config import Config -from alembic import command -from langflow.services.database import models # noqa +from sqlalchemy import inspect +from sqlalchemy.exc import OperationalError +from sqlmodel import Session, SQLModel, create_engine if TYPE_CHECKING: from sqlalchemy.engine import Engine @@ -225,14 +227,8 @@ class DatabaseService(Service): settings_service = get_settings_service() # remove the default superuser if auto_login is enabled # using the SUPERUSER to get the user - if settings_service.auth_settings.AUTO_LOGIN: - logger.debug("Removing default superuser") - username = settings_service.auth_settings.SUPERUSER - with Session(self.engine) as session: - user = get_user_by_username(session, username) - session.delete(user) - session.commit() - logger.debug("Default superuser removed") + with Session(self.engine) as session: + teardown_superuser(settings_service, session) except Exception as exc: logger.error(f"Error tearing down database: {exc}")