fix: ensure unique flow names when assigning to superuser and avoid adding starter projects (#4751)
* fix: ensure unique flow names when assigning to superuser * refactor: Creates a new function with logic to rename flows * refactor: simplify function * Extract STARTER_FOLDER constants to a separate module * Filter orphaned flows by excluding those in the starter folder --------- Co-authored-by: Gabriel Luiz Freitas Almeida <gabriel@langflow.org>
This commit is contained in:
Ítalo Johnny
GitHub
parent
12407f00af
commit
e728e40561
7 changed files with 76 additions and 51 deletions
|
|
@ -25,7 +25,7 @@ from langflow.api.utils import (
|
|||
validate_is_component,
|
||||
)
|
||||
from langflow.api.v1.schemas import FlowListCreate
|
||||
from langflow.initial_setup.setup import STARTER_FOLDER_NAME
|
||||
from langflow.initial_setup.constants import STARTER_FOLDER_NAME
|
||||
from langflow.services.database.models.flow import Flow, FlowCreate, FlowRead, FlowUpdate
|
||||
from langflow.services.database.models.flow.model import FlowHeader
|
||||
from langflow.services.database.models.flow.utils import get_webhook_component_in_flow
|
||||
|
|
|
|||
|
|
@ -19,7 +19,7 @@ from langflow.api.v1.flows import create_flows
|
|||
from langflow.api.v1.schemas import FlowListCreate
|
||||
from langflow.helpers.flow import generate_unique_flow_name
|
||||
from langflow.helpers.folders import generate_unique_folder_name
|
||||
from langflow.initial_setup.setup import STARTER_FOLDER_NAME
|
||||
from langflow.initial_setup.constants import STARTER_FOLDER_NAME
|
||||
from langflow.services.database.models.flow.model import Flow, FlowCreate, FlowRead
|
||||
from langflow.services.database.models.folder.constants import DEFAULT_FOLDER_NAME
|
||||
from langflow.services.database.models.folder.model import (
|
||||
|
|
|
|||
2
src/backend/base/langflow/initial_setup/constants.py
Normal file
2
src/backend/base/langflow/initial_setup/constants.py
Normal file
|
|
@ -0,0 +1,2 @@
|
|||
STARTER_FOLDER_NAME = "Starter Projects"
|
||||
STARTER_FOLDER_DESCRIPTION = "Starter projects to help you get started in Langflow."
|
||||
|
|
@ -14,11 +14,8 @@ from loguru import logger
|
|||
from sqlalchemy.exc import NoResultFound
|
||||
from sqlmodel import select
|
||||
|
||||
from langflow.base.constants import (
|
||||
FIELD_FORMAT_ATTRIBUTES,
|
||||
NODE_FORMAT_ATTRIBUTES,
|
||||
ORJSON_OPTIONS,
|
||||
)
|
||||
from langflow.base.constants import FIELD_FORMAT_ATTRIBUTES, NODE_FORMAT_ATTRIBUTES, ORJSON_OPTIONS
|
||||
from langflow.initial_setup.constants import STARTER_FOLDER_DESCRIPTION, STARTER_FOLDER_NAME
|
||||
from langflow.services.auth.utils import create_super_user
|
||||
from langflow.services.database.models.flow.model import Flow, FlowCreate
|
||||
from langflow.services.database.models.folder.model import Folder, FolderCreate
|
||||
|
|
@ -37,9 +34,6 @@ from langflow.services.deps import (
|
|||
from langflow.template.field.prompt import DEFAULT_PROMPT_INTUT_TYPES
|
||||
from langflow.utils.util import escape_json_dump
|
||||
|
||||
STARTER_FOLDER_NAME = "Starter Projects"
|
||||
STARTER_FOLDER_DESCRIPTION = "Starter projects to help you get started in Langflow."
|
||||
|
||||
# In the folder ./starter_projects we have a few JSON files that represent
|
||||
# starter projects. We want to load these into the database so that users
|
||||
# can use them as a starting point for their own projects.
|
||||
|
|
|
|||
|
|
@ -21,13 +21,11 @@ from sqlalchemy.ext.asyncio import AsyncEngine, create_async_engine
|
|||
from sqlmodel import Session, SQLModel, create_engine, select, text
|
||||
from sqlmodel.ext.asyncio.session import AsyncSession
|
||||
|
||||
from langflow.initial_setup.constants import STARTER_FOLDER_NAME
|
||||
from langflow.services.base import Service
|
||||
from langflow.services.database import models
|
||||
from langflow.services.database.models.user.crud import get_user_by_username
|
||||
from langflow.services.database.utils import (
|
||||
Result,
|
||||
TableResults,
|
||||
)
|
||||
from langflow.services.database.utils import Result, TableResults
|
||||
from langflow.services.deps import get_settings_service
|
||||
from langflow.services.utils import teardown_superuser
|
||||
|
||||
|
|
@ -147,44 +145,77 @@ class DatabaseService(Service):
|
|||
yield session
|
||||
|
||||
async def assign_orphaned_flows_to_superuser(self) -> None:
|
||||
"""Assign flows without a user ID to the default superuser if auto_login is enabled."""
|
||||
# Get the settings service to check if auto_login is enabled
|
||||
"""Assign orphaned flows to the default superuser when auto login is enabled."""
|
||||
settings_service = get_settings_service()
|
||||
if settings_service.auth_settings.AUTO_LOGIN:
|
||||
async with self.with_async_session() as session:
|
||||
# Select flows where user_id is NULL (orphaned flows)
|
||||
stmt = select(models.Flow).where(models.Flow.user_id == None) # noqa: E711
|
||||
orphaned_flows = (await session.exec(stmt)).all()
|
||||
|
||||
if orphaned_flows:
|
||||
logger.debug("Assigning orphaned flows to the default superuser")
|
||||
if not settings_service.auth_settings.AUTO_LOGIN:
|
||||
return
|
||||
|
||||
# Get the default superuser
|
||||
superuser_username = settings_service.auth_settings.SUPERUSER
|
||||
superuser = await get_user_by_username(session, superuser_username)
|
||||
async with self.with_async_session() as session:
|
||||
# Fetch orphaned flows
|
||||
stmt = (
|
||||
select(models.Flow)
|
||||
.join(models.Folder)
|
||||
.where(
|
||||
models.Flow.user_id == None, # noqa: E711
|
||||
models.Folder.name != STARTER_FOLDER_NAME,
|
||||
)
|
||||
)
|
||||
orphaned_flows = (await session.exec(stmt)).all()
|
||||
|
||||
if not superuser:
|
||||
logger.error("Default superuser not found")
|
||||
msg = "Default superuser not found"
|
||||
raise RuntimeError(msg)
|
||||
if not orphaned_flows:
|
||||
return
|
||||
|
||||
stmt = select(models.Flow.name).where(models.Flow.user_id == superuser.id)
|
||||
result = await session.exec(stmt)
|
||||
superuser_flows_names = result.all()
|
||||
# Assign each orphaned flow to the superuser
|
||||
for flow in orphaned_flows:
|
||||
flow.user_id = superuser.id
|
||||
if flow.name in superuser_flows_names:
|
||||
name_match = re.search(r"\((\d+)\)$", flow.name)
|
||||
if not name_match:
|
||||
flow.name = f"{flow.name} (1)"
|
||||
else:
|
||||
num = int(name_match.group(1)) + 1
|
||||
flow.name = re.sub(r"\(\d+\)$", f"({num})", flow.name)
|
||||
logger.debug("Assigning orphaned flows to the default superuser")
|
||||
|
||||
# Commit the changes to the database
|
||||
await session.commit()
|
||||
logger.debug("Successfully assigned orphaned flows to the default superuser")
|
||||
# Retrieve superuser
|
||||
superuser_username = settings_service.auth_settings.SUPERUSER
|
||||
superuser = await get_user_by_username(session, superuser_username)
|
||||
|
||||
if not superuser:
|
||||
error_message = "Default superuser not found"
|
||||
logger.error(error_message)
|
||||
raise RuntimeError(error_message)
|
||||
|
||||
# Get existing flow names for the superuser
|
||||
existing_names: set[str] = set(
|
||||
(await session.exec(select(models.Flow.name).where(models.Flow.user_id == superuser.id))).all()
|
||||
)
|
||||
|
||||
# Process orphaned flows
|
||||
for flow in orphaned_flows:
|
||||
flow.user_id = superuser.id
|
||||
flow.name = self._generate_unique_flow_name(flow.name, existing_names)
|
||||
existing_names.add(flow.name)
|
||||
|
||||
# Commit changes
|
||||
await session.commit()
|
||||
logger.debug("Successfully assigned orphaned flows to the default superuser")
|
||||
|
||||
def _generate_unique_flow_name(self, original_name: str, existing_names: set[str]) -> str:
|
||||
"""Generate a unique flow name by adding or incrementing a suffix."""
|
||||
if original_name not in existing_names:
|
||||
return original_name
|
||||
|
||||
match = re.search(r"^(.*) \((\d+)\)$", original_name)
|
||||
if match:
|
||||
base_name, current_number = match.groups()
|
||||
new_name = f"{base_name} ({int(current_number) + 1})"
|
||||
else:
|
||||
new_name = f"{original_name} (1)"
|
||||
|
||||
# Ensure unique name by incrementing suffix
|
||||
while new_name in existing_names:
|
||||
match = re.match(r"^(.*) \((\d+)\)$", new_name)
|
||||
if match is not None:
|
||||
base_name, current_number = match.groups()
|
||||
else:
|
||||
error_message = "Invalid format: match is None"
|
||||
raise ValueError(error_message)
|
||||
|
||||
new_name = f"{base_name} ({int(current_number) + 1})"
|
||||
|
||||
return new_name
|
||||
|
||||
def check_schema_health(self) -> bool:
|
||||
inspector = inspect(self.engine)
|
||||
|
|
|
|||
|
|
@ -17,7 +17,7 @@ from dotenv import load_dotenv
|
|||
from fastapi.testclient import TestClient
|
||||
from httpx import ASGITransport, AsyncClient
|
||||
from langflow.graph import Graph
|
||||
from langflow.initial_setup.setup import STARTER_FOLDER_NAME
|
||||
from langflow.initial_setup.constants import STARTER_FOLDER_NAME
|
||||
from langflow.services.auth.utils import get_password_hash
|
||||
from langflow.services.database.models.api_key.model import ApiKey
|
||||
from langflow.services.database.models.flow.model import Flow, FlowCreate
|
||||
|
|
|
|||
|
|
@ -3,11 +3,9 @@ from datetime import datetime
|
|||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
from langflow.custom.directory_reader.utils import (
|
||||
abuild_custom_component_list_from_path,
|
||||
)
|
||||
from langflow.custom.directory_reader.utils import abuild_custom_component_list_from_path
|
||||
from langflow.initial_setup.constants import STARTER_FOLDER_NAME
|
||||
from langflow.initial_setup.setup import (
|
||||
STARTER_FOLDER_NAME,
|
||||
get_project_data,
|
||||
load_starter_projects,
|
||||
update_projects_components_with_latest_component_versions,
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue