From 1c549ddbaa39edebcabbe3ab2ea5f62338c41ebb Mon Sep 17 00:00:00 2001 From: zhangx1n Date: Wed, 9 Apr 2025 15:44:28 +0800 Subject: [PATCH] feat: add workspace/member limit errors and enforce email login restriction - Add WorkspacesLimitExceededError and WorkspaceMembersLimitExceededError - Add workspace creation limit for email login --- api/controllers/console/auth/forgot_password.py | 8 ++++++-- api/controllers/console/workspace/members.py | 2 +- api/controllers/inner_api/mail.py | 2 +- api/services/account_service.py | 4 +--- api/services/enterprise/mail_service.py | 8 +------- api/services/errors/workspace.py | 2 +- api/tasks/mail_enterprise_task.py | 8 ++------ 7 files changed, 13 insertions(+), 21 deletions(-) diff --git a/api/controllers/console/auth/forgot_password.py b/api/controllers/console/auth/forgot_password.py index a9c4cc77fd..750be197df 100644 --- a/api/controllers/console/auth/forgot_password.py +++ b/api/controllers/console/auth/forgot_password.py @@ -7,7 +7,12 @@ from flask_restful import Resource, reqparse # type: ignore from constants.languages import languages from controllers.console import api from controllers.console.auth.error import EmailCodeError, InvalidEmailError, InvalidTokenError, PasswordMismatchError -from controllers.console.error import AccountInFreezeError, AccountNotFound, EmailSendIpLimitError +from controllers.console.error import ( + AccountInFreezeError, + AccountNotFound, + EmailSendIpLimitError, + WorkspacesLimitExceeded, +) from controllers.console.wraps import setup_required from events.tenant_event import tenant_was_created from extensions.ext_database import db @@ -18,7 +23,6 @@ from services.account_service import AccountService, TenantService from services.errors.account import AccountRegisterError from services.errors.workspace import WorkSpaceNotAllowedCreateError, WorkspacesLimitExceededError from services.feature_service import FeatureService -from controllers.console.error import WorkspacesLimitExceeded class ForgotPasswordSendEmailApi(Resource): diff --git a/api/controllers/console/workspace/members.py b/api/controllers/console/workspace/members.py index 4f93406de2..56e1efde62 100644 --- a/api/controllers/console/workspace/members.py +++ b/api/controllers/console/workspace/members.py @@ -60,7 +60,7 @@ class MemberInviteEmailApi(Resource): workspace_members = FeatureService.get_features(tenant_id=inviter.current_tenant.id).workspace_members if ( FeatureService.get_system_features().license.product_id == "DIFY_ENTERPRISE_STANDARD" - and workspace_members.limit != 0 # if limit == 0, it means unlimited + and workspace_members.limit != 0 # if limit == 0, it means unlimited and len(invitee_emails) > workspace_members.limit - workspace_members.size ): raise WorkspaceMembersLimitExceeded() diff --git a/api/controllers/inner_api/mail.py b/api/controllers/inner_api/mail.py index b98356becb..47cbcb713c 100644 --- a/api/controllers/inner_api/mail.py +++ b/api/controllers/inner_api/mail.py @@ -14,7 +14,7 @@ class EnterpriseMail(Resource): @inner_api_only def post(self): parser = reqparse.RequestParser() - parser.add_argument("to", type=str, action='append', required=True) + parser.add_argument("to", type=str, action="append", required=True) parser.add_argument("subject", type=str, required=True) parser.add_argument("body", type=str, required=True) parser.add_argument("substitutions", type=dict, required=False) diff --git a/api/services/account_service.py b/api/services/account_service.py index 6ba1bc2dfb..c1b2d4f9ce 100644 --- a/api/services/account_service.py +++ b/api/services/account_service.py @@ -53,9 +53,7 @@ from services.errors.workspace import WorkSpaceNotAllowedCreateError, Workspaces from services.feature_service import FeatureService from tasks.delete_account_task import delete_account_task from tasks.mail_account_deletion_task import send_account_deletion_verification_code -from tasks.mail_email_code_login import send_email_code_login_mail_task from tasks.mail_invite_member_task import send_invite_member_mail_task -from tasks.mail_reset_password_task import send_reset_password_mail_task class TokenPair(BaseModel): @@ -590,7 +588,7 @@ class TenantService: workspaces = FeatureService.get_system_features().license.workspaces if ( FeatureService.get_system_features().license.product_id == "DIFY_ENTERPRISE_STANDARD" - and workspaces.limit != 0 # if limit == 0, it means unlimited + and workspaces.limit != 0 # if limit == 0, it means unlimited and workspaces.limit - workspaces.size <= 0 ): raise WorkspacesLimitExceededError() diff --git a/api/services/enterprise/mail_service.py b/api/services/enterprise/mail_service.py index 24b22008a1..a070a3fa60 100644 --- a/api/services/enterprise/mail_service.py +++ b/api/services/enterprise/mail_service.py @@ -1,4 +1,3 @@ - from typing import Dict, List from pydantic import BaseModel @@ -14,13 +13,8 @@ class DifyMail(BaseModel): class EnterpriseMailService: - @classmethod def send_mail(cls, mail: DifyMail): - send_enterprise_email_task.delay( - to=mail.to, - subject=mail.subject, - body=mail.body, - substitutions=mail.substitutions + to=mail.to, subject=mail.subject, body=mail.body, substitutions=mail.substitutions ) diff --git a/api/services/errors/workspace.py b/api/services/errors/workspace.py index f2973674e9..577238507f 100644 --- a/api/services/errors/workspace.py +++ b/api/services/errors/workspace.py @@ -10,4 +10,4 @@ class WorkSpaceNotFoundError(BaseServiceError): class WorkspacesLimitExceededError(BaseServiceError): - pass \ No newline at end of file + pass diff --git a/api/tasks/mail_enterprise_task.py b/api/tasks/mail_enterprise_task.py index 67475185db..b9d8fd55df 100644 --- a/api/tasks/mail_enterprise_task.py +++ b/api/tasks/mail_enterprise_task.py @@ -13,9 +13,7 @@ def send_enterprise_email_task(to, subject, body, substitutions): if not mail.is_inited(): return - logging.info( - click.style("Start enterprise mail to {} with subject {}".format(to, subject), fg="green") - ) + logging.info(click.style("Start enterprise mail to {} with subject {}".format(to, subject), fg="green")) start_at = time.perf_counter() try: @@ -29,9 +27,7 @@ def send_enterprise_email_task(to, subject, body, substitutions): end_at = time.perf_counter() logging.info( - click.style( - "Send enterprise mail to {} succeeded: latency: {}".format(to, end_at - start_at), fg="green" - ) + click.style("Send enterprise mail to {} succeeded: latency: {}".format(to, end_at - start_at), fg="green") ) except Exception: logging.exception("Send enterprise mail to {} failed".format(to))