From c390a6ca1d3f3a6204df489fe7ba2c71c9239ced Mon Sep 17 00:00:00 2001
From: Gabriel Luiz Freitas Almeida <gabriel@logspace.ai>
Date: Mon, 18 Sep 2023 15:43:18 -0300
Subject: [PATCH] =?UTF-8?q?=F0=9F=94=A7=20chore(utils.py):=20refactor=20ge?=
 =?UTF-8?q?t=5Fsettings=5Fmanager,=20get=5Fdb=5Fmanager,=20get=5Fsession,?=
 =?UTF-8?q?=20and=20get=5Fchat=5Fmanager=20functions=20to=20use=20service?=
 =?UTF-8?q?=5Fmanager=20for=20better=20code=20organization=20and=20maintai?=
 =?UTF-8?q?nability=20=E2=9C=A8=20feat(utils.py):=20add=20setup=5Fsuperuse?=
 =?UTF-8?q?r=20function=20to=20create=20a=20superuser=20if=20it=20does=20n?=
 =?UTF-8?q?ot=20exist=20and=20reset=20superuser=20credentials=20after=20cr?=
 =?UTF-8?q?eation?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 src/backend/langflow/services/utils.py | 66 +++++++++++++++++---------
 1 file changed, 44 insertions(+), 22 deletions(-)

diff --git a/src/backend/langflow/services/utils.py b/src/backend/langflow/services/utils.py
index 8b32aef02..4b16edc69 100644
--- a/src/backend/langflow/services/utils.py
+++ b/src/backend/langflow/services/utils.py
@@ -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.")