diff --git a/src/backend/langflow/services/utils.py b/src/backend/langflow/services/utils.py index 8b32aef02..4b16edc69 100644 --- a/src/backend/langflow/services/utils.py +++ b/src/backend/langflow/services/utils.py @@ -1,26 +1,48 @@ -from langflow.services import ServiceType, service_manager -from typing import TYPE_CHECKING, Generator +from langflow.services.auth.utils import create_super_user +from langflow.services.settings.constants import ( + DEFAULT_SUPERUSER, + DEFAULT_SUPERUSER_PASSWORD, +) +from .getters import get_session, get_settings_manager +from loguru import logger -if TYPE_CHECKING: - from langflow.services.database.manager import DatabaseManager - from langflow.services.settings.manager import SettingsManager - from langflow.services.chat.manager import ChatManager - from sqlmodel import Session +def setup_superuser(): + """ + Setup the superuser. + """ + # We will use the FIRST_SUPERUSER and FIRST_SUPERUSER_PASSWORD + # vars on settings_manager.auth_settings to create the superuser + # if it does not exist. + settings_manager = get_settings_manager() + session = next(get_session()) + username = settings_manager.auth_settings.SUPERUSER + password = settings_manager.auth_settings.SUPERUSER_PASSWORD + if username == DEFAULT_SUPERUSER and password == DEFAULT_SUPERUSER_PASSWORD: + logger.debug( + "Using default superuser credentials. Please change them in production." + ) + return + try: + from langflow.services.database.models.user.user import User -def get_settings_manager() -> "SettingsManager": - return service_manager.get(ServiceType.SETTINGS_MANAGER) - - -def get_db_manager() -> "DatabaseManager": - return service_manager.get(ServiceType.DATABASE_MANAGER) - - -def get_session() -> Generator["Session", None, None]: - db_manager = service_manager.get(ServiceType.DATABASE_MANAGER) - yield from db_manager.get_session() - - -def get_chat_manager() -> "ChatManager": - return service_manager.get(ServiceType.CHAT_MANAGER) + user = session.query(User).filter(User.username == username).first() + if user and user.is_superuser: + return + except Exception as exc: + logger.exception(exc) + raise RuntimeError( + "Could not create superuser. Please create a superuser manually." + ) from exc + try: + # create superuser + create_super_user(db=session, username=username, password=password) + except Exception as exc: + logger.exception(exc) + raise RuntimeError( + "Could not create superuser. Please create a superuser manually." + ) from exc + # reset superuser credentials + settings_manager.auth_settings.reset_credentials() + logger.debug("Superuser created successfully.")