feat: Publish Flow, API code update and UI components (#6140)
* refactor: Update flowToolbarComponent with FlowToolbarOptions component and ENABLE_PUBLISH feature flag * [autofix.ci] apply automated fixes * refactor: Update ENABLE_PUBLISH feature flag in feature-flags.ts * add external link for hover state * Refactor routes.tsx to enable PlaygroundPage * Refactor deploy-dropdown.tsx to add ShadTooltipComponent and handle cases where there is no IO * [autofix.ci] apply automated fixes * add colorfull langflow icon * Add playgroundPage prop to IOModalPropsType * Refactor IOModal in PlaygroundPage component * Refactor IOModal component and add publish options * Add LangflowButtonRedirectTarget utility function * Refactor IOModal component and add LangflowButtonRedirectTarget utility function * fix: remove feature flag for playground button name * fix: rename DeployDropdown to PublishDropdown in FlowToolbarOptions * fix: rename DeployDropdown to PublishDropdown and update related functionality * fix: update classNames utility import and refactor class assignment in FlowToolbar * [autofix.ci] apply automated fixes * fix: enhance hover effects and accessibility in PublishDropdown component * [autofix.ci] apply automated fixes * fix: update Playground title in IOModal and ChatViewWrapper components * fix: improve layout and visibility of session information in ChatViewWrapper component * add neutral icon to playground * fix: add playgroundPage prop to ContentBlockDisplay and conditionally render elements in ChatMessage * fix: pass playgroundPage prop to ContentDisplay and conditionally render duration * fix: remove playgroundTitle display from ChatViewWrapper component * fix: adjust padding and alignment in ChatViewWrapper and IOModal components based on playgroundPage prop * fix: update alignment and responsiveness in ChatViewWrapper component based on playgroundPage and sidebarOpen states * fix: update document title based on currentSavedFlow in PlaygroundPage component * [autofix.ci] apply automated fixes * feat: add ENABLE_WIDGET flag to conditionally render embed option in PublishDropdown * feat: add EmbedModal component for copying embed code * feat: integrate EmbedModal in PublishDropdown for embed code sharing * feat: enhance EmbedModal integration in PublishDropdown with dynamic embed code generation * feat: add switch for publishing state in PublishDropdown component and update FlowType * [autofix.ci] apply automated fixes * [autofix.ci] apply automated fixes * feat: add optional copy property to GetCodeType for enhanced tweak management * feat: make description property optional in modalHeaderType for improved flexibility * feat: update background color in code styles for improved visibility * feat: add support for 'retangular' size in switchCaseModalSize helper * feat: make description property optional in Header component and add 'retangular' size option in BaseModalProps * feat: add optional copy parameter to getWidgetCode for customizable script output * feat: refactor EmbedModal props for improved clarity and consistency * feat: add ApiModal component for custom API code generation * feat: adjust icon size in EmbedModal for better visual consistency * feat: integrate ApiModal into PublishDropdown for enhanced API code generation * feat: update ApiModal header to include API key creation instructions * feat: add APITabsComponent for displaying code snippets in tabs * feat: add APITabType and tabsArrayType for enhanced API tab management * feat: replace CodeTabsComponent with APITabsComponent in ApiModal for improved tab management * feat: add JSIcon component for JavaScript icon representation * feat: add JSIcon to nodeIconsLucide for JavaScript representation * style: format code for consistency in EmbedModal component * feat: change ApiModal size from x-large to medium for better usability * feat: enhance APITabsComponent with light theme support and improve button positioning * feat: adjust APITabsComponent tab styling for improved layout and readability * feat: add BWPython icon and update references in APITabsComponent and styleUtils * fix: correct export name for BWSvgPython icon component * feat: add new column 'public' to flow table * raw: Sample code for creating a new unauthenticated endpoint as a PoCh Note: this code is blatantly copied from another code snippet. * feat: add dark mode support to JSIcon component * feat: add dark mode support to BWPython icon component * feat: add new column 'access_type' to flow * chore: update code * chore: remove unnecessary migration * feat: add getNewCurlCode function for generating cURL commands with dynamic payloads * feat: integrate dynamic code generation for Python, JavaScript, and cURL in APITabsComponent * feat: add getNewJsApiCode function for generating JavaScript API code with dynamic payloads * feat: add getNewPythonApiCode function for generating Python API code with dynamic payloads * fix color on sintax highlight * feat: enhance APITabsComponent with dynamic streaming and authentication state * feat: update APITabsComponent to handle dynamic input and output types based on flow data * [autofix.ci] apply automated fixes * feat: add input and output validation in PlaygroundPage to redirect if none exist * [autofix.ci] apply automated fixes * fix: ensure navigation only occurs when currentSavedFlow data is present * feat: add access_type field to FlowUpdate model * feat: add playgroundPage parameter to buildFlowVertices for conditional URL construction * feat: add playgroundPage state and setter to FlowStoreType * feat: add playgroundPage state and setter to FlowStoreType * feat: add access_type field to FlowType for improved access control * feat: modify useGetMessagesQuery to handle playgroundPage state for conditional message retrieval * feat: update IPatchUpdateFlow interface to make fields optional and add access_type for enhanced flexibility * feat: implement publish toggle functionality in PublishDropdown component for dynamic access control * feat: integrate playgroundPage state management in IOModal for improved session handling * add: new endpoint to public_flow * refactor: remove unused current_user parameter from read_public_flow function * feat: add ContextWrapper to PlaygroundPage route for improved context management * refactor: simplify flow retrieval logic in PlaygroundPage component * feat: add support for public flow retrieval in useGetFlow hook * feat: add PUBLIC_FLOW constant to URLs for public flow retrieval * fix: add whitespace for improved code readability in AuthSettingsGuard component * fix: add whitespace for improved code readability in ProtectedAdminRoute component * [autofix.ci] apply automated fixes * fix: update redirect condition in PlaygroundPage for non-public access types * [autofix.ci] apply automated fixes * persist session name update * fix: enhance message update logic to handle playground state and local storage * fix: remove debugger statement from PlaygroundPage initialization * fix: manage dark mode class in App component and remove redundant logic from AppInitPage * [autofix.ci] apply automated fixes * feat: add access_type field to FlowHeader model for flow access control * fix: refactor flow access handling in PublishDropdown component for improved readability and async operation * feat: enhance FlowMenu component with swatch color display based on flow gradient * feat: add swatch color display in IOModal based on flow gradient * [autofix.ci] apply automated fixes * Update copyCode to use dynamic API code generation for Python, JavaScript, and cURL tabs * Add optional session tracking to JavaScript API code generation * Enhance Python API code generation with detailed comments and error handling * Fix SVG attribute casing and update dark mode state handling in Python icon components * Fix SVG clip-path casing and ensure dark mode state is a string in JS icon components * Fix SVG fill color handling for dark mode in Python icon components * Fix SVG filter handling for dark mode in JS icon components * Add tweaks management and update functionality in tweaks store * Add normal font style to line numbers in CSS * Add debug log for flow retrieval in FlowPage component * Refactor APITabsComponent to remove unused props and integrate tweaks management * Enhance ApiModal to support tweaks management and improve API access UI * Remove flow prop from ApiModal in PublishDropdown component * update package lock * [autofix.ci] apply automated fixes * Update ChatViewWrapper to adjust layout based on visibleSession state * [autofix.ci] apply automated fixes * Update icon fallback in FlowMenu to use "Workflow" * Refactor EmbedModal button styles for consistency and clarity * Add useEffect to reset copied state on active tab change and clean up button styles * Comment out DropdownMenuItem in PublishDropdown for future reference * refactor: remove duplicated code from route * refactor: remove duplicated code from route * [autofix.ci] apply automated fixes * Increase font size for code blocks in classes.css for better readability * Adjust padding and height for deploy dropdown items for improved layout * Update minWidth for medium modal size to include max-width constraint * Refactor ApiModal to conditionally render button and adjust styles for improved layout * Add margin-top to API modal tabs content for improved spacing * Refactor deploy dropdown to include API access and Embed options, and rename 'Standalone app' to 'Shareable Playground' * Enhance API code generation to include environment variable checks for API key in curl, JavaScript, and Python examples * Fix authentication check logic and clean up modal class names * Refactor authentication logic in APITabsComponent to improve clarity and functionality * [autofix.ci] apply automated fixes * Update environment variable references in API code examples to use LANGFLOW_API_KEY * [autofix.ci] apply automated fixes * Update API key references to use LANGFLOW_API_KEY in curl and JS code examples * Remove streaming parameter from API code examples in JavaScript and Python * Adjust button padding and separator margin in API modal for improved layout * Add transparent background to scrollbar corner in Tailwind config * Update publish dropdown to display sharing status based on flow publication state * Add playgroundPage prop to ChatInput and conditionally render file upload button * Refactor ChatViewWrapper layout logic for improved responsiveness * [autofix.ci] apply automated fixes * Add closeButtonClassName prop to BaseModal and DialogContent for customization * [autofix.ci] apply automated fixes * [autofix.ci] apply automated fixes * feat: use new deterministic flow id for public flow calls (#6314) * Enhance public flow building with deterministic flow ID and name handling * Handle asyncio.CancelledError in message table commit operation * Add UUID v5 generation for flow IDs in chat view and modal components * Add comment explaining CancelledError handling in message table commit * Simplify public flow ID and name generation logic * Refactor UUID v5 import and flow ID generation in IOModal components * Update Alembic migration down_revision identifier --------- Co-authored-by: anovazzi1 <otavio2204@gmail.com> * feat: add playgroundPage prop to SessionSelector and SidebarOpenView components * feat: switch from localStorage to sessionStorage for message handling in playground mode * feat: Add public flow settings with cleanup and expiration configurations Add new settings for managing public temporary flows, including: - Configurable cleanup interval (default 1 hour) - Configurable flow expiration time (default 24 hours) - Minimum interval of 10 minutes for both settings * feat: Add public flow expiration settings to ConfigResponse Update ConfigResponse schema to include new configuration parameters for public flow management: - public_flow_cleanup_interval: Interval for cleaning up public flows - public_flow_expiration: Duration for public flow retention * feat: Add temporary public flow cleanup worker Implement a background worker to manage and clean up expired public flows: - Add CleanupWorker class to handle periodic cleanup tasks - Integrate cleanup worker into application lifespan - Implement cleanup logic for removing expired public flow data from database and storage - Add start and stop methods for graceful worker management * feat: implement client ID management using cookies in Playground and update flow ID generation * refactor: Optimize public flow cleanup worker with targeted execution Improve the temporary public flow cleanup process by: - Adding a pre-check to only run cleanup when public flows exist - Passing public flows and session directly to cleanup function - Simplifying the cleanup logic to reduce nested session management * refactor: Enhance database cleanup worker with comprehensive record management Improve the cleanup worker to handle both expired public flows and orphaned records: - Add function to clean up expired public flows with detailed logging - Implement orphaned record cleanup across multiple database tables - Enhance error handling and logging for storage file deletion - Simplify worker run method to execute both cleanup tasks sequentially * feat: Add session cookie validation for public flow generation Enhance public flow building by: - Requiring a session cookie for generating temporary public flows - Incorporating the session cookie into the flow ID generation process - Adding explicit error handling for missing session cookies * fix: Update session cookie retrieval in public flow generation Change cookie key from "session" to "client_id" to align with recent client ID management implementation * fix: Correct flow ID generation by adding an underscore separator between client ID and real flow ID * [autofix.ci] apply automated fixes * fix: add options to the fetch call and add docs Update getNewJsApiCode function to: - Add comprehensive JSDoc documentation - Include fetch options in API call - Simplify code generation logic - Ensure proper payload and options handling * fix: update label for temporary overrides to tweaks in API modal * update package lock * ensure individual instances of contexts * fix: add data-testid attributes for testing in PublishDropdown component * fix: handle authentication errors for public API requests in ApiInterceptor * test: add publish feature test using Playwright * [autofix.ci] apply automated fixes * refactor: optimize temp flow cleanup with improved file and logging management * test: add unit tests for temp flow cleanup service * chore: remove unnecessary console logs and comments for cleaner code * Update src/frontend/tests/core/features/publish-flow.spec.ts Co-authored-by: Cristhian Zanforlin Lousa <cristhian.lousa@gmail.com> * Update src/frontend/tests/core/features/publish-flow.spec.ts Co-authored-by: Cristhian Zanforlin Lousa <cristhian.lousa@gmail.com> * Update src/frontend/tests/core/features/publish-flow.spec.ts Co-authored-by: Cristhian Zanforlin Lousa <cristhian.lousa@gmail.com> * Update src/frontend/tests/core/features/publish-flow.spec.ts Co-authored-by: Cristhian Zanforlin Lousa <cristhian.lousa@gmail.com> * Update src/frontend/tests/core/features/publish-flow.spec.ts Co-authored-by: Cristhian Zanforlin Lousa <cristhian.lousa@gmail.com> * Update src/frontend/tests/core/features/publish-flow.spec.ts Co-authored-by: Cristhian Zanforlin Lousa <cristhian.lousa@gmail.com> * Update src/frontend/tests/core/features/publish-flow.spec.ts Co-authored-by: Cristhian Zanforlin Lousa <cristhian.lousa@gmail.com> * Update src/frontend/src/modals/apiModal/utils/get-python-api-code.tsx Co-authored-by: Cristhian Zanforlin Lousa <cristhian.lousa@gmail.com> * Update src/frontend/tests/core/features/publish-flow.spec.ts Co-authored-by: Cristhian Zanforlin Lousa <cristhian.lousa@gmail.com> * Update src/frontend/tests/core/features/publish-flow.spec.ts Co-authored-by: Cristhian Zanforlin Lousa <cristhian.lousa@gmail.com> * Update src/frontend/src/modals/apiModal/utils/get-curl-code.tsx Co-authored-by: Cristhian Zanforlin Lousa <cristhian.lousa@gmail.com> * Update src/frontend/src/modals/apiModal/utils/get-curl-code.tsx Co-authored-by: Cristhian Zanforlin Lousa <cristhian.lousa@gmail.com> * Update src/frontend/src/modals/apiModal/utils/get-curl-code.tsx Co-authored-by: Cristhian Zanforlin Lousa <cristhian.lousa@gmail.com> * Update src/frontend/src/modals/apiModal/utils/get-python-api-code.tsx Co-authored-by: Cristhian Zanforlin Lousa <cristhian.lousa@gmail.com> * Update src/frontend/src/modals/apiModal/utils/get-python-api-code.tsx Co-authored-by: Cristhian Zanforlin Lousa <cristhian.lousa@gmail.com> * fix: adds edge filtering only to parameters not hidden (#6270) Adds advanced filter to only filter by showing fields on click of edge * Add config parameter to Graph run method with configuration application * fix: remove unnecessary await from flow directory existence check * refactor: Remove expired public flows cleanup method * refactor: Move data_dir initialization to base StorageService * fix: Update temp flow cleanup to use async file operations - Add explicit type hint for storage_service - Use async methods for file and directory existence checks - Improve error handling for file deletion during flow cleanup * refactor: Add explicit type hint for tables in temp flow cleanup - Improve type annotations for tables list in cleanup_orphaned_records - Explicitly define the types of tables to be processed * [autofix.ci] apply automated fixes * refactor: Remove tests for expired public flows from temp flow cleanup * test: Update test_cleanup_orphaned_records_no_orphans to use fixtures * feat: Add utility function for verifying public flow access Implement async function to validate public flow requests with: - Client ID verification - Flow existence and public access check - Deterministic flow ID generation - User retrieval for permission handling * feat: build_public_tmp to use the jobqueue Refactored the build_public_tmp endpoint to: - Add comprehensive docstring explaining endpoint functionality - Improve error handling and logging - Simplify flow verification and user retrieval process - Use new verify_public_flow_and_get_user utility function - Streamline job creation and error management * chore: Add anyio import to local storage service Import anyio library in preparation for potential async storage operations * style: run formatter * changed endpoint * [autofix.ci] apply automated fixes * Add size for tweaks * Add size for tweaks * Change tweaks modal * Fix switch design * [autofix.ci] apply automated fixes * fix: mypy erros * fix: alembic multiple heads error * fix: ruff error * refactor: update test cleanup for orphaned records to use fixtures Changed the test for cleanup of orphaned records to utilize the "client" fixture instead of the "asyncio" marker, enhancing test organization and clarity. * fix: cli test --------- Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com> Co-authored-by: italojohnny <italojohnnydosanjos@gmail.com> Co-authored-by: Gabriel Luiz Freitas Almeida <gabriel@langflow.org> Co-authored-by: Cristhian Zanforlin Lousa <cristhian.lousa@gmail.com> Co-authored-by: Lucas Oliveira <62335616+lucaseduoli@users.noreply.github.com> Co-authored-by: Lucas Oliveira <lucas.edu.oli@hotmail.com>
This commit is contained in:
anovazzi1
GitHub
parent
63a2e3349f
commit
7aca264fec
86 changed files with 2252 additions and 1465 deletions
|
|
@ -0,0 +1,33 @@
|
|||
"""add column 'access_type' to flow
|
||||
|
||||
Revision ID: f3b2d1f1002d
|
||||
Revises: 93e2705fa8d6
|
||||
Create Date: 2025-02-05 14:35:29.658101
|
||||
|
||||
"""
|
||||
from typing import Sequence, Union
|
||||
|
||||
from alembic import op
|
||||
import sqlalchemy as sa
|
||||
from langflow.utils import migration
|
||||
|
||||
|
||||
# revision identifiers, used by Alembic.
|
||||
revision: str = 'f3b2d1f1002d'
|
||||
down_revision: Union[str, None] = '93e2705fa8d6'
|
||||
branch_labels: Union[str, Sequence[str], None] = None
|
||||
depends_on: Union[str, Sequence[str], None] = None
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
conn = op.get_bind()
|
||||
with op.batch_alter_table('flow', schema=None) as batch_op:
|
||||
if not migration.column_exists(table_name='flow', column_name='access_type', conn=conn):
|
||||
batch_op.add_column(sa.Column('access_type', sa.Enum('PRIVATE', 'PUBLIC', name='access_type_enum'), server_default='private', nullable=False))
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
conn = op.get_bind()
|
||||
with op.batch_alter_table('flow', schema=None) as batch_op:
|
||||
if migration.column_exists(table_name='flow', column_name='access_type', conn=conn):
|
||||
batch_op.drop_column('access_type')
|
||||
|
|
@ -50,6 +50,7 @@ async def start_flow_build(
|
|||
log_builds: bool,
|
||||
current_user: CurrentActiveUser,
|
||||
queue_service: JobQueueService,
|
||||
flow_name: str | None = None,
|
||||
) -> str:
|
||||
"""Start the flow build process by setting up the queue and starting the build task.
|
||||
|
||||
|
|
@ -70,6 +71,7 @@ async def start_flow_build(
|
|||
start_component_id=start_component_id,
|
||||
log_builds=log_builds,
|
||||
current_user=current_user,
|
||||
flow_name=flow_name,
|
||||
)
|
||||
queue_service.start_job(job_id, task_coro)
|
||||
except Exception as e:
|
||||
|
|
@ -154,6 +156,7 @@ async def generate_flow_events(
|
|||
start_component_id: str | None,
|
||||
log_builds: bool,
|
||||
current_user: CurrentActiveUser,
|
||||
flow_name: str | None = None,
|
||||
) -> None:
|
||||
"""Generate events for flow building process.
|
||||
|
||||
|
|
@ -175,7 +178,7 @@ async def generate_flow_events(
|
|||
flow_id_str = str(flow_id)
|
||||
# Create a fresh session for database operations
|
||||
async with session_scope() as fresh_session:
|
||||
graph = await create_graph(fresh_session, flow_id_str)
|
||||
graph = await create_graph(fresh_session, flow_id_str, flow_name)
|
||||
|
||||
graph.validate_stream()
|
||||
first_layer = sort_vertices(graph)
|
||||
|
|
@ -214,7 +217,7 @@ async def generate_flow_events(
|
|||
),
|
||||
)
|
||||
|
||||
async def create_graph(fresh_session, flow_id_str: str) -> Graph:
|
||||
async def create_graph(fresh_session, flow_id_str: str, flow_name: str | None) -> Graph:
|
||||
if inputs is not None and getattr(inputs, "session", None) is not None:
|
||||
effective_session_id = inputs.session
|
||||
else:
|
||||
|
|
@ -229,8 +232,9 @@ async def generate_flow_events(
|
|||
session_id=effective_session_id,
|
||||
)
|
||||
|
||||
result = await fresh_session.exec(select(Flow.name).where(Flow.id == flow_id))
|
||||
flow_name = result.first()
|
||||
if not flow_name:
|
||||
result = await fresh_session.exec(select(Flow.name).where(Flow.id == flow_id))
|
||||
flow_name = result.first()
|
||||
|
||||
return await build_graph_from_data(
|
||||
flow_id=flow_id_str,
|
||||
|
|
|
|||
|
|
@ -303,3 +303,63 @@ def custom_params(
|
|||
if page is None and size is None:
|
||||
return None
|
||||
return Params(page=page or MIN_PAGE_SIZE, size=size or MAX_PAGE_SIZE)
|
||||
|
||||
|
||||
async def verify_public_flow_and_get_user(flow_id: uuid.UUID, client_id: str | None) -> tuple[User, uuid.UUID]:
|
||||
"""Verify a public flow request and generate a deterministic flow ID.
|
||||
|
||||
This utility function:
|
||||
1. Checks that a client_id cookie is provided
|
||||
2. Verifies the flow exists and is marked as PUBLIC
|
||||
3. Creates a deterministic UUID based on client_id and original flow_id
|
||||
4. Retrieves the flow owner user for permission purposes
|
||||
|
||||
This function is used to support public flow endpoints that don't require
|
||||
authentication but still need to operate within the permission model.
|
||||
|
||||
Args:
|
||||
flow_id: The original flow ID to verify
|
||||
client_id: The client ID from the request cookie
|
||||
|
||||
Returns:
|
||||
tuple: (flow owner user, deterministic flow ID for tracking)
|
||||
|
||||
Raises:
|
||||
HTTPException:
|
||||
- 400 if no client_id is provided
|
||||
- 403 if flow doesn't exist or isn't public
|
||||
- 403 if unable to retrieve the flow owner user
|
||||
- 403 if user is not found for public flow
|
||||
"""
|
||||
if not client_id:
|
||||
raise HTTPException(status_code=400, detail="No client_id cookie found")
|
||||
|
||||
# Check if the flow is public
|
||||
async with session_scope() as session:
|
||||
from sqlmodel import select
|
||||
|
||||
from langflow.services.database.models.flow.model import AccessTypeEnum, Flow
|
||||
|
||||
flow = (await session.exec(select(Flow).where(Flow.id == flow_id))).first()
|
||||
if not flow or flow.access_type is not AccessTypeEnum.PUBLIC:
|
||||
raise HTTPException(status_code=403, detail="Flow is not public")
|
||||
|
||||
# Create a new flow ID using the client_id and flow_id
|
||||
new_id = f"{client_id}_{flow_id}"
|
||||
new_flow_id = uuid.uuid5(uuid.NAMESPACE_DNS, new_id)
|
||||
|
||||
# Get the user associated with the flow
|
||||
try:
|
||||
from langflow.helpers.user import get_user_by_flow_id_or_endpoint_name
|
||||
|
||||
user = await get_user_by_flow_id_or_endpoint_name(str(flow_id))
|
||||
|
||||
except Exception as exc:
|
||||
logger.exception(f"Error getting user for public flow {flow_id}")
|
||||
raise HTTPException(status_code=403, detail="Flow is not accessible") from exc
|
||||
|
||||
if not user:
|
||||
msg = f"User not found for public flow {flow_id}"
|
||||
raise HTTPException(status_code=403, detail=msg)
|
||||
|
||||
return user, new_flow_id
|
||||
|
|
|
|||
|
|
@ -6,7 +6,15 @@ import traceback
|
|||
import uuid
|
||||
from typing import TYPE_CHECKING, Annotated
|
||||
|
||||
from fastapi import APIRouter, BackgroundTasks, Body, Depends, HTTPException, status
|
||||
from fastapi import (
|
||||
APIRouter,
|
||||
BackgroundTasks,
|
||||
Body,
|
||||
Depends,
|
||||
HTTPException,
|
||||
Request,
|
||||
status,
|
||||
)
|
||||
from fastapi.responses import StreamingResponse
|
||||
from loguru import logger
|
||||
|
||||
|
|
@ -25,6 +33,7 @@ from langflow.api.utils import (
|
|||
format_exception_message,
|
||||
get_top_level_vertices,
|
||||
parse_exception,
|
||||
verify_public_flow_and_get_user,
|
||||
)
|
||||
from langflow.api.v1.schemas import (
|
||||
CancelFlowResponse,
|
||||
|
|
@ -142,8 +151,29 @@ async def build_flow(
|
|||
log_builds: bool = True,
|
||||
current_user: CurrentActiveUser,
|
||||
queue_service: Annotated[JobQueueService, Depends(get_queue_service)],
|
||||
flow_name: str | None = None,
|
||||
):
|
||||
"""Build and process a flow, returning a job ID for event polling."""
|
||||
"""Build and process a flow, returning a job ID for event polling.
|
||||
|
||||
This endpoint requires authentication through the CurrentActiveUser dependency.
|
||||
For public flows that don't require authentication, use the /build_public_tmp/{flow_id}/flow endpoint.
|
||||
|
||||
Args:
|
||||
flow_id: UUID of the flow to build
|
||||
background_tasks: Background tasks manager
|
||||
inputs: Optional input values for the flow
|
||||
data: Optional flow data
|
||||
files: Optional files to include
|
||||
stop_component_id: Optional ID of component to stop at
|
||||
start_component_id: Optional ID of component to start from
|
||||
log_builds: Whether to log the build process
|
||||
current_user: The authenticated user
|
||||
queue_service: Queue service for job management
|
||||
flow_name: Optional name for the flow
|
||||
|
||||
Returns:
|
||||
Dict with job_id that can be used to poll for build status
|
||||
"""
|
||||
# First verify the flow exists
|
||||
async with session_scope() as session:
|
||||
flow = await session.get(Flow, flow_id)
|
||||
|
|
@ -161,6 +191,7 @@ async def build_flow(
|
|||
log_builds=log_builds,
|
||||
current_user=current_user,
|
||||
queue_service=queue_service,
|
||||
flow_name=flow_name,
|
||||
)
|
||||
return {"job_id": job_id}
|
||||
|
||||
|
|
@ -254,7 +285,9 @@ async def build_vertex(
|
|||
# If there's no cache
|
||||
logger.warning(f"No cache found for {flow_id_str}. Building graph starting at {vertex_id}")
|
||||
graph = await build_graph_from_db(
|
||||
flow_id=flow_id, session=await anext(get_session()), chat_service=chat_service
|
||||
flow_id=flow_id,
|
||||
session=await anext(get_session()),
|
||||
chat_service=chat_service,
|
||||
)
|
||||
else:
|
||||
graph = cache.get("result")
|
||||
|
|
@ -450,7 +483,11 @@ async def _stream_vertex(flow_id: str, vertex_id: str, chat_service: ChatService
|
|||
yield str(StreamData(event="close", data={"message": "Stream closed"}))
|
||||
|
||||
|
||||
@router.get("/build/{flow_id}/{vertex_id}/stream", response_class=StreamingResponse, deprecated=True)
|
||||
@router.get(
|
||||
"/build/{flow_id}/{vertex_id}/stream",
|
||||
response_class=StreamingResponse,
|
||||
deprecated=True,
|
||||
)
|
||||
async def build_vertex_stream(
|
||||
flow_id: uuid.UUID,
|
||||
vertex_id: str,
|
||||
|
|
@ -482,7 +519,80 @@ async def build_vertex_stream(
|
|||
"""
|
||||
try:
|
||||
return StreamingResponse(
|
||||
_stream_vertex(str(flow_id), vertex_id, get_chat_service()), media_type="text/event-stream"
|
||||
_stream_vertex(str(flow_id), vertex_id, get_chat_service()),
|
||||
media_type="text/event-stream",
|
||||
)
|
||||
except Exception as exc:
|
||||
raise HTTPException(status_code=500, detail="Error building Component") from exc
|
||||
|
||||
|
||||
@router.post("/build_public_tmp/{flow_id}/flow")
|
||||
async def build_public_tmp(
|
||||
*,
|
||||
background_tasks: LimitVertexBuildBackgroundTasks,
|
||||
flow_id: uuid.UUID,
|
||||
inputs: Annotated[InputValueRequest | None, Body(embed=True)] = None,
|
||||
data: Annotated[FlowDataRequest | None, Body(embed=True)] = None,
|
||||
files: list[str] | None = None,
|
||||
stop_component_id: str | None = None,
|
||||
start_component_id: str | None = None,
|
||||
log_builds: bool | None = True,
|
||||
flow_name: str | None = None,
|
||||
request: Request,
|
||||
queue_service: Annotated[JobQueueService, Depends(get_queue_service)],
|
||||
):
|
||||
"""Build a public flow without requiring authentication.
|
||||
|
||||
This endpoint is specifically for public flows that don't require authentication.
|
||||
It uses a client_id cookie to create a deterministic flow ID for tracking purposes.
|
||||
|
||||
The endpoint:
|
||||
1. Verifies the requested flow is marked as public in the database
|
||||
2. Creates a deterministic UUID based on client_id and flow_id
|
||||
3. Uses the flow owner's permissions to build the flow
|
||||
|
||||
Requirements:
|
||||
- The flow must be marked as PUBLIC in the database
|
||||
- The request must include a client_id cookie
|
||||
|
||||
Args:
|
||||
flow_id: UUID of the public flow to build
|
||||
background_tasks: Background tasks manager
|
||||
inputs: Optional input values for the flow
|
||||
data: Optional flow data
|
||||
files: Optional files to include
|
||||
stop_component_id: Optional ID of component to stop at
|
||||
start_component_id: Optional ID of component to start from
|
||||
log_builds: Whether to log the build process
|
||||
flow_name: Optional name for the flow
|
||||
request: FastAPI request object (needed for cookie access)
|
||||
queue_service: Queue service for job management
|
||||
|
||||
Returns:
|
||||
Dict with job_id that can be used to poll for build status
|
||||
"""
|
||||
try:
|
||||
# Verify this is a public flow and get the associated user
|
||||
client_id = request.cookies.get("client_id")
|
||||
owner_user, new_flow_id = await verify_public_flow_and_get_user(flow_id=flow_id, client_id=client_id)
|
||||
|
||||
# Start the flow build using the new flow ID
|
||||
job_id = await start_flow_build(
|
||||
flow_id=new_flow_id,
|
||||
background_tasks=background_tasks,
|
||||
inputs=inputs,
|
||||
data=data,
|
||||
files=files,
|
||||
stop_component_id=stop_component_id,
|
||||
start_component_id=start_component_id,
|
||||
log_builds=log_builds or False,
|
||||
current_user=owner_user,
|
||||
queue_service=queue_service,
|
||||
flow_name=flow_name or f"{client_id}_{flow_id}",
|
||||
)
|
||||
except Exception as exc:
|
||||
logger.exception("Error building public flow")
|
||||
if isinstance(exc, HTTPException):
|
||||
raise
|
||||
raise HTTPException(status_code=500, detail=str(exc)) from exc
|
||||
return {"job_id": job_id}
|
||||
|
|
|
|||
|
|
@ -21,10 +21,11 @@ from sqlmodel.ext.asyncio.session import AsyncSession
|
|||
|
||||
from langflow.api.utils import CurrentActiveUser, DbSession, cascade_delete_flow, remove_api_keys, validate_is_component
|
||||
from langflow.api.v1.schemas import FlowListCreate
|
||||
from langflow.helpers.user import get_user_by_flow_id_or_endpoint_name
|
||||
from langflow.initial_setup.constants import STARTER_FOLDER_NAME
|
||||
from langflow.logging import logger
|
||||
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.model import AccessTypeEnum, FlowHeader
|
||||
from langflow.services.database.models.flow.utils import get_webhook_component_in_flow
|
||||
from langflow.services.database.models.folder.constants import DEFAULT_FOLDER_NAME
|
||||
from langflow.services.database.models.folder.model import Folder
|
||||
|
|
@ -278,6 +279,21 @@ async def read_flow(
|
|||
raise HTTPException(status_code=404, detail="Flow not found")
|
||||
|
||||
|
||||
@router.get("/public_flow/{flow_id}", response_model=FlowRead, status_code=200)
|
||||
async def read_public_flow(
|
||||
*,
|
||||
session: DbSession,
|
||||
flow_id: UUID,
|
||||
):
|
||||
"""Read a public flow."""
|
||||
access_type = (await session.exec(select(Flow.access_type).where(Flow.id == flow_id))).first()
|
||||
if access_type is not AccessTypeEnum.PUBLIC:
|
||||
raise HTTPException(status_code=403, detail="Flow is not public")
|
||||
|
||||
current_user = await get_user_by_flow_id_or_endpoint_name(str(flow_id))
|
||||
return await read_flow(session=session, flow_id=flow_id, current_user=current_user)
|
||||
|
||||
|
||||
@router.patch("/{flow_id}", response_model=FlowRead, status_code=200)
|
||||
async def update_flow(
|
||||
*,
|
||||
|
|
|
|||
|
|
@ -4,7 +4,14 @@ from pathlib import Path
|
|||
from typing import Any, Literal
|
||||
from uuid import UUID
|
||||
|
||||
from pydantic import BaseModel, ConfigDict, Field, field_serializer, field_validator, model_serializer
|
||||
from pydantic import (
|
||||
BaseModel,
|
||||
ConfigDict,
|
||||
Field,
|
||||
field_serializer,
|
||||
field_validator,
|
||||
model_serializer,
|
||||
)
|
||||
|
||||
from langflow.graph.schema import RunOutputs
|
||||
from langflow.schema import dotdict
|
||||
|
|
@ -377,6 +384,8 @@ class ConfigResponse(BaseModel):
|
|||
health_check_max_retries: int
|
||||
max_file_size_upload: int
|
||||
webhook_polling_interval: int
|
||||
public_flow_cleanup_interval: int
|
||||
public_flow_expiration: int
|
||||
event_delivery: Literal["polling", "streaming"]
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -341,6 +341,7 @@ class Graph:
|
|||
self,
|
||||
inputs: list[dict] | None = None,
|
||||
max_iterations: int | None = None,
|
||||
config: StartConfigDict | None = None,
|
||||
event_manager: EventManager | None = None,
|
||||
):
|
||||
if not self._prepared:
|
||||
|
|
@ -348,6 +349,8 @@ class Graph:
|
|||
raise ValueError(msg)
|
||||
# The idea is for this to return a generator that yields the result of
|
||||
# each step call and raise StopIteration when the graph is done
|
||||
if config is not None:
|
||||
self.__apply_config(config)
|
||||
for _input in inputs or []:
|
||||
for key, value in _input.items():
|
||||
vertex = self.get_vertex(key)
|
||||
|
|
|
|||
|
|
@ -274,6 +274,7 @@ def create_app():
|
|||
FastAPIInstrumentor.instrument_app(app)
|
||||
|
||||
add_pagination(app)
|
||||
|
||||
return app
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -1,3 +1,4 @@
|
|||
import asyncio
|
||||
import json
|
||||
from collections.abc import Sequence
|
||||
from uuid import UUID
|
||||
|
|
@ -154,9 +155,19 @@ async def aadd_messagetables(messages: list[MessageTable], session: AsyncSession
|
|||
try:
|
||||
for message in messages:
|
||||
session.add(message)
|
||||
await session.commit()
|
||||
try:
|
||||
await session.commit()
|
||||
# This is a hack.
|
||||
# We are doing this because build_public_tmp causes the CancelledError to be raised
|
||||
# while build_flow does not.
|
||||
except asyncio.CancelledError:
|
||||
await session.commit()
|
||||
for message in messages:
|
||||
await session.refresh(message)
|
||||
except asyncio.CancelledError as e:
|
||||
logger.exception(e)
|
||||
error_msg = "Operation cancelled"
|
||||
raise ValueError(error_msg) from e
|
||||
except Exception as e:
|
||||
logger.exception(e)
|
||||
raise
|
||||
|
|
|
|||
|
|
@ -2,6 +2,7 @@
|
|||
|
||||
import re
|
||||
from datetime import datetime, timezone
|
||||
from enum import Enum
|
||||
from typing import TYPE_CHECKING, Optional
|
||||
from uuid import UUID, uuid4
|
||||
|
||||
|
|
@ -15,7 +16,8 @@ from pydantic import (
|
|||
field_serializer,
|
||||
field_validator,
|
||||
)
|
||||
from sqlalchemy import Text, UniqueConstraint
|
||||
from sqlalchemy import Enum as SQLEnum
|
||||
from sqlalchemy import Text, UniqueConstraint, text
|
||||
from sqlmodel import JSON, Column, Field, Relationship, SQLModel
|
||||
|
||||
from langflow.schema import Data
|
||||
|
|
@ -30,6 +32,11 @@ if TYPE_CHECKING:
|
|||
HEX_COLOR_LENGTH = 7
|
||||
|
||||
|
||||
class AccessTypeEnum(str, Enum):
|
||||
PRIVATE = "private"
|
||||
PUBLIC = "public"
|
||||
|
||||
|
||||
class FlowBase(SQLModel):
|
||||
name: str = Field(index=True)
|
||||
description: str | None = Field(default=None, sa_column=Column(Text, index=True, nullable=True))
|
||||
|
|
@ -43,6 +50,18 @@ class FlowBase(SQLModel):
|
|||
endpoint_name: str | None = Field(default=None, nullable=True, index=True)
|
||||
tags: list[str] | None = None
|
||||
locked: bool | None = Field(default=False, nullable=True)
|
||||
access_type: AccessTypeEnum = Field(
|
||||
default=AccessTypeEnum.PRIVATE,
|
||||
sa_column=Column(
|
||||
SQLEnum(
|
||||
AccessTypeEnum,
|
||||
name="access_type_enum",
|
||||
values_callable=lambda enum: [member.value for member in enum],
|
||||
),
|
||||
nullable=False,
|
||||
server_default=text("'private'"),
|
||||
),
|
||||
)
|
||||
|
||||
@field_validator("endpoint_name")
|
||||
@classmethod
|
||||
|
|
@ -218,6 +237,7 @@ class FlowHeader(BaseModel):
|
|||
endpoint_name: str | None = Field(None, description="The name of the endpoint associated with this flow")
|
||||
description: str | None = Field(None, description="A description of the flow")
|
||||
data: dict | None = Field(None, description="The data of the component, if is_component is True")
|
||||
access_type: AccessTypeEnum | None = Field(None, description="The access type of the flow")
|
||||
tags: list[str] | None = Field(None, description="The tags of the flow")
|
||||
|
||||
@field_validator("data", mode="before")
|
||||
|
|
@ -235,6 +255,7 @@ class FlowUpdate(SQLModel):
|
|||
folder_id: UUID | None = None
|
||||
endpoint_name: str | None = None
|
||||
locked: bool | None = None
|
||||
access_type: AccessTypeEnum | None = None
|
||||
fs_path: str | None = None
|
||||
|
||||
@field_validator("endpoint_name")
|
||||
|
|
|
|||
|
|
@ -10,9 +10,14 @@ import orjson
|
|||
import yaml
|
||||
from aiofile import async_open
|
||||
from loguru import logger
|
||||
from pydantic import field_validator
|
||||
from pydantic import Field, field_validator
|
||||
from pydantic.fields import FieldInfo
|
||||
from pydantic_settings import BaseSettings, EnvSettingsSource, PydanticBaseSettingsSource, SettingsConfigDict
|
||||
from pydantic_settings import (
|
||||
BaseSettings,
|
||||
EnvSettingsSource,
|
||||
PydanticBaseSettingsSource,
|
||||
SettingsConfigDict,
|
||||
)
|
||||
from typing_extensions import override
|
||||
|
||||
from langflow.services.settings.constants import VARIABLES_TO_GET_FROM_ENVIRONMENT
|
||||
|
|
@ -224,6 +229,13 @@ class Settings(BaseSettings):
|
|||
mcp_server_enable_progress_notifications: bool = False
|
||||
"""If set to False, Langflow will not send progress notifications in the MCP server."""
|
||||
|
||||
# Public Flow Settings
|
||||
public_flow_cleanup_interval: int = Field(default=3600, gt=600)
|
||||
"""The interval in seconds at which public temporary flows will be cleaned up.
|
||||
Default is 1 hour (3600 seconds). Minimum is 600 seconds (10 minutes)."""
|
||||
public_flow_expiration: int = Field(default=86400, gt=600)
|
||||
"""The time in seconds after which a public temporary flow will be considered expired and eligible for cleanup.
|
||||
Default is 24 hours (86400 seconds). Minimum is 600 seconds (10 minutes)."""
|
||||
event_delivery: Literal["polling", "streaming"] = "polling"
|
||||
"""How to deliver build events to the frontend. Can be 'polling' or 'streaming'."""
|
||||
|
||||
|
|
|
|||
|
|
@ -11,7 +11,6 @@ class LocalStorageService(StorageService):
|
|||
def __init__(self, session_service, settings_service) -> None:
|
||||
"""Initialize the local storage service with session and settings services."""
|
||||
super().__init__(session_service, settings_service)
|
||||
self.data_dir = anyio.Path(settings_service.settings.config_dir)
|
||||
self.set_ready()
|
||||
|
||||
def build_full_path(self, flow_id: str, file_name: str) -> str:
|
||||
|
|
|
|||
|
|
@ -3,6 +3,8 @@ from __future__ import annotations
|
|||
from abc import abstractmethod
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
import anyio
|
||||
|
||||
from langflow.services.base import Service
|
||||
|
||||
if TYPE_CHECKING:
|
||||
|
|
@ -16,6 +18,7 @@ class StorageService(Service):
|
|||
def __init__(self, session_service: SessionService, settings_service: SettingsService):
|
||||
self.settings_service = settings_service
|
||||
self.session_service = session_service
|
||||
self.data_dir: anyio.Path = anyio.Path(settings_service.settings.config_dir)
|
||||
self.set_ready()
|
||||
|
||||
def build_full_path(self, flow_id: str, file_name: str) -> str:
|
||||
|
|
|
|||
136
src/backend/base/langflow/services/task/temp_flow_cleanup.py
Normal file
136
src/backend/base/langflow/services/task/temp_flow_cleanup.py
Normal file
|
|
@ -0,0 +1,136 @@
|
|||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import contextlib
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
from loguru import logger
|
||||
from sqlmodel import col, delete, select
|
||||
|
||||
from langflow.services.database.models.message.model import MessageTable
|
||||
from langflow.services.database.models.transactions.model import TransactionTable
|
||||
from langflow.services.database.models.vertex_builds.model import VertexBuildTable
|
||||
from langflow.services.deps import get_settings_service, get_storage_service, session_scope
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from langflow.services.storage.service import StorageService
|
||||
|
||||
|
||||
async def cleanup_orphaned_records() -> None:
|
||||
"""Clean up all records that reference non-existent flows."""
|
||||
from langflow.services.database.models.flow.model import Flow
|
||||
|
||||
async with session_scope() as session:
|
||||
# Create a subquery of existing flow IDs
|
||||
flow_ids_subquery = select(Flow.id)
|
||||
|
||||
# Tables that have flow_id foreign keys
|
||||
tables: list[type[VertexBuildTable | MessageTable | TransactionTable]] = [
|
||||
MessageTable,
|
||||
VertexBuildTable,
|
||||
TransactionTable,
|
||||
]
|
||||
|
||||
for table in tables:
|
||||
try:
|
||||
# Get distinct orphaned flow IDs from the table
|
||||
orphaned_flow_ids = (
|
||||
await session.exec(
|
||||
select(col(table.flow_id).distinct()).where(col(table.flow_id).not_in(flow_ids_subquery))
|
||||
)
|
||||
).all()
|
||||
|
||||
if orphaned_flow_ids:
|
||||
logger.debug(f"Found {len(orphaned_flow_ids)} orphaned flow IDs in {table.__name__}")
|
||||
|
||||
# Delete all orphaned records in a single query
|
||||
await session.exec(delete(table).where(col(table.flow_id).in_(orphaned_flow_ids)))
|
||||
|
||||
# Clean up any associated storage files
|
||||
storage_service: StorageService = get_storage_service()
|
||||
for flow_id in orphaned_flow_ids:
|
||||
try:
|
||||
files = await storage_service.list_files(str(flow_id))
|
||||
for file in files:
|
||||
try:
|
||||
await storage_service.delete_file(str(flow_id), file)
|
||||
except Exception as exc: # noqa: BLE001
|
||||
logger.error(f"Failed to delete file {file} for flow {flow_id}: {exc!s}")
|
||||
# Delete the flow directory after all files are deleted
|
||||
flow_dir = storage_service.data_dir / str(flow_id)
|
||||
if await flow_dir.exists():
|
||||
await flow_dir.rmdir()
|
||||
except Exception as exc: # noqa: BLE001
|
||||
logger.error(f"Failed to list files for flow {flow_id}: {exc!s}")
|
||||
|
||||
await session.commit()
|
||||
logger.debug(f"Successfully deleted orphaned records from {table.__name__}")
|
||||
|
||||
except Exception as exc: # noqa: BLE001
|
||||
logger.error(f"Error cleaning up orphaned records in {table.__name__}: {exc!s}")
|
||||
await session.rollback()
|
||||
|
||||
|
||||
class CleanupWorker:
|
||||
def __init__(self) -> None:
|
||||
self._stop_event = asyncio.Event()
|
||||
self._task: asyncio.Task | None = None
|
||||
|
||||
async def start(self):
|
||||
"""Start the cleanup worker."""
|
||||
if self._task is not None:
|
||||
logger.warning("Cleanup worker is already running")
|
||||
return
|
||||
|
||||
self._task = asyncio.create_task(self._run())
|
||||
logger.debug("Started database cleanup worker")
|
||||
|
||||
async def stop(self):
|
||||
"""Stop the cleanup worker gracefully."""
|
||||
if self._task is None:
|
||||
logger.warning("Cleanup worker is not running")
|
||||
return
|
||||
|
||||
logger.debug("Stopping database cleanup worker...")
|
||||
self._stop_event.set()
|
||||
await self._task
|
||||
self._task = None
|
||||
logger.debug("Database cleanup worker stopped")
|
||||
|
||||
async def _run(self):
|
||||
"""Run the cleanup worker until stopped."""
|
||||
settings = get_settings_service().settings
|
||||
while not self._stop_event.is_set():
|
||||
try:
|
||||
# Clean up any orphaned records
|
||||
await cleanup_orphaned_records()
|
||||
except Exception as exc: # noqa: BLE001
|
||||
logger.error(f"Error in cleanup worker: {exc!s}")
|
||||
|
||||
try:
|
||||
# Create a task for the timeout
|
||||
sleep_task = asyncio.create_task(asyncio.sleep(settings.public_flow_cleanup_interval))
|
||||
# Create a task for the stop event
|
||||
stop_task = asyncio.create_task(self._stop_event.wait())
|
||||
|
||||
# Wait for either the timeout or the stop event
|
||||
done, pending = await asyncio.wait([sleep_task, stop_task], return_when=asyncio.FIRST_COMPLETED)
|
||||
|
||||
# Cancel any pending tasks
|
||||
for task in pending:
|
||||
task.cancel()
|
||||
with contextlib.suppress(asyncio.CancelledError):
|
||||
await task
|
||||
|
||||
# If the stop event completed, break the loop
|
||||
if stop_task in done:
|
||||
break
|
||||
|
||||
except Exception as exc: # noqa: BLE001
|
||||
logger.error(f"Error in cleanup worker sleep: {exc!s}")
|
||||
# Sleep a minimum amount in case of errors
|
||||
await asyncio.sleep(60)
|
||||
|
||||
|
||||
# Create a global instance of the worker
|
||||
cleanup_worker = CleanupWorker()
|
||||
0
src/backend/tests/unit/services/tasks/__init__.py
Normal file
0
src/backend/tests/unit/services/tasks/__init__.py
Normal file
109
src/backend/tests/unit/services/tasks/test_temp_flow_cleanup.py
Normal file
109
src/backend/tests/unit/services/tasks/test_temp_flow_cleanup.py
Normal file
|
|
@ -0,0 +1,109 @@
|
|||
from __future__ import annotations
|
||||
|
||||
import datetime
|
||||
from datetime import timezone
|
||||
from uuid import uuid4
|
||||
|
||||
import pytest
|
||||
from langflow.services.database.models.flow import Flow as FlowTable
|
||||
from langflow.services.database.models.message.model import MessageTable
|
||||
from langflow.services.deps import get_settings_service, get_storage_service, session_scope
|
||||
from langflow.services.task.temp_flow_cleanup import (
|
||||
CleanupWorker,
|
||||
cleanup_orphaned_records,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.usefixtures("client")
|
||||
async def test_cleanup_orphaned_records_no_orphans():
|
||||
"""Test cleanup when there are no orphaned records."""
|
||||
storage_service = get_storage_service()
|
||||
flow_id = uuid4()
|
||||
|
||||
async with session_scope() as session:
|
||||
# Create a flow and associated message
|
||||
flow = FlowTable(
|
||||
id=flow_id,
|
||||
name="Test Flow",
|
||||
data="null",
|
||||
updated_at=datetime.datetime.now(timezone.utc),
|
||||
)
|
||||
message = MessageTable(
|
||||
id=uuid4(),
|
||||
flow_id=flow_id,
|
||||
sender="test_user",
|
||||
sender_name="Test User",
|
||||
timestamp=datetime.datetime.now(timezone.utc),
|
||||
session_id=str(uuid4()),
|
||||
)
|
||||
session.add(flow)
|
||||
session.add(message)
|
||||
await session.commit()
|
||||
|
||||
# Write a file for the flow
|
||||
await storage_service.save_file(str(flow_id), "test.json", b"test data")
|
||||
|
||||
# Run cleanup
|
||||
async with session_scope() as session:
|
||||
await cleanup_orphaned_records()
|
||||
|
||||
# Verify message still exists
|
||||
async with session_scope() as session:
|
||||
message = await session.get(MessageTable, message.id)
|
||||
assert message is not None
|
||||
|
||||
|
||||
@pytest.mark.usefixtures("client")
|
||||
async def test_cleanup_orphaned_records_with_orphans():
|
||||
"""Test cleanup when there are orphaned records."""
|
||||
orphaned_flow_id = uuid4()
|
||||
|
||||
async with session_scope() as session:
|
||||
# Create orphaned records without an associated flow
|
||||
message = MessageTable(
|
||||
id=uuid4(),
|
||||
flow_id=orphaned_flow_id,
|
||||
sender="test_user",
|
||||
sender_name="Test User",
|
||||
timestamp=datetime.datetime.now(timezone.utc),
|
||||
session_id=str(uuid4()),
|
||||
)
|
||||
session.add(message)
|
||||
await session.commit()
|
||||
|
||||
# Run cleanup
|
||||
async with session_scope() as session:
|
||||
await cleanup_orphaned_records()
|
||||
|
||||
# Verify orphaned message was deleted
|
||||
async with session_scope() as session:
|
||||
message = await session.get(MessageTable, message.id)
|
||||
assert message is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_cleanup_worker_start_stop():
|
||||
"""Test CleanupWorker start and stop functionality."""
|
||||
worker = CleanupWorker()
|
||||
await worker.start()
|
||||
assert worker._task is not None
|
||||
assert not worker._stop_event.is_set()
|
||||
await worker.stop()
|
||||
assert worker._task is None
|
||||
assert worker._stop_event.is_set()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_cleanup_worker_run_with_exception(caplog):
|
||||
"""Test CleanupWorker handles exceptions gracefully."""
|
||||
settings = get_settings_service().settings
|
||||
settings.public_flow_cleanup_interval = 601 # Minimum valid interval
|
||||
worker = CleanupWorker()
|
||||
|
||||
# Start worker and let it run briefly
|
||||
await worker.start()
|
||||
await worker.stop()
|
||||
|
||||
# Check logs for expected messages
|
||||
assert any("Started database cleanup worker" in record.message for record in caplog.records)
|
||||
assert any("Stopping database cleanup worker" in record.message for record in caplog.records)
|
||||
Loading…
Add table
Add a link
Reference in a new issue