mirror of
https://github.com/langgenius/dify.git
synced 2026-02-24 03:41:21 +08:00
422 lines
18 KiB
Python
422 lines
18 KiB
Python
import logging
|
|
import time
|
|
from collections.abc import Callable
|
|
from enum import StrEnum, auto
|
|
from functools import wraps
|
|
from typing import Concatenate, ParamSpec, TypeVar
|
|
|
|
from flask import current_app, request
|
|
from flask_login import user_logged_in
|
|
from flask_restx import Resource
|
|
from pydantic import BaseModel
|
|
from sqlalchemy import select
|
|
from sqlalchemy.orm import Session
|
|
from werkzeug.exceptions import Forbidden, NotFound, Unauthorized
|
|
|
|
from enums.cloud_plan import CloudPlan
|
|
from extensions.ext_database import db
|
|
from extensions.ext_redis import redis_client
|
|
from libs.api_token_cache import ApiTokenCache
|
|
from libs.api_token_updater import update_token_last_used_at
|
|
from libs.datetime_utils import naive_utc_now
|
|
from libs.login import current_user
|
|
from models import Account, Tenant, TenantAccountJoin, TenantStatus
|
|
from models.dataset import Dataset, RateLimitLog
|
|
from models.model import ApiToken, App
|
|
from services.end_user_service import EndUserService
|
|
from services.feature_service import FeatureService
|
|
|
|
P = ParamSpec("P")
|
|
R = TypeVar("R")
|
|
T = TypeVar("T")
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class WhereisUserArg(StrEnum):
|
|
"""
|
|
Enum for whereis_user_arg.
|
|
"""
|
|
|
|
QUERY = auto()
|
|
JSON = auto()
|
|
FORM = auto()
|
|
|
|
|
|
class FetchUserArg(BaseModel):
|
|
fetch_from: WhereisUserArg
|
|
required: bool = False
|
|
|
|
|
|
def validate_app_token(view: Callable[P, R] | None = None, *, fetch_user_arg: FetchUserArg | None = None):
|
|
def decorator(view_func: Callable[P, R]):
|
|
@wraps(view_func)
|
|
def decorated_view(*args: P.args, **kwargs: P.kwargs):
|
|
api_token = validate_and_get_api_token("app")
|
|
|
|
app_model = db.session.query(App).where(App.id == api_token.app_id).first()
|
|
if not app_model:
|
|
raise Forbidden("The app no longer exists.")
|
|
|
|
if app_model.status != "normal":
|
|
raise Forbidden("The app's status is abnormal.")
|
|
|
|
if not app_model.enable_api:
|
|
raise Forbidden("The app's API service has been disabled.")
|
|
|
|
tenant = db.session.query(Tenant).where(Tenant.id == app_model.tenant_id).first()
|
|
if tenant is None:
|
|
raise ValueError("Tenant does not exist.")
|
|
if tenant.status == TenantStatus.ARCHIVE:
|
|
raise Forbidden("The workspace's status is archived.")
|
|
|
|
kwargs["app_model"] = app_model
|
|
|
|
# If caller needs end-user context, attach EndUser to current_user
|
|
if fetch_user_arg:
|
|
user_id = None
|
|
match fetch_user_arg.fetch_from:
|
|
case WhereisUserArg.QUERY:
|
|
user_id = request.args.get("user")
|
|
case WhereisUserArg.JSON:
|
|
user_id = request.get_json().get("user")
|
|
case WhereisUserArg.FORM:
|
|
user_id = request.form.get("user")
|
|
|
|
if not user_id and fetch_user_arg.required:
|
|
raise ValueError("Arg user must be provided.")
|
|
|
|
if user_id:
|
|
user_id = str(user_id)
|
|
|
|
end_user = EndUserService.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
|
|
current_app.login_manager._update_request_context_with_user(end_user) # type: ignore
|
|
user_logged_in.send(current_app._get_current_object(), user=end_user) # type: ignore
|
|
else:
|
|
# For service API without end-user context, ensure an Account is logged in
|
|
# so services relying on current_account_with_tenant() work correctly.
|
|
tenant_owner_info = (
|
|
db.session.query(Tenant, Account)
|
|
.join(TenantAccountJoin, Tenant.id == TenantAccountJoin.tenant_id)
|
|
.join(Account, TenantAccountJoin.account_id == Account.id)
|
|
.where(
|
|
Tenant.id == app_model.tenant_id,
|
|
TenantAccountJoin.role == "owner",
|
|
Tenant.status == TenantStatus.NORMAL,
|
|
)
|
|
.one_or_none()
|
|
)
|
|
|
|
if tenant_owner_info:
|
|
tenant_model, account = tenant_owner_info
|
|
account.current_tenant = tenant_model
|
|
current_app.login_manager._update_request_context_with_user(account) # type: ignore
|
|
user_logged_in.send(current_app._get_current_object(), user=current_user) # type: ignore
|
|
else:
|
|
raise Unauthorized("Tenant owner account not found or tenant is not active.")
|
|
|
|
return view_func(*args, **kwargs)
|
|
|
|
return decorated_view
|
|
|
|
if view is None:
|
|
return decorator
|
|
else:
|
|
return decorator(view)
|
|
|
|
|
|
def cloud_edition_billing_resource_check(resource: str, api_token_type: str):
|
|
def interceptor(view: Callable[P, R]):
|
|
def decorated(*args: P.args, **kwargs: P.kwargs):
|
|
api_token = validate_and_get_api_token(api_token_type)
|
|
features = FeatureService.get_features(api_token.tenant_id)
|
|
|
|
if features.billing.enabled:
|
|
members = features.members
|
|
apps = features.apps
|
|
vector_space = features.vector_space
|
|
documents_upload_quota = features.documents_upload_quota
|
|
|
|
if resource == "members" and 0 < members.limit <= members.size:
|
|
raise Forbidden("The number of members has reached the limit of your subscription.")
|
|
elif resource == "apps" and 0 < apps.limit <= apps.size:
|
|
raise Forbidden("The number of apps has reached the limit of your subscription.")
|
|
elif resource == "vector_space" and 0 < vector_space.limit <= vector_space.size:
|
|
raise Forbidden("The capacity of the vector space has reached the limit of your subscription.")
|
|
elif resource == "documents" and 0 < documents_upload_quota.limit <= documents_upload_quota.size:
|
|
raise Forbidden("The number of documents has reached the limit of your subscription.")
|
|
else:
|
|
return view(*args, **kwargs)
|
|
|
|
return view(*args, **kwargs)
|
|
|
|
return decorated
|
|
|
|
return interceptor
|
|
|
|
|
|
def cloud_edition_billing_knowledge_limit_check(resource: str, api_token_type: str):
|
|
def interceptor(view: Callable[P, R]):
|
|
@wraps(view)
|
|
def decorated(*args: P.args, **kwargs: P.kwargs):
|
|
api_token = validate_and_get_api_token(api_token_type)
|
|
features = FeatureService.get_features(api_token.tenant_id)
|
|
if features.billing.enabled:
|
|
if resource == "add_segment":
|
|
if features.billing.subscription.plan == CloudPlan.SANDBOX:
|
|
raise Forbidden(
|
|
"To unlock this feature and elevate your Dify experience, please upgrade to a paid plan."
|
|
)
|
|
else:
|
|
return view(*args, **kwargs)
|
|
|
|
return view(*args, **kwargs)
|
|
|
|
return decorated
|
|
|
|
return interceptor
|
|
|
|
|
|
def cloud_edition_billing_rate_limit_check(resource: str, api_token_type: str):
|
|
def interceptor(view: Callable[P, R]):
|
|
@wraps(view)
|
|
def decorated(*args: P.args, **kwargs: P.kwargs):
|
|
api_token = validate_and_get_api_token(api_token_type)
|
|
|
|
if resource == "knowledge":
|
|
knowledge_rate_limit = FeatureService.get_knowledge_rate_limit(api_token.tenant_id)
|
|
if knowledge_rate_limit.enabled:
|
|
current_time = int(time.time() * 1000)
|
|
key = f"rate_limit_{api_token.tenant_id}"
|
|
|
|
redis_client.zadd(key, {current_time: current_time})
|
|
|
|
redis_client.zremrangebyscore(key, 0, current_time - 60000)
|
|
|
|
request_count = redis_client.zcard(key)
|
|
|
|
if request_count > knowledge_rate_limit.limit:
|
|
# add ratelimit record
|
|
rate_limit_log = RateLimitLog(
|
|
tenant_id=api_token.tenant_id,
|
|
subscription_plan=knowledge_rate_limit.subscription_plan,
|
|
operation="knowledge",
|
|
)
|
|
db.session.add(rate_limit_log)
|
|
db.session.commit()
|
|
raise Forbidden(
|
|
"Sorry, you have reached the knowledge base request rate limit of your subscription."
|
|
)
|
|
return view(*args, **kwargs)
|
|
|
|
return decorated
|
|
|
|
return interceptor
|
|
|
|
|
|
def validate_dataset_token(view: Callable[Concatenate[T, P], R] | None = None):
|
|
def decorator(view: Callable[Concatenate[T, P], R]):
|
|
@wraps(view)
|
|
def decorated(*args: P.args, **kwargs: P.kwargs):
|
|
# get url path dataset_id from positional args or kwargs
|
|
# Flask passes URL path parameters as positional arguments
|
|
dataset_id = None
|
|
|
|
# First try to get from kwargs (explicit parameter)
|
|
dataset_id = kwargs.get("dataset_id")
|
|
|
|
# If not in kwargs, try to extract from positional args
|
|
if not dataset_id and args:
|
|
# For class methods: args[0] is self, args[1] is dataset_id (if exists)
|
|
# Check if first arg is likely a class instance (has __dict__ or __class__)
|
|
if len(args) > 1 and hasattr(args[0], "__dict__"):
|
|
# This is a class method, dataset_id should be in args[1]
|
|
potential_id = args[1]
|
|
# Validate it's a string-like UUID, not another object
|
|
try:
|
|
# Try to convert to string and check if it's a valid UUID format
|
|
str_id = str(potential_id)
|
|
# Basic check: UUIDs are 36 chars with hyphens
|
|
if len(str_id) == 36 and str_id.count("-") == 4:
|
|
dataset_id = str_id
|
|
except Exception:
|
|
logger.exception("Failed to parse dataset_id from class method args")
|
|
elif len(args) > 0:
|
|
# Not a class method, check if args[0] looks like a UUID
|
|
potential_id = args[0]
|
|
try:
|
|
str_id = str(potential_id)
|
|
if len(str_id) == 36 and str_id.count("-") == 4:
|
|
dataset_id = str_id
|
|
except Exception:
|
|
logger.exception("Failed to parse dataset_id from positional args")
|
|
|
|
# Validate dataset if dataset_id is provided
|
|
if dataset_id:
|
|
dataset_id = str(dataset_id)
|
|
dataset = db.session.query(Dataset).where(Dataset.id == dataset_id).first()
|
|
if not dataset:
|
|
raise NotFound("Dataset not found.")
|
|
if not dataset.enable_api:
|
|
raise Forbidden("Dataset api access is not enabled.")
|
|
api_token = validate_and_get_api_token("dataset")
|
|
tenant_account_join = (
|
|
db.session.query(Tenant, TenantAccountJoin)
|
|
.where(Tenant.id == api_token.tenant_id)
|
|
.where(TenantAccountJoin.tenant_id == Tenant.id)
|
|
.where(TenantAccountJoin.role.in_(["owner"]))
|
|
.where(Tenant.status == TenantStatus.NORMAL)
|
|
.one_or_none()
|
|
) # TODO: only owner information is required, so only one is returned.
|
|
if tenant_account_join:
|
|
tenant, ta = tenant_account_join
|
|
account = db.session.query(Account).where(Account.id == ta.account_id).first()
|
|
# Login admin
|
|
if account:
|
|
account.current_tenant = tenant
|
|
current_app.login_manager._update_request_context_with_user(account) # type: ignore
|
|
user_logged_in.send(current_app._get_current_object(), user=current_user) # type: ignore
|
|
else:
|
|
raise Unauthorized("Tenant owner account does not exist.")
|
|
else:
|
|
raise Unauthorized("Tenant does not exist.")
|
|
return view(api_token.tenant_id, *args, **kwargs)
|
|
|
|
return decorated
|
|
|
|
if view:
|
|
return decorator(view)
|
|
|
|
# if view is None, it means that the decorator is used without parentheses
|
|
# use the decorator as a function for method_decorators
|
|
return decorator
|
|
|
|
|
|
def validate_and_get_api_token(scope: str | None = None):
|
|
"""
|
|
Validate and get API token with Redis caching.
|
|
|
|
This function uses a two-tier approach:
|
|
1. First checks Redis cache for the token
|
|
2. If not cached, queries database and caches the result
|
|
|
|
The last_used_at field is updated asynchronously via Celery task
|
|
to avoid blocking the request.
|
|
"""
|
|
auth_header = request.headers.get("Authorization")
|
|
if auth_header is None or " " not in auth_header:
|
|
raise Unauthorized("Authorization header must be provided and start with 'Bearer'")
|
|
|
|
auth_scheme, auth_token = auth_header.split(None, 1)
|
|
auth_scheme = auth_scheme.lower()
|
|
|
|
if auth_scheme != "bearer":
|
|
raise Unauthorized("Authorization scheme must be 'Bearer'")
|
|
|
|
# Try to get token from cache first
|
|
# Returns a CachedApiToken (plain Python object), not a SQLAlchemy model
|
|
cached_token = ApiTokenCache.get(auth_token, scope)
|
|
if cached_token is not None:
|
|
logger.debug("Token validation served from cache for scope: %s", scope)
|
|
# Asynchronously update last_used_at (non-blocking)
|
|
_async_update_token_last_used_at(auth_token, scope)
|
|
return cached_token
|
|
|
|
# Cache miss - use Redis lock for single-flight mode
|
|
# This ensures only one request queries DB for the same token concurrently
|
|
return _fetch_token_with_single_flight(auth_token, scope)
|
|
|
|
|
|
def _query_token_from_db(auth_token: str, scope: str | None) -> ApiToken:
|
|
"""
|
|
Query API token from database, update last_used_at, and cache the result.
|
|
|
|
Raises Unauthorized if token is invalid.
|
|
"""
|
|
with Session(db.engine, expire_on_commit=False) as session:
|
|
current_time = naive_utc_now()
|
|
update_token_last_used_at(auth_token, scope, current_time, session=session)
|
|
|
|
stmt = select(ApiToken).where(ApiToken.token == auth_token, ApiToken.type == scope)
|
|
api_token = session.scalar(stmt)
|
|
|
|
if not api_token:
|
|
ApiTokenCache.set(auth_token, scope, None)
|
|
raise Unauthorized("Access token is invalid")
|
|
|
|
ApiTokenCache.set(auth_token, scope, api_token)
|
|
return api_token
|
|
|
|
|
|
def _fetch_token_with_single_flight(auth_token: str, scope: str | None) -> ApiToken:
|
|
"""
|
|
Fetch token from DB with single-flight pattern using Redis lock.
|
|
|
|
Ensures only one concurrent request queries the database for the same token.
|
|
Falls back to direct query if lock acquisition fails.
|
|
"""
|
|
logger.debug("Token cache miss, attempting to acquire query lock for scope: %s", scope)
|
|
|
|
lock_key = f"api_token_query_lock:{scope}:{auth_token}"
|
|
lock = redis_client.lock(lock_key, timeout=10, blocking_timeout=5)
|
|
|
|
try:
|
|
if lock.acquire(blocking=True):
|
|
try:
|
|
# Double-check cache after acquiring lock
|
|
# (another concurrent request might have already cached it)
|
|
cached_token = ApiTokenCache.get(auth_token, scope)
|
|
if cached_token is not None:
|
|
logger.debug("Token cached by concurrent request, using cached version")
|
|
return cached_token
|
|
|
|
return _query_token_from_db(auth_token, scope)
|
|
finally:
|
|
lock.release()
|
|
else:
|
|
logger.warning("Lock timeout for token: %s, proceeding with direct query", auth_token[:10])
|
|
return _query_token_from_db(auth_token, scope)
|
|
except Unauthorized:
|
|
raise
|
|
except Exception as e:
|
|
logger.warning("Redis lock failed for token query: %s, proceeding anyway", e)
|
|
return _query_token_from_db(auth_token, scope)
|
|
|
|
|
|
def _async_update_token_last_used_at(auth_token: str, scope: str | None):
|
|
"""
|
|
Asynchronously update the last_used_at timestamp for a token.
|
|
|
|
This schedules a Celery task to update the database without blocking
|
|
the current request. The update time is passed to ensure only older
|
|
records are updated, providing natural concurrency control.
|
|
"""
|
|
try:
|
|
from tasks.update_api_token_last_used_task import update_api_token_last_used_task
|
|
|
|
# Record the update time for concurrency control
|
|
update_time = naive_utc_now()
|
|
update_time_iso = update_time.isoformat()
|
|
|
|
# Fire and forget - don't wait for result
|
|
update_api_token_last_used_task.delay(auth_token, scope, update_time_iso)
|
|
logger.debug("Scheduled async update for last_used_at (scope: %s, update_time: %s)", scope, update_time_iso)
|
|
except Exception as e:
|
|
# Don't fail the request if task scheduling fails
|
|
logger.warning("Failed to schedule last_used_at update task: %s", e)
|
|
|
|
|
|
class DatasetApiResource(Resource):
|
|
method_decorators = [validate_dataset_token]
|
|
|
|
def get_dataset(self, dataset_id: str, tenant_id: str) -> Dataset:
|
|
dataset = db.session.query(Dataset).where(Dataset.id == dataset_id, Dataset.tenant_id == tenant_id).first()
|
|
|
|
if not dataset:
|
|
raise NotFound("Dataset not found.")
|
|
|
|
return dataset
|