From b21ddd1c93e9e76c8289782123d6cb57319339a4 Mon Sep 17 00:00:00 2001
From: Gabriel Luiz Freitas Almeida <gabriel@logspace.ai>
Date: Tue, 3 Oct 2023 14:15:29 -0300
Subject: [PATCH] =?UTF-8?q?=F0=9F=94=A7=20fix(utils.py):=20refactor=20setu?=
 =?UTF-8?q?p=5Fsuperuser=20function=20to=20improve=20readability=20and=20m?=
 =?UTF-8?q?aintainability=20=E2=9C=A8=20feat(utils.py):=20add=20get=5For?=
 =?UTF-8?q?=5Fcreate=5Fsuper=5Fuser=20function=20to=20handle=20the=20creat?=
 =?UTF-8?q?ion=20of=20superuser=20and=20handle=20different=20scenarios=20b?=
 =?UTF-8?q?ased=20on=20existing=20users=20and=20credentials?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 src/backend/langflow/services/utils.py | 95 +++++++++++++++-----------
 1 file changed, 56 insertions(+), 39 deletions(-)

diff --git a/src/backend/langflow/services/utils.py b/src/backend/langflow/services/utils.py
index b85155d8a..6a6696faf 100644
--- a/src/backend/langflow/services/utils.py
+++ b/src/backend/langflow/services/utils.py
@@ -10,58 +10,75 @@ from .getters import get_session, get_settings_service
 from loguru import logger
 
 
+def get_or_create_super_user(session, username, password, is_default):
+    from langflow.services.database.models.user.user import User
+
+    user = session.query(User).filter(User.username == username).first()
+
+    if user and user.is_superuser and verify_password(password, user.password):
+        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(session, username, password)
+
+
 def setup_superuser(settings_service, session):
-    """
-    Setup the superuser.
-    """
-    # We will use the SUPERUSER and SUPERUSER_PASSWORD
-    # vars on settings_manager.auth_settings to create the superuser
-    # if it does not exist.
     if settings_service.auth_settings.AUTO_LOGIN:
         logger.debug("AUTO_LOGIN is set to True. Creating default superuser.")
 
     username = settings_service.auth_settings.SUPERUSER
     password = settings_service.auth_settings.SUPERUSER_PASSWORD
-    if username == DEFAULT_SUPERUSER and password == DEFAULT_SUPERUSER_PASSWORD:
-        logger.debug("Default superuser credentials detected.")
-        logger.debug("Creating default superuser.")
-    else:
-        logger.debug("Creating superuser.")
+
+    is_default = (username == DEFAULT_SUPERUSER) and (
+        password == DEFAULT_SUPERUSER_PASSWORD
+    )
 
     try:
-        from langflow.services.database.models.user.user import User
-
-        user = session.query(User).filter(User.username == username).first()
-        if user and user.is_superuser is True:
-            return
+        user = get_or_create_super_user(session, username, password, is_default)
+        if user is not None:
+            logger.debug("Superuser created successfully.")
     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:
-        if "UNIQUE constraint failed: user.username" in str(exc):
-            # check if the password is valid, if it is, we will
-            # continue normally
-            # if not, we will raise an error
-            user = session.query(User).filter(User.username == username).first()
-            if (
-                user
-                and user.is_superuser is True
-                and verify_password(password, user.password)
-            ):
-                return  # superuser already exists
-            else:
-                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.")
+    finally:
+        settings_service.auth_settings.reset_credentials()
 
 
 def teardown_superuser(settings_service, session):