🔧 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 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.")