KMS/dify
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix: add billing service health check

Browse Source
This commit is contained in:
GareArc 2025-02-26 23:45:50 -05:00
parent 10c40c4286
commit a4738d290e
2 changed files with 27 additions and 15 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

28
api/controllers/console/workspace/account.py
View File

@ -3,24 +3,20 @@ import datetime
import pytz import pytz
from flask import request from flask import request
from flask_login import current_user # type: ignore from flask_login import current_user # type: ignore
from flask_restful import Resource, fields, marshal_with, reqparse # type: ignore from flask_restful import (Resource, fields, marshal_with, # type: ignore
reqparse)
from configs import dify_config from configs import dify_config
from constants.languages import supported_language from constants.languages import supported_language
from controllers.console import api from controllers.console import api
from controllers.console.workspace.error import ( from controllers.console.workspace.error import (
AccountAlreadyInitedError, AccountAlreadyInitedError, CurrentPasswordIncorrectError,
CurrentPasswordIncorrectError, InvalidAccountDeletionCodeError, InvalidInvitationCodeError,
InvalidAccountDeletionCodeError, RepeatPasswordNotMatchError)
InvalidInvitationCodeError, from controllers.console.wraps import (account_initialization_required,
RepeatPasswordNotMatchError, cloud_edition_billing_enabled,
) enterprise_license_required,
from controllers.console.wraps import ( only_edition_cloud, setup_required)
account_initialization_required,
enterprise_license_required,
only_edition_cloud,
setup_required,
)
from extensions.ext_database import db from extensions.ext_database import db
from fields.member_fields import account_fields from fields.member_fields import account_fields
from libs.helper import TimestampField, timezone from libs.helper import TimestampField, timezone
@ -28,7 +24,8 @@ from libs.login import login_required
from models import AccountIntegrate, InvitationCode from models import AccountIntegrate, InvitationCode
from services.account_service import AccountService from services.account_service import AccountService
from services.billing_service import BillingService from services.billing_service import BillingService
from services.errors.account import CurrentPasswordIncorrectError as ServiceCurrentPasswordIncorrectError from services.errors.account import \
CurrentPasswordIncorrectError as ServiceCurrentPasswordIncorrectError
class AccountInitApi(Resource): class AccountInitApi(Resource):
@ -306,6 +303,7 @@ class EducationVerifyApi(Resource):
@login_required @login_required
@account_initialization_required @account_initialization_required
@only_edition_cloud @only_edition_cloud
@cloud_edition_billing_enabled
@marshal_with(verify_fields) @marshal_with(verify_fields)
def get(self): def get(self):
account = current_user account = current_user
@ -322,6 +320,7 @@ class EducationApi(Resource):
@login_required @login_required
@account_initialization_required @account_initialization_required
@only_edition_cloud @only_edition_cloud
@cloud_edition_billing_enabled
def post(self): def post(self):
account = current_user account = current_user
@ -335,6 +334,7 @@ class EducationApi(Resource):
@login_required @login_required
@account_initialization_required @account_initialization_required
@only_edition_cloud @only_edition_cloud
@cloud_edition_billing_enabled
@marshal_with(status_fields) @marshal_with(status_fields)
def get(self): def get(self):
account = current_user account = current_user

14
api/controllers/console/wraps.py
View File

@ -12,7 +12,8 @@ from models.model import DifySetup
from services.feature_service import FeatureService, LicenseStatus from services.feature_service import FeatureService, LicenseStatus
from services.operation_service import OperationService from services.operation_service import OperationService
from .error import NotInitValidateError, NotSetupError, UnauthorizedAndForceLogout from .error import (NotInitValidateError, NotSetupError,
UnauthorizedAndForceLogout)
def account_initialization_required(view): def account_initialization_required(view):
@ -51,6 +52,17 @@ def only_edition_self_hosted(view):
return decorated return decorated
def cloud_edition_billing_enabled(view):
@wraps(view)
def decorated(*args, **kwargs):
features = FeatureService.get_features(current_user.current_tenant_id)
if not features.billing.enabled:
abort(403, "The billing feature is not enabled.")
return view(*args, **kwargs)
return decorated
def cloud_edition_billing_resource_check(resource: str): def cloud_edition_billing_resource_check(resource: str):
def interceptor(view): def interceptor(view):
@wraps(view) @wraps(view)