mirror of
https://github.com/langgenius/dify.git
synced 2026-05-03 17:08:03 +08:00
fix(trigger): enhance error handling and refactor end user creation in trigger workflows
- Improved error handling in `TriggerSubscriptionListApi` to return a 404 response for ValueErrors. - Refactored end user creation logic in `service_api/wraps.py` to use `get_or_create_end_user` for better clarity and consistency. - Introduced a new method `create_end_user_batch` for batch creation of end users, optimizing database interactions. - Updated various trigger-related services to utilize the new end user handling, ensuring proper user context during trigger dispatching.
This commit is contained in:
@ -1,5 +1,5 @@
|
||||
import time
|
||||
from collections.abc import Callable
|
||||
from collections.abc import Callable, Mapping
|
||||
from datetime import timedelta
|
||||
from enum import StrEnum, auto
|
||||
from functools import wraps
|
||||
@ -13,6 +13,7 @@ from sqlalchemy import select, update
|
||||
from sqlalchemy.orm import Session
|
||||
from werkzeug.exceptions import Forbidden, NotFound, Unauthorized
|
||||
|
||||
from core.app.entities.app_invoke_entities import InvokeFrom
|
||||
from extensions.ext_database import db
|
||||
from extensions.ext_redis import redis_client
|
||||
from libs.datetime_utils import naive_utc_now
|
||||
@ -83,7 +84,7 @@ def validate_app_token(view: Callable[P, R] | None = None, *, fetch_user_arg: Fe
|
||||
if user_id:
|
||||
user_id = str(user_id)
|
||||
|
||||
end_user = create_or_update_end_user_for_user_id(app_model, user_id)
|
||||
end_user = get_or_create_end_user(app_model, user_id)
|
||||
kwargs["end_user"] = end_user
|
||||
|
||||
# Set EndUser as current logged-in user for flask_login.current_user
|
||||
@ -308,10 +309,13 @@ def validate_and_get_api_token(scope: str | None = None):
|
||||
return api_token
|
||||
|
||||
|
||||
def create_or_update_end_user_for_user_id(app_model: App, user_id: str | None = None) -> EndUser:
|
||||
"""
|
||||
Create or update session terminal based on user ID.
|
||||
"""
|
||||
def get_or_create_end_user(app_model: App, user_id: str | None = None) -> EndUser:
|
||||
return get_or_create_end_user_by_type(InvokeFrom.SERVICE_API, app_model.tenant_id, app_model.id, user_id)
|
||||
|
||||
|
||||
def get_or_create_end_user_by_type(
|
||||
type: InvokeFrom, tenant_id: str, app_id: str, user_id: str | None = None
|
||||
) -> EndUser:
|
||||
if not user_id:
|
||||
user_id = DefaultEndUserSessionID.DEFAULT_SESSION_ID
|
||||
|
||||
@ -319,21 +323,22 @@ def create_or_update_end_user_for_user_id(app_model: App, user_id: str | None =
|
||||
end_user = (
|
||||
session.query(EndUser)
|
||||
.where(
|
||||
EndUser.tenant_id == app_model.tenant_id,
|
||||
EndUser.app_id == app_model.id,
|
||||
EndUser.tenant_id == tenant_id,
|
||||
EndUser.app_id == app_id,
|
||||
EndUser.session_id == user_id,
|
||||
EndUser.type == "service_api",
|
||||
EndUser.type == type,
|
||||
)
|
||||
.first()
|
||||
)
|
||||
|
||||
if end_user is None:
|
||||
end_user = EndUser(
|
||||
tenant_id=app_model.tenant_id,
|
||||
app_id=app_model.id,
|
||||
type="service_api",
|
||||
tenant_id=tenant_id,
|
||||
app_id=app_id,
|
||||
type=type,
|
||||
is_anonymous=user_id == DefaultEndUserSessionID.DEFAULT_SESSION_ID,
|
||||
session_id=user_id,
|
||||
external_user_id=user_id,
|
||||
)
|
||||
session.add(end_user)
|
||||
session.commit()
|
||||
@ -341,6 +346,87 @@ def create_or_update_end_user_for_user_id(app_model: App, user_id: str | None =
|
||||
return end_user
|
||||
|
||||
|
||||
def create_end_user_batch(type: InvokeFrom, tenant_id: str, app_ids: list[str], user_id: str) -> Mapping[str, EndUser]:
|
||||
"""Create end users in batch.
|
||||
|
||||
Creates end users in batch for the specified tenant and application IDs in O(1) time.
|
||||
|
||||
This batch creation is necessary because trigger subscriptions can span multiple applications,
|
||||
and trigger events may be dispatched to multiple applications simultaneously.
|
||||
|
||||
For each app_id in app_ids, check if an `EndUser` with the given
|
||||
`user_id` (as session_id/external_user_id) already exists for the
|
||||
tenant/app and type `type`. If it exists, return it; otherwise,
|
||||
create it. Operates with minimal DB I/O by querying and inserting in
|
||||
batches.
|
||||
|
||||
Returns a mapping of `app_id -> EndUser`.
|
||||
"""
|
||||
|
||||
# Normalize user_id to default if empty
|
||||
if not user_id:
|
||||
user_id = DefaultEndUserSessionID.DEFAULT_SESSION_ID
|
||||
|
||||
# Deduplicate app_ids while preserving input order
|
||||
seen: set[str] = set()
|
||||
unique_app_ids: list[str] = []
|
||||
for app_id in app_ids:
|
||||
if app_id not in seen:
|
||||
seen.add(app_id)
|
||||
unique_app_ids.append(app_id)
|
||||
|
||||
# Result is a simple app_id -> EndUser mapping
|
||||
result: dict[str, EndUser] = {}
|
||||
if not unique_app_ids:
|
||||
return result
|
||||
|
||||
with Session(db.engine, expire_on_commit=False) as session:
|
||||
# Fetch existing end users for all target apps in a single query
|
||||
existing_end_users: list[EndUser] = (
|
||||
session.query(EndUser)
|
||||
.where(
|
||||
EndUser.tenant_id == tenant_id,
|
||||
EndUser.app_id.in_(unique_app_ids),
|
||||
EndUser.session_id == user_id,
|
||||
EndUser.type == type,
|
||||
)
|
||||
.all()
|
||||
)
|
||||
|
||||
found_app_ids: set[str] = set()
|
||||
for eu in existing_end_users:
|
||||
# If duplicates exist due to weak DB constraints, prefer the first
|
||||
if eu.app_id not in result:
|
||||
result[eu.app_id] = eu
|
||||
found_app_ids.add(eu.app_id)
|
||||
|
||||
# Determine which apps still need an EndUser created
|
||||
missing_app_ids = [app_id for app_id in unique_app_ids if app_id not in found_app_ids]
|
||||
|
||||
if missing_app_ids:
|
||||
new_end_users: list[EndUser] = []
|
||||
is_anonymous = user_id == DefaultEndUserSessionID.DEFAULT_SESSION_ID
|
||||
for app_id in missing_app_ids:
|
||||
new_end_users.append(
|
||||
EndUser(
|
||||
tenant_id=tenant_id,
|
||||
app_id=app_id,
|
||||
type=type,
|
||||
is_anonymous=is_anonymous,
|
||||
session_id=user_id,
|
||||
external_user_id=user_id,
|
||||
)
|
||||
)
|
||||
|
||||
session.add_all(new_end_users)
|
||||
session.commit()
|
||||
|
||||
for eu in new_end_users:
|
||||
result[eu.app_id] = eu
|
||||
|
||||
return result
|
||||
|
||||
|
||||
class DatasetApiResource(Resource):
|
||||
method_decorators = [validate_dataset_token]
|
||||
|
||||
|
||||
Reference in New Issue
Block a user