jyapayne/dify
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

feat: billing enhancement 20231204 (#1691)

Browse source 
Co-authored-by: jyong <jyong@dify.ai>
This commit is contained in:
Garfield Dai 2023-12-05 16:53:55 +08:00 • committed by GitHub
commit 7b8a10f3ea
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
12 changed files with 86 additions and 81 deletions
Download patch file Download diff file Expand all files Collapse all files

28
api/controllers/service_api/wraps.py
View file

@ -11,6 +11,7 @@ from libs.login import _get_user
from extensions.ext_database import db
from models.account import Tenant, TenantAccountJoin, Account
from models.model import ApiToken, App
from services.billing_service import BillingService
def validate_app_token(view=None):
@ -40,6 +41,33 @@ def validate_app_token(view=None):
return decorator
def cloud_edition_billing_resource_check(resource: str,
api_token_type: str,
error_msg: str = "You have reached the limit of your subscription."):
def interceptor(view):
def decorated(*args, **kwargs):
if current_app.config['EDITION'] == 'CLOUD':
api_token = validate_and_get_api_token(api_token_type)
billing_info = BillingService.get_info(api_token.tenant_id)
members = billing_info['members']
apps = billing_info['apps']
vector_space = billing_info['vector_space']
if resource == 'members' and 0 < members['limit'] <= members['size']:
raise Unauthorized(error_msg)
elif resource == 'apps' and 0 < apps['limit'] <= apps['size']:
raise Unauthorized(error_msg)
elif resource == 'vector_space' and 0 < vector_space['limit'] <= vector_space['size']:
raise Unauthorized(error_msg)
else:
return view(*args, **kwargs)
return view(*args, **kwargs)
return decorated
return interceptor
def validate_dataset_token(view=None):
def decorator(view):
@wraps(view)