🔧 chore(utils.py): refactor get_settings_manager, get_db_manager, get_session, and get_chat_manager functions to use service_manager for better code organization and maintainability

✨ feat(utils.py): add setup_superuser function to create a superuser if it does not exist and reset superuser credentials after creation
This commit is contained in:
Gabriel Luiz Freitas Almeida 2023-09-18 15:43:18 -03:00
commit c390a6ca1d

View file

@ -1,26 +1,48 @@
from langflow.services import ServiceType, service_manager from langflow.services.auth.utils import create_super_user
from typing import TYPE_CHECKING, Generator 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: def setup_superuser():
from langflow.services.database.manager import DatabaseManager """
from langflow.services.settings.manager import SettingsManager Setup the superuser.
from langflow.services.chat.manager import ChatManager """
from sqlmodel import Session # 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": user = session.query(User).filter(User.username == username).first()
return service_manager.get(ServiceType.SETTINGS_MANAGER) if user and user.is_superuser:
return
except Exception as exc:
def get_db_manager() -> "DatabaseManager": logger.exception(exc)
return service_manager.get(ServiceType.DATABASE_MANAGER) raise RuntimeError(
"Could not create superuser. Please create a superuser manually."
) from exc
def get_session() -> Generator["Session", None, None]: try:
db_manager = service_manager.get(ServiceType.DATABASE_MANAGER) # create superuser
yield from db_manager.get_session() create_super_user(db=session, username=username, password=password)
except Exception as exc:
logger.exception(exc)
def get_chat_manager() -> "ChatManager": raise RuntimeError(
return service_manager.get(ServiceType.CHAT_MANAGER) "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.")