From f7884af3afd175ecbe7c2b0c39d4966ce860e769 Mon Sep 17 00:00:00 2001
From: Gabriel Luiz Freitas Almeida <gabriel@logspace.ai>
Date: Tue, 3 Oct 2023 17:05:46 -0300
Subject: [PATCH] =?UTF-8?q?=F0=9F=90=9B=20fix(manager.py):=20migrate=20flo?=
 =?UTF-8?q?ws=20to=20default=20superuser=20if=20they=20don't=20have=20a=20?=
 =?UTF-8?q?user=20id=20associated=20with=20them=20when=20auto=5Flogin=20is?=
 =?UTF-8?q?=20enabled=20=F0=9F=90=9B=20fix(utils.py):=20handle=20case=20wh?=
 =?UTF-8?q?en=20superuser=20already=20exists=20during=20creation=20to=20pr?=
 =?UTF-8?q?event=20UNIQUE=20constraint=20violation?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 .../langflow/services/database/manager.py     | 19 +++++++++++++++++++
 src/backend/langflow/services/utils.py        | 15 ++++++++++++---
 2 files changed, 31 insertions(+), 3 deletions(-)

diff --git a/src/backend/langflow/services/database/manager.py b/src/backend/langflow/services/database/manager.py
index 3a388f694..0521ef583 100644
--- a/src/backend/langflow/services/database/manager.py
+++ b/src/backend/langflow/services/database/manager.py
@@ -58,6 +58,25 @@ class DatabaseService(Service):
         with Session(self.engine) as session:
             yield session
 
+    def migrate_flows_if_auto_login(self):
+        # if auto_login is enabled, we need to migrate the flows
+        # to the default superuser if they don't have a user id
+        # associated with them
+        settings_service = get_settings_service()
+        if settings_service.auth_settings.AUTO_LOGIN:
+            with Session(self.engine) as session:
+                flows = (
+                    session.query(models.Flow).filter(models.Flow.user_id == None).all()
+                )
+                if flows:
+                    logger.debug("Migrating flows to default superuser")
+                    username = settings_service.auth_settings.SUPERUSER
+                    user = get_user_by_username(session, username)
+                    for flow in flows:
+                        flow.user_id = user.id
+                    session.commit()
+                    logger.debug("Flows migrated successfully")
+
     def check_schema_health(self) -> bool:
         inspector = inspect(self.engine)
 
diff --git a/src/backend/langflow/services/utils.py b/src/backend/langflow/services/utils.py
index 20c23f8c0..86cfed92e 100644
--- a/src/backend/langflow/services/utils.py
+++ b/src/backend/langflow/services/utils.py
@@ -7,7 +7,7 @@ from langflow.services.settings.constants import (
     DEFAULT_SUPERUSER_PASSWORD,
 )
 from sqlmodel import Session
-from .getters import get_session, get_settings_service
+from .getters import get_db_service, get_session, get_settings_service
 from loguru import logger
 
 
@@ -54,8 +54,12 @@ def get_or_create_super_user(session: Session, username, password, is_default):
         logger.debug("Creating default superuser.")
     else:
         logger.debug("Creating superuser.")
-
-    return create_super_user(username, password, db=session)
+    try:
+        return create_super_user(username, password, db=session)
+    except Exception as exc:
+        if "UNIQUE constraint failed: user.username" in str(exc):
+            logger.debug("Superuser already exists.")
+            return None
 
 
 def setup_superuser(settings_service, session: Session):
@@ -193,3 +197,8 @@ def initialize_services():
     setup_superuser(
         service_manager.get(ServiceType.SETTINGS_SERVICE), next(get_session())
     )
+    try:
+        get_db_service().migrate_flows_if_auto_login()
+    except Exception as exc:
+        logger.error(f"Error migrating flows: {exc}")
+        raise RuntimeError("Error migrating flows") from exc