refactor: Renames flask_context_manager to preserve_flask_contexts

Signed-off-by: -LAN- <laipz8200@outlook.com>
pull/21061/head
-LAN- 11 months ago
parent b011777935
commit 225c5fe2a3
No known key found for this signature in database
GPG Key ID: 6BA0D108DED011FF

@ -31,7 +31,7 @@ from core.workflow.repositories.workflow_execution_repository import WorkflowExe
from core.workflow.repositories.workflow_node_execution_repository import WorkflowNodeExecutionRepository from core.workflow.repositories.workflow_node_execution_repository import WorkflowNodeExecutionRepository
from extensions.ext_database import db from extensions.ext_database import db
from factories import file_factory from factories import file_factory
from libs.flask_utils import flask_context_manager from libs.flask_utils import preserve_flask_contexts
from models import Account, App, Conversation, EndUser, Message, Workflow, WorkflowNodeExecutionTriggeredFrom from models import Account, App, Conversation, EndUser, Message, Workflow, WorkflowNodeExecutionTriggeredFrom
from models.enums import WorkflowRunTriggeredFrom from models.enums import WorkflowRunTriggeredFrom
from services.conversation_service import ConversationService from services.conversation_service import ConversationService
@ -448,7 +448,7 @@ class AdvancedChatAppGenerator(MessageBasedAppGenerator):
:return: :return:
""" """
with flask_context_manager(flask_app, context_vars=context): with preserve_flask_contexts(flask_app, context_vars=context):
try: try:
# get conversation and message # get conversation and message
conversation = self._get_conversation(conversation_id) conversation = self._get_conversation(conversation_id)

@ -23,7 +23,7 @@ from core.model_runtime.errors.invoke import InvokeAuthorizationError
from core.ops.ops_trace_manager import TraceQueueManager from core.ops.ops_trace_manager import TraceQueueManager
from extensions.ext_database import db from extensions.ext_database import db
from factories import file_factory from factories import file_factory
from libs.flask_utils import flask_context_manager from libs.flask_utils import preserve_flask_contexts
from models import Account, App, EndUser from models import Account, App, EndUser
from services.conversation_service import ConversationService from services.conversation_service import ConversationService
from services.errors.message import MessageNotExistsError from services.errors.message import MessageNotExistsError
@ -228,7 +228,7 @@ class AgentChatAppGenerator(MessageBasedAppGenerator):
:return: :return:
""" """
with flask_context_manager(flask_app, context_vars=context): with preserve_flask_contexts(flask_app, context_vars=context):
try: try:
# get conversation and message # get conversation and message
conversation = self._get_conversation(conversation_id) conversation = self._get_conversation(conversation_id)

@ -29,7 +29,7 @@ from core.workflow.repositories.workflow_execution_repository import WorkflowExe
from core.workflow.repositories.workflow_node_execution_repository import WorkflowNodeExecutionRepository from core.workflow.repositories.workflow_node_execution_repository import WorkflowNodeExecutionRepository
from extensions.ext_database import db from extensions.ext_database import db
from factories import file_factory from factories import file_factory
from libs.flask_utils import flask_context_manager from libs.flask_utils import preserve_flask_contexts
from models import Account, App, EndUser, Workflow, WorkflowNodeExecutionTriggeredFrom from models import Account, App, EndUser, Workflow, WorkflowNodeExecutionTriggeredFrom
from models.enums import WorkflowRunTriggeredFrom from models.enums import WorkflowRunTriggeredFrom
@ -407,7 +407,7 @@ class WorkflowAppGenerator(BaseAppGenerator):
:return: :return:
""" """
with flask_context_manager(flask_app, context_vars=context): with preserve_flask_contexts(flask_app, context_vars=context):
try: try:
# workflow app # workflow app
runner = WorkflowAppRunner( runner = WorkflowAppRunner(

@ -53,7 +53,7 @@ from core.workflow.nodes.end.end_stream_processor import EndStreamProcessor
from core.workflow.nodes.enums import ErrorStrategy, FailBranchSourceHandle from core.workflow.nodes.enums import ErrorStrategy, FailBranchSourceHandle
from core.workflow.nodes.event import RunCompletedEvent, RunRetrieverResourceEvent, RunStreamChunkEvent from core.workflow.nodes.event import RunCompletedEvent, RunRetrieverResourceEvent, RunStreamChunkEvent
from core.workflow.nodes.node_mapping import NODE_TYPE_CLASSES_MAPPING from core.workflow.nodes.node_mapping import NODE_TYPE_CLASSES_MAPPING
from libs.flask_utils import flask_context_manager from libs.flask_utils import preserve_flask_contexts
from models.enums import UserFrom from models.enums import UserFrom
from models.workflow import WorkflowType from models.workflow import WorkflowType
@ -539,7 +539,7 @@ class GraphEngine:
Run parallel nodes Run parallel nodes
""" """
with flask_context_manager(flask_app, context_vars=context): with preserve_flask_contexts(flask_app, context_vars=context):
try: try:
q.put( q.put(
ParallelBranchRunStartedEvent( ParallelBranchRunStartedEvent(

@ -37,7 +37,7 @@ from core.workflow.nodes.base import BaseNode
from core.workflow.nodes.enums import NodeType from core.workflow.nodes.enums import NodeType
from core.workflow.nodes.event import NodeEvent, RunCompletedEvent from core.workflow.nodes.event import NodeEvent, RunCompletedEvent
from core.workflow.nodes.iteration.entities import ErrorHandleMode, IterationNodeData from core.workflow.nodes.iteration.entities import ErrorHandleMode, IterationNodeData
from libs.flask_utils import flask_context_manager from libs.flask_utils import preserve_flask_contexts
from .exc import ( from .exc import (
InvalidIteratorValueError, InvalidIteratorValueError,
@ -585,7 +585,7 @@ class IterationNode(BaseNode[IterationNodeData]):
run single iteration in parallel mode run single iteration in parallel mode
""" """
with flask_context_manager(flask_app, context_vars=context): with preserve_flask_contexts(flask_app, context_vars=context):
parallel_mode_run_id = uuid.uuid4().hex parallel_mode_run_id = uuid.uuid4().hex
graph_engine_copy = graph_engine.create_copy() graph_engine_copy = graph_engine.create_copy()
variable_pool_copy = graph_engine_copy.graph_runtime_state.variable_pool variable_pool_copy = graph_engine_copy.graph_runtime_state.variable_pool

@ -9,7 +9,7 @@ T = TypeVar("T")
@contextmanager @contextmanager
def flask_context_manager( def preserve_flask_contexts(
flask_app: Flask, flask_app: Flask,
context_vars: contextvars.Context, context_vars: contextvars.Context,
) -> Iterator[None]: ) -> Iterator[None]:
@ -36,7 +36,7 @@ def flask_context_manager(
Example: Example:
```python ```python
with flask_context_manager(flask_app, context_vars=context_vars): with preserve_flask_contexts(flask_app, context_vars=context_vars):
# Code that needs Flask app context and context variables # Code that needs Flask app context and context variables
# Current user will be preserved if available # Current user will be preserved if available
``` ```

@ -6,7 +6,7 @@ import pytest
from flask import Flask from flask import Flask
from flask_login import LoginManager, UserMixin, current_user, login_user from flask_login import LoginManager, UserMixin, current_user, login_user
from libs.flask_utils import flask_context_manager from libs.flask_utils import preserve_flask_contexts
class User(UserMixin): class User(UserMixin):
@ -45,9 +45,9 @@ def test_user() -> User:
def test_current_user_not_accessible_across_threads(login_app: Flask, test_user: User): def test_current_user_not_accessible_across_threads(login_app: Flask, test_user: User):
""" """
Test that current_user is not accessible in a different thread without flask_context_manager. Test that current_user is not accessible in a different thread without preserve_flask_contexts.
This test demonstrates that without the flask_context_manager, we cannot access This test demonstrates that without the preserve_flask_contexts, we cannot access
current_user in a different thread, even with app_context. current_user in a different thread, even with app_context.
""" """
# Log in the user in the main thread # Log in the user in the main thread
@ -65,7 +65,7 @@ def test_current_user_not_accessible_across_threads(login_app: Flask, test_user:
# Try to access current_user in a different thread with app_context # Try to access current_user in a different thread with app_context
with login_app.app_context(): with login_app.app_context():
# This should fail because current_user is not accessible across threads # This should fail because current_user is not accessible across threads
# without flask_context_manager # without preserve_flask_contexts
result["user_accessible"] = current_user.is_authenticated result["user_accessible"] = current_user.is_authenticated
except Exception as e: except Exception as e:
result["error"] = str(e) # type: ignore result["error"] = str(e) # type: ignore
@ -79,11 +79,11 @@ def test_current_user_not_accessible_across_threads(login_app: Flask, test_user:
assert result["error"] is not None or (result["user_accessible"] is not None and not result["user_accessible"]) assert result["error"] is not None or (result["user_accessible"] is not None and not result["user_accessible"])
def test_current_user_accessible_with_flask_context_manager(login_app: Flask, test_user: User): def test_current_user_accessible_with_preserve_flask_contexts(login_app: Flask, test_user: User):
""" """
Test that current_user is accessible in a different thread with flask_context_manager. Test that current_user is accessible in a different thread with preserve_flask_contexts.
This test demonstrates that with the flask_context_manager, we can access This test demonstrates that with the preserve_flask_contexts, we can access
current_user in a different thread. current_user in a different thread.
""" """
# Log in the user in the main thread # Log in the user in the main thread
@ -101,8 +101,8 @@ def test_current_user_accessible_with_flask_context_manager(login_app: Flask, te
# Define a function to run in a separate thread # Define a function to run in a separate thread
def check_user_in_thread_with_manager(): def check_user_in_thread_with_manager():
try: try:
# Use flask_context_manager to access current_user in a different thread # Use preserve_flask_contexts to access current_user in a different thread
with flask_context_manager(login_app, context_vars): with preserve_flask_contexts(login_app, context_vars):
from flask_login import current_user from flask_login import current_user
if current_user: if current_user:

Loading…
Cancel
Save