🔧 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:
parent
0c3d59a804
commit
c390a6ca1d
1 changed files with 44 additions and 22 deletions
|
|
@ -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.")
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue