Fixes super user creation (#996)

This commit is contained in:
Gabriel Luiz Freitas Almeida 2023-10-03 14:42:40 -03:00 • committed by GitHub
commit 7198197924
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23

View file

@ -1,4 +1,4 @@
from langflow.services.auth.utils import create_super_user from langflow.services.auth.utils import create_super_user, verify_password
from langflow.services.database.utils import initialize_database from langflow.services.database.utils import initialize_database
from langflow.services.manager import service_manager from langflow.services.manager import service_manager
from langflow.services.schema import ServiceType from langflow.services.schema import ServiceType
@ -6,50 +6,82 @@ from langflow.services.settings.constants import (
DEFAULT_SUPERUSER, DEFAULT_SUPERUSER,
DEFAULT_SUPERUSER_PASSWORD, DEFAULT_SUPERUSER_PASSWORD,
) )
from sqlmodel import Session
from .getters import get_session, get_settings_service from .getters import get_session, get_settings_service
from loguru import logger from loguru import logger
def setup_superuser(settings_service, session): def get_or_create_super_user(session: Session, username, password, is_default):
""" from langflow.services.database.models.user.user import User
Setup the superuser.
""" user = session.query(User).filter(User.username == username).first()
# We will use the SUPERUSER and SUPERUSER_PASSWORD
# vars on settings_manager.auth_settings to create the superuser if user and user.is_superuser and verify_password(password, user.password):
# if it does not exist. return None # Superuser already exists
if user and is_default:
if user.is_superuser:
if verify_password(password, user.password):
return None
else:
# Superuser exists but password is incorrect
# which means that the user has changed the
# base superuser credentials.
# This means that the user has already created
# a superuser and changed the password in the UI
# so we don't need to do anything.
logger.debug(
"Superuser exists but password is incorrect. "
"This means that the user has changed the "
"base superuser credentials."
)
return None
else:
logger.debug(
"User with superuser credentials exists but is not a superuser."
)
return None
if user:
if verify_password(password, user.password):
raise ValueError(
"User with superuser credentials exists but is not a superuser."
)
else:
raise ValueError("Incorrect superuser credentials")
if is_default:
logger.debug("Creating default superuser.")
else:
logger.debug("Creating superuser.")
return create_super_user(username, password, db=session)
def setup_superuser(settings_service, session: Session):
if settings_service.auth_settings.AUTO_LOGIN: if settings_service.auth_settings.AUTO_LOGIN:
logger.debug("AUTO_LOGIN is set to True. Creating default superuser.") logger.debug("AUTO_LOGIN is set to True. Creating default superuser.")
username = settings_service.auth_settings.SUPERUSER username = settings_service.auth_settings.SUPERUSER
password = settings_service.auth_settings.SUPERUSER_PASSWORD password = settings_service.auth_settings.SUPERUSER_PASSWORD
if username == DEFAULT_SUPERUSER and password == DEFAULT_SUPERUSER_PASSWORD:
logger.debug("Default superuser credentials detected.") is_default = (username == DEFAULT_SUPERUSER) and (
logger.debug("Creating default superuser.") password == DEFAULT_SUPERUSER_PASSWORD
else: )
logger.debug("Creating superuser.")
try: try:
from langflow.services.database.models.user.user import User user = get_or_create_super_user(
session=session, username=username, password=password, is_default=is_default
user = session.query(User).filter(User.username == username).first() )
if user and user.is_superuser is True: if user is not None:
return logger.debug("Superuser created successfully.")
except Exception as exc: except Exception as exc:
logger.exception(exc) logger.exception(exc)
raise RuntimeError( raise RuntimeError(
"Could not create superuser. Please create a superuser manually." "Could not create superuser. Please create a superuser manually."
) from exc ) from exc
try: finally:
# create superuser settings_service.auth_settings.reset_credentials()
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_service.auth_settings.reset_credentials()
logger.debug("Superuser created successfully.")
def teardown_superuser(settings_service, session): def teardown_superuser(settings_service, session):
@ -124,30 +156,40 @@ def initialize_services():
from langflow.services.task import factory as task_factory from langflow.services.task import factory as task_factory
from langflow.services.session import factory as session_service_factory # type: ignore from langflow.services.session import factory as session_service_factory # type: ignore
service_manager.register_factory(settings_factory.SettingsServiceFactory()) factory_and_dependencies = [
service_manager.register_factory( (settings_factory.SettingsServiceFactory(), []),
auth_factory.AuthServiceFactory(), dependencies=[ServiceType.SETTINGS_SERVICE] (
) auth_factory.AuthServiceFactory(),
service_manager.register_factory( [ServiceType.SETTINGS_SERVICE],
database_factory.DatabaseServiceFactory(), ),
dependencies=[ServiceType.SETTINGS_SERVICE], (
) database_factory.DatabaseServiceFactory(),
service_manager.register_factory( [ServiceType.SETTINGS_SERVICE],
cache_factory.CacheServiceFactory(), dependencies=[ServiceType.SETTINGS_SERVICE] ),
) (
service_manager.register_factory(chat_factory.ChatServiceFactory()) cache_factory.CacheServiceFactory(),
[ServiceType.SETTINGS_SERVICE],
),
(chat_factory.ChatServiceFactory(), []),
(task_factory.TaskServiceFactory(), []),
(
session_service_factory.SessionServiceFactory(),
[ServiceType.CACHE_SERVICE],
),
]
for factory, dependencies in factory_and_dependencies:
try:
service_manager.register_factory(factory, dependencies=dependencies)
except Exception as exc:
logger.exception(exc)
raise RuntimeError(
"Could not initialize services. Please check your settings."
) from exc
service_manager.register_factory(task_factory.TaskServiceFactory())
service_manager.register_factory(
session_service_factory.SessionServiceFactory(),
dependencies=[ServiceType.CACHE_SERVICE],
)
# Test cache connection # Test cache connection
service_manager.get(ServiceType.CACHE_SERVICE) service_manager.get(ServiceType.CACHE_SERVICE)
# Test database connection
service_manager.get(ServiceType.DATABASE_SERVICE)
# Setup the superuser # Setup the superuser
initialize_database() initialize_database()
session = next(get_session()) setup_superuser(
setup_superuser(service_manager.get(ServiceType.SETTINGS_SERVICE), session) service_manager.get(ServiceType.SETTINGS_SERVICE), next(get_session())
)