mirror of
https://github.com/langgenius/dify.git
synced 2026-06-05 15:40:14 +08:00
fix(tools): use short-lived sessions for icon lookups to prevent idle-in-transaction (#36903)
Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com>
This commit is contained in:
committed by
GitHub
parent
a3265f722e
commit
c64d3e98c4
@@ -962,34 +962,41 @@ class ToolManager:
|
||||
@classmethod
|
||||
def generate_workflow_tool_icon_url(cls, tenant_id: str, provider_id: str) -> EmojiIconDict:
|
||||
try:
|
||||
workflow_provider: WorkflowToolProvider | None = db.session.scalar(
|
||||
select(WorkflowToolProvider)
|
||||
.where(WorkflowToolProvider.tenant_id == tenant_id, WorkflowToolProvider.id == provider_id)
|
||||
.limit(1)
|
||||
)
|
||||
# Use a short-lived session to avoid holding a database transaction
|
||||
# during long-running nested workflow execution.
|
||||
# Fixes: idle in transaction when Workflow Tool runs (#36902)
|
||||
with Session(db.engine, expire_on_commit=False) as session:
|
||||
workflow_provider: WorkflowToolProvider | None = session.scalar(
|
||||
select(WorkflowToolProvider)
|
||||
.where(WorkflowToolProvider.tenant_id == tenant_id, WorkflowToolProvider.id == provider_id)
|
||||
.limit(1)
|
||||
)
|
||||
|
||||
if workflow_provider is None:
|
||||
raise ToolProviderNotFoundError(f"workflow provider {provider_id} not found")
|
||||
if workflow_provider is None:
|
||||
raise ToolProviderNotFoundError(f"workflow provider {provider_id} not found")
|
||||
|
||||
icon = emoji_icon_adapter.validate_json(workflow_provider.icon)
|
||||
return icon
|
||||
icon = emoji_icon_adapter.validate_json(workflow_provider.icon)
|
||||
return icon
|
||||
except Exception:
|
||||
return {"background": "#252525", "content": "\ud83d\ude01"}
|
||||
|
||||
@classmethod
|
||||
def generate_api_tool_icon_url(cls, tenant_id: str, provider_id: str) -> EmojiIconDict:
|
||||
try:
|
||||
api_provider: ApiToolProvider | None = db.session.scalar(
|
||||
select(ApiToolProvider)
|
||||
.where(ApiToolProvider.tenant_id == tenant_id, ApiToolProvider.id == provider_id)
|
||||
.limit(1)
|
||||
)
|
||||
# Use a short-lived session to avoid holding a database transaction
|
||||
# during long-running tool execution.
|
||||
with Session(db.engine, expire_on_commit=False) as session:
|
||||
api_provider: ApiToolProvider | None = session.scalar(
|
||||
select(ApiToolProvider)
|
||||
.where(ApiToolProvider.tenant_id == tenant_id, ApiToolProvider.id == provider_id)
|
||||
.limit(1)
|
||||
)
|
||||
|
||||
if api_provider is None:
|
||||
raise ToolProviderNotFoundError(f"api provider {provider_id} not found")
|
||||
if api_provider is None:
|
||||
raise ToolProviderNotFoundError(f"api provider {provider_id} not found")
|
||||
|
||||
icon = emoji_icon_adapter.validate_json(api_provider.icon)
|
||||
return icon
|
||||
icon = emoji_icon_adapter.validate_json(api_provider.icon)
|
||||
return icon
|
||||
except Exception:
|
||||
return {"background": "#252525", "content": "\ud83d\ude01"}
|
||||
|
||||
|
||||
@@ -6,7 +6,7 @@ import json
|
||||
import threading
|
||||
from types import SimpleNamespace
|
||||
from typing import Any
|
||||
from unittest.mock import Mock, patch
|
||||
from unittest.mock import MagicMock, Mock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
@@ -802,17 +802,36 @@ def test_generate_tool_icon_urls_for_builtin_and_plugin():
|
||||
def test_generate_tool_icon_urls_for_workflow_and_api():
|
||||
workflow_provider = SimpleNamespace(icon='{"background": "#222", "content": "W"}')
|
||||
api_provider = SimpleNamespace(icon='{"background": "#333", "content": "A"}')
|
||||
mock_engine = object()
|
||||
with patch("core.tools.tool_manager.db") as mock_db:
|
||||
mock_db.session.scalar.side_effect = [workflow_provider, api_provider]
|
||||
assert ToolManager.generate_workflow_tool_icon_url("tenant-1", "wf-1") == {"background": "#222", "content": "W"}
|
||||
assert ToolManager.generate_api_tool_icon_url("tenant-1", "api-1") == {"background": "#333", "content": "A"}
|
||||
mock_db.engine = mock_engine
|
||||
with patch("core.tools.tool_manager.Session") as mock_session_cls:
|
||||
mock_session = MagicMock()
|
||||
mock_session.scalar.side_effect = [workflow_provider, api_provider]
|
||||
mock_session_cls.return_value.__enter__ = MagicMock(return_value=mock_session)
|
||||
mock_session_cls.return_value.__exit__ = MagicMock(return_value=False)
|
||||
assert ToolManager.generate_workflow_tool_icon_url("tenant-1", "wf-1") == {
|
||||
"background": "#222",
|
||||
"content": "W",
|
||||
}
|
||||
assert ToolManager.generate_api_tool_icon_url("tenant-1", "api-1") == {"background": "#333", "content": "A"}
|
||||
# Verify sessions are created with the engine
|
||||
assert mock_session_cls.call_count == 2
|
||||
mock_session_cls.assert_called_with(mock_engine, expire_on_commit=False)
|
||||
|
||||
|
||||
def test_generate_tool_icon_urls_missing_workflow_and_api_use_default():
|
||||
mock_engine = object()
|
||||
with patch("core.tools.tool_manager.db") as mock_db:
|
||||
mock_db.session.scalar.return_value = None
|
||||
assert ToolManager.generate_workflow_tool_icon_url("tenant-1", "missing")["background"] == "#252525"
|
||||
assert ToolManager.generate_api_tool_icon_url("tenant-1", "missing")["background"] == "#252525"
|
||||
mock_db.engine = mock_engine
|
||||
with patch("core.tools.tool_manager.Session") as mock_session_cls:
|
||||
mock_session = MagicMock()
|
||||
mock_session.scalar.return_value = None
|
||||
mock_session_cls.return_value.__enter__ = MagicMock(return_value=mock_session)
|
||||
mock_session_cls.return_value.__exit__ = MagicMock(return_value=False)
|
||||
assert ToolManager.generate_workflow_tool_icon_url("tenant-1", "missing")["background"] == "#252525"
|
||||
assert ToolManager.generate_api_tool_icon_url("tenant-1", "missing")["background"] == "#252525"
|
||||
assert mock_session_cls.call_count == 2
|
||||
|
||||
|
||||
def test_get_tool_icon_for_builtin_provider_variants():
|
||||
|
||||
Reference in New Issue
Block a user