Skip to content

refactor: port api/services/auth/api_key_auth_service.py api/services/datasource_provider_service.py api/services/saved_message_service.py api/services/web_conversation_service.py to session #32003

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
asukaminato0721 wants to merge 7 commits into
base: main
Choose a base branch
from
+298 −236
Open

refactor: port api/services/auth/api_key_auth_service.py api/services/datasource_provider_service.py api/services/saved_message_service.py api/services/web_conversation_service.py to session #32003

Show file tree
Hide file tree
Changes from 4 commits
Commits
Show all changes
7 commits
Select commit Hold shift + click to select a range
769cd87
port some files
asukaminato0721 Feb 5, 2026
c6e065e
Update api/services/auth/api_key_auth_service.py
asukaminato0721 Feb 5, 2026
130b082
Update api/services/saved_message_service.py
asukaminato0721 Feb 5, 2026
0269231
Merge branch 'main' into 24115
asukaminato0721 Feb 5, 2026
85c8544
Merge branch 'main' into 24115
asukaminato0721 Feb 6, 2026
7061d4c
fix test
asukaminato0721 Feb 6, 2026
54a2f68
Merge branch 'main' into 24115
asukaminato0721 Feb 6, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
80 changes: 41 additions & 39 deletions api/services/auth/api_key_auth_service.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,66 +2,68 @@

from sqlalchemy import select

from core.db.session_factory import session_factory
from core.helper import encrypter
from extensions.ext_database import db
from models.source import DataSourceApiKeyAuthBinding
from services.auth.api_key_auth_factory import ApiKeyAuthFactory


class ApiKeyAuthService:
@staticmethod
def get_provider_auth_list(tenant_id: str):
data_source_api_key_bindings = db.session.scalars(
select(DataSourceApiKeyAuthBinding).where(
DataSourceApiKeyAuthBinding.tenant_id == tenant_id, DataSourceApiKeyAuthBinding.disabled.is_(False)
)
).all()
return data_source_api_key_bindings
with session_factory.create_session() as session:
data_source_api_key_bindings = session.scalars(
select(DataSourceApiKeyAuthBinding).where(
DataSourceApiKeyAuthBinding.tenant_id == tenant_id, DataSourceApiKeyAuthBinding.disabled.is_(False)
)
).all()
return data_source_api_key_bindings

@staticmethod
def create_provider_auth(tenant_id: str, args: dict):
auth_result = ApiKeyAuthFactory(args["provider"], args["credentials"]).validate_credentials()
if auth_result:
# Encrypt the api key
api_key = encrypter.encrypt_token(tenant_id, args["credentials"]["config"]["api_key"])
args["credentials"]["config"]["api_key"] = api_key
with session_factory.create_session() as session, session.begin():
# Encrypt the api key
api_key = encrypter.encrypt_token(tenant_id, args["credentials"]["config"]["api_key"])
args["credentials"]["config"]["api_key"] = api_key

data_source_api_key_binding = DataSourceApiKeyAuthBinding(
tenant_id=tenant_id, category=args["category"], provider=args["provider"]
)
data_source_api_key_binding.credentials = json.dumps(args["credentials"], ensure_ascii=False)
db.session.add(data_source_api_key_binding)
db.session.commit()
data_source_api_key_binding = DataSourceApiKeyAuthBinding(
tenant_id=tenant_id, category=args["category"], provider=args["provider"]
)
data_source_api_key_binding.credentials = json.dumps(args["credentials"], ensure_ascii=False)
session.add(data_source_api_key_binding)

@staticmethod
def get_auth_credentials(tenant_id: str, category: str, provider: str):
data_source_api_key_bindings = (
db.session.query(DataSourceApiKeyAuthBinding)
.where(
DataSourceApiKeyAuthBinding.tenant_id == tenant_id,
DataSourceApiKeyAuthBinding.category == category,
DataSourceApiKeyAuthBinding.provider == provider,
DataSourceApiKeyAuthBinding.disabled.is_(False),
with session_factory.create_session() as session:
data_source_api_key_bindings = (
session.query(DataSourceApiKeyAuthBinding)
.where(
DataSourceApiKeyAuthBinding.tenant_id == tenant_id,
DataSourceApiKeyAuthBinding.category == category,
DataSourceApiKeyAuthBinding.provider == provider,
DataSourceApiKeyAuthBinding.disabled.is_(False),
)
.first()
)
.first()
)
if not data_source_api_key_bindings:
return None
if not data_source_api_key_bindings.credentials:
return None
credentials = json.loads(data_source_api_key_bindings.credentials)
return credentials
if not data_source_api_key_bindings:
return None
if not data_source_api_key_bindings.credentials:
return None
credentials = json.loads(data_source_api_key_bindings.credentials)
return credentials
Comment on lines +39 to +55
Copy link

Copilot AI Feb 5, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The data_source_api_key_bindings object accessed on lines 50-54 is detached from the session (since it's accessed after being queried within a session context that will exit). While the scalar attributes like credentials are already loaded and accessible, following the established pattern in the codebase (e.g., api/core/tools/workflow_as_tool/tool.py:226), the object should be explicitly expunged using session.expunge(data_source_api_key_bindings) before being used outside the session context. This makes it clear that the object is intended to be used as a detached object.

Copilot uses AI. Check for mistakes.

@staticmethod
def delete_provider_auth(tenant_id: str, binding_id: str):
data_source_api_key_binding = (
db.session.query(DataSourceApiKeyAuthBinding)
.where(DataSourceApiKeyAuthBinding.tenant_id == tenant_id, DataSourceApiKeyAuthBinding.id == binding_id)
.first()
)
if data_source_api_key_binding:
db.session.delete(data_source_api_key_binding)
db.session.commit()
with session_factory.create_session() as session, session.begin():
data_source_api_key_binding = (
session.query(DataSourceApiKeyAuthBinding)
.where(DataSourceApiKeyAuthBinding.tenant_id == tenant_id, DataSourceApiKeyAuthBinding.id == binding_id)
.first()
)
if data_source_api_key_binding:
session.delete(data_source_api_key_binding)

@classmethod
def validate_api_key_auth_args(cls, args):
Expand Down
17 changes: 9 additions & 8 deletions api/services/datasource_provider_service.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,7 @@

from configs import dify_config
from constants import HIDDEN_VALUE, UNKNOWN_VALUE
from core.db.session_factory import session_factory
from core.helper import encrypter
from core.helper.name_generator import generate_incremental_name
from core.helper.provider_cache import NoOpProviderCredentialCache
Expand Down Expand Up @@ -986,11 +987,11 @@ def remove_datasource_credentials(self, tenant_id: str, auth_id: str, provider:
:param plugin_id: plugin id
:return:
"""
datasource_provider = (
db.session.query(DatasourceProvider)
.filter_by(tenant_id=tenant_id, id=auth_id, provider=provider, plugin_id=plugin_id)
.first()
)
if datasource_provider:
db.session.delete(datasource_provider)
db.session.commit()
with session_factory.create_session() as session, session.begin():
datasource_provider = (
session.query(DatasourceProvider)
.filter_by(tenant_id=tenant_id, id=auth_id, provider=provider, plugin_id=plugin_id)
.first()
)
if datasource_provider:
session.delete(datasource_provider)
88 changes: 45 additions & 43 deletions api/services/saved_message_service.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
from typing import Union

from extensions.ext_database import db
from core.db.session_factory import session_factory
from libs.infinite_scroll_pagination import InfiniteScrollPagination
from models import Account
from models.model import App, EndUser
Expand All @@ -15,16 +15,17 @@ def pagination_by_last_id(
) -> InfiniteScrollPagination:
if not user:
raise ValueError("User is required")
saved_messages = (
db.session.query(SavedMessage)
.where(
SavedMessage.app_id == app_model.id,
SavedMessage.created_by_role == ("account" if isinstance(user, Account) else "end_user"),
SavedMessage.created_by == user.id,
with session_factory.create_session() as session:
saved_messages = (
session.query(SavedMessage)
.where(
SavedMessage.app_id == app_model.id,
SavedMessage.created_by_role == ("account" if isinstance(user, Account) else "end_user"),
SavedMessage.created_by == user.id,
)
.order_by(SavedMessage.created_at.desc())
.all()
)
.order_by(SavedMessage.created_at.desc())
.all()
)
message_ids = [sm.message_id for sm in saved_messages]

return MessageService.pagination_by_last_id(
Expand All @@ -35,49 +36,50 @@ def pagination_by_last_id(
def save(cls, app_model: App, user: Union[Account, EndUser] | None, message_id: str):
if not user:
return
saved_message = (
db.session.query(SavedMessage)
.where(
SavedMessage.app_id == app_model.id,
SavedMessage.message_id == message_id,
SavedMessage.created_by_role == ("account" if isinstance(user, Account) else "end_user"),
SavedMessage.created_by == user.id,

with session_factory.create_session() as session, session.begin():
saved_message = (
session.query(SavedMessage)
.where(
SavedMessage.app_id == app_model.id,
SavedMessage.message_id == message_id,
SavedMessage.created_by_role == ("account" if isinstance(user, Account) else "end_user"),
SavedMessage.created_by == user.id,
)
.first()
)
.first()
)

if saved_message:
return
if saved_message:
return
Comment on lines +52 to +53
Copy link

Copilot AI Feb 5, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The early return on line 53 occurs within a session.begin() transaction context. When exiting the context manager via return, an empty transaction is committed. While not incorrect, this is inefficient. Consider restructuring to avoid early returns within the transaction block, or move the transaction begin to after the early return check.

Copilot uses AI. Check for mistakes.

message = MessageService.get_message(app_model=app_model, user=user, message_id=message_id)
message = MessageService.get_message(app_model=app_model, user=user, message_id=message_id)

saved_message = SavedMessage(
app_id=app_model.id,
message_id=message.id,
created_by_role="account" if isinstance(user, Account) else "end_user",
created_by=user.id,
)
saved_message = SavedMessage(
app_id=app_model.id,
message_id=message.id,
created_by_role="account" if isinstance(user, Account) else "end_user",
created_by=user.id,
)

db.session.add(saved_message)
db.session.commit()
session.add(saved_message)

@classmethod
def delete(cls, app_model: App, user: Union[Account, EndUser] | None, message_id: str):
if not user:
return
saved_message = (
db.session.query(SavedMessage)
.where(
SavedMessage.app_id == app_model.id,
SavedMessage.message_id == message_id,
SavedMessage.created_by_role == ("account" if isinstance(user, Account) else "end_user"),
SavedMessage.created_by == user.id,
with session_factory.create_session() as session, session.begin():
saved_message = (
session.query(SavedMessage)
.where(
SavedMessage.app_id == app_model.id,
SavedMessage.message_id == message_id,
SavedMessage.created_by_role == ("account" if isinstance(user, Account) else "end_user"),
SavedMessage.created_by == user.id,
)
.first()
)
.first()
)

if not saved_message:
return
if not saved_message:
return
Comment on lines +82 to +83
Copy link

Copilot AI Feb 5, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The early return on line 83 occurs within a session.begin() transaction context. When exiting the context manager via return, an empty transaction is committed. While not incorrect, this is inefficient. Consider restructuring to avoid early returns within the transaction block, or move the transaction begin to after the early return check.

Copilot uses AI. Check for mistakes.

db.session.delete(saved_message)
db.session.commit()
session.delete(saved_message)
72 changes: 36 additions & 36 deletions api/services/web_conversation_service.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,7 @@
from sqlalchemy.orm import Session

from core.app.entities.app_invoke_entities import InvokeFrom
from extensions.ext_database import db
from core.db.session_factory import session_factory
from libs.infinite_scroll_pagination import InfiniteScrollPagination
from models import Account
from models.model import App, EndUser
Expand Down Expand Up @@ -63,51 +63,51 @@ def pagination_by_last_id(
def pin(cls, app_model: App, conversation_id: str, user: Union[Account, EndUser] | None):
if not user:
return
pinned_conversation = (
db.session.query(PinnedConversation)
.where(
PinnedConversation.app_id == app_model.id,
PinnedConversation.conversation_id == conversation_id,
PinnedConversation.created_by_role == ("account" if isinstance(user, Account) else "end_user"),
PinnedConversation.created_by == user.id,
with session_factory.create_session() as session, session.begin():
pinned_conversation = (
session.query(PinnedConversation)
.where(
PinnedConversation.app_id == app_model.id,
PinnedConversation.conversation_id == conversation_id,
PinnedConversation.created_by_role == ("account" if isinstance(user, Account) else "end_user"),
PinnedConversation.created_by == user.id,
)
.first()
)
.first()
)

if pinned_conversation:
return
if pinned_conversation:
return
Comment on lines +78 to +79
Copy link

Copilot AI Feb 5, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The early return on line 79 occurs within a session.begin() transaction context. When exiting the context manager via return, an empty transaction is committed. While not incorrect, this is inefficient. Consider restructuring to avoid early returns within the transaction block, or move the transaction begin to after the early return check.

Copilot uses AI. Check for mistakes.

conversation = ConversationService.get_conversation(
app_model=app_model, conversation_id=conversation_id, user=user
)
conversation = ConversationService.get_conversation(
app_model=app_model, conversation_id=conversation_id, user=user
)

pinned_conversation = PinnedConversation(
app_id=app_model.id,
conversation_id=conversation.id,
created_by_role="account" if isinstance(user, Account) else "end_user",
created_by=user.id,
)
pinned_conversation = PinnedConversation(
app_id=app_model.id,
conversation_id=conversation.id,
created_by_role="account" if isinstance(user, Account) else "end_user",
created_by=user.id,
)

db.session.add(pinned_conversation)
db.session.commit()
session.add(pinned_conversation)

@classmethod
def unpin(cls, app_model: App, conversation_id: str, user: Union[Account, EndUser] | None):
if not user:
return
pinned_conversation = (
db.session.query(PinnedConversation)
.where(
PinnedConversation.app_id == app_model.id,
PinnedConversation.conversation_id == conversation_id,
PinnedConversation.created_by_role == ("account" if isinstance(user, Account) else "end_user"),
PinnedConversation.created_by == user.id,
with session_factory.create_session() as session, session.begin():
pinned_conversation = (
session.query(PinnedConversation)
.where(
PinnedConversation.app_id == app_model.id,
PinnedConversation.conversation_id == conversation_id,
PinnedConversation.created_by_role == ("account" if isinstance(user, Account) else "end_user"),
PinnedConversation.created_by == user.id,
)
.first()
)
.first()
)

if not pinned_conversation:
return
if not pinned_conversation:
return
Comment on lines +110 to +111
Copy link

Copilot AI Feb 5, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The early return on line 111 occurs within a session.begin() transaction context. When exiting the context manager via return, an empty transaction is committed. While not incorrect, this is inefficient. Consider restructuring to avoid early returns within the transaction block, or move the transaction begin to after the early return check.

Copilot uses AI. Check for mistakes.

db.session.delete(pinned_conversation)
db.session.commit()
session.delete(pinned_conversation)
Loading