From 85ad6a3f2278733b9f5ac80718abcf1d6e3902d5 Mon Sep 17 00:00:00 2001 From: Garfield Dai Date: Sat, 19 Oct 2024 00:54:01 +0800 Subject: [PATCH] refactor. refactor. --- .../console/auth/forgot_password.py | 4 +-- api/controllers/console/auth/login.py | 12 ++++---- api/controllers/console/auth/oauth.py | 4 +-- api/controllers/console/error.py | 6 ++-- api/services/account_service.py | 8 +++--- api/services/feature_service.py | 28 ++++++++++--------- 6 files changed, 32 insertions(+), 30 deletions(-) diff --git a/api/controllers/console/auth/forgot_password.py b/api/controllers/console/auth/forgot_password.py index c01d1c82e9..7fea610610 100644 --- a/api/controllers/console/auth/forgot_password.py +++ b/api/controllers/console/auth/forgot_password.py @@ -44,7 +44,7 @@ class ForgotPasswordSendEmailApi(Resource): account = Account.query.filter_by(email=args["email"]).first() token = None if account is None: - if FeatureService.system_features.is_allow_register: + if FeatureService.get_system_features().is_allow_register: token = AccountService.send_reset_password_email(email=args["email"], language=language) return {"result": "fail", "data": token, "code": "account_not_found"} else: @@ -114,7 +114,7 @@ class ForgotPasswordResetApi(Resource): account.password_salt = base64_salt db.session.commit() tenant = TenantService.get_join_tenants(account) - if not tenant and not FeatureService.system_features.is_allow_create_workspace: + if not tenant and not FeatureService.get_system_features().is_allow_create_workspace: tenant = TenantService.create_tenant(f"{account.name}'s Workspace") TenantService.create_tenant_member(tenant, account, role="owner") account.current_tenant = tenant diff --git a/api/controllers/console/auth/login.py b/api/controllers/console/auth/login.py index 0bfa5d72ce..4821c543b7 100644 --- a/api/controllers/console/auth/login.py +++ b/api/controllers/console/auth/login.py @@ -16,7 +16,7 @@ from controllers.console.auth.error import ( InvalidTokenError, ) from controllers.console.error import ( - AccountBannedOrClosedError, + AccountBannedError, EmailSendIpLimitError, NotAllowedCreateWorkspace, NotAllowedRegister, @@ -68,12 +68,12 @@ class LoginApi(Resource): else: account = AccountService.authenticate(args["email"], args["password"]) except services.errors.account.AccountLoginError: - raise AccountBannedOrClosedError() + raise AccountBannedError() except services.errors.account.AccountPasswordError: AccountService.add_login_error_rate_limit(args["email"]) raise EmailOrPasswordMismatchError() except services.errors.account.AccountNotFoundError: - if FeatureService.system_features.is_allow_register: + if FeatureService.get_system_features().is_allow_register: token = AccountService.send_reset_password_email(email=args["email"], language=language) return {"result": "fail", "data": token, "code": "account_not_found"} else: @@ -117,7 +117,7 @@ class ResetPasswordSendEmailApi(Resource): account = AccountService.get_user_through_email(args["email"]) if account is None: - if FeatureService.system_features.is_allow_register: + if FeatureService.get_system_features().is_allow_register: token = AccountService.send_reset_password_email(email=args["email"], language=language) else: raise NotAllowedRegister() @@ -146,7 +146,7 @@ class EmailCodeLoginSendEmailApi(Resource): account = AccountService.get_user_through_email(args["email"]) if account is None: - if FeatureService.system_features.is_allow_register: + if FeatureService.get_system_features().is_allow_register: token = AccountService.send_email_code_login_email(email=args["email"], language=language) else: raise NotAllowedRegister() @@ -182,7 +182,7 @@ class EmailCodeLoginApi(Resource): if account: tenant = TenantService.get_join_tenants(account) if not tenant: - if not FeatureService.system_features.is_allow_create_workspace: + if not FeatureService.get_system_features().is_allow_create_workspace: raise NotAllowedCreateWorkspace() else: tenant = TenantService.create_tenant(f"{account.name}'s Workspace") diff --git a/api/controllers/console/auth/oauth.py b/api/controllers/console/auth/oauth.py index 77b46bba98..ba0e07cd16 100644 --- a/api/controllers/console/auth/oauth.py +++ b/api/controllers/console/auth/oauth.py @@ -146,7 +146,7 @@ def _generate_account(provider: str, user_info: OAuthUserInfo): if account: tenant = TenantService.get_join_tenants(account) if not tenant: - if not FeatureService.system_features.is_allow_create_workspace: + if not FeatureService.get_system_features().is_allow_create_workspace: raise WorkSpaceNotAllowedCreateError() else: tenant = TenantService.create_tenant(f"{account.name}'s Workspace") @@ -155,7 +155,7 @@ def _generate_account(provider: str, user_info: OAuthUserInfo): tenant_was_created.send(tenant) if not account: - if not FeatureService.system_features.is_allow_register: + if not FeatureService.get_system_features().is_allow_register: raise AccountNotFoundError() account_name = user_info.name or "Dify" account = RegisterService.register( diff --git a/api/controllers/console/error.py b/api/controllers/console/error.py index 8daa781d66..a6d4c8e8ec 100644 --- a/api/controllers/console/error.py +++ b/api/controllers/console/error.py @@ -46,9 +46,9 @@ class NotAllowedCreateWorkspace(BaseHTTPException): code = 400 -class AccountBannedOrClosedError(BaseHTTPException): - error_code = "account_banned_or_closed" - description = "Account is banned or closed." +class AccountBannedError(BaseHTTPException): + error_code = "account_banned" + description = "Account is banned." code = 400 diff --git a/api/services/account_service.py b/api/services/account_service.py index e3a5c0f1fc..529b716773 100644 --- a/api/services/account_service.py +++ b/api/services/account_service.py @@ -197,7 +197,7 @@ class AccountService: is_setup: Optional[bool] = False, ) -> Account: """create account""" - if not FeatureService.system_features.is_allow_register and not is_setup: + if not FeatureService.get_system_features().is_allow_register and not is_setup: from controllers.console.error import NotAllowedRegister raise NotAllowedRegister() @@ -488,7 +488,7 @@ class TenantService: @staticmethod def create_tenant(name: str, is_setup: Optional[bool] = False) -> Tenant: """Create tenant""" - if not FeatureService.system_features.is_allow_create_workspace and not is_setup: + if not FeatureService.get_system_features().is_allow_create_workspace and not is_setup: from controllers.console.error import NotAllowedCreateWorkspace raise NotAllowedCreateWorkspace() @@ -506,7 +506,7 @@ class TenantService: account: Account, name: Optional[str] = None, is_setup: Optional[bool] = False ): """Create owner tenant if not exist""" - if not FeatureService.system_features.is_allow_create_workspace and not is_setup: + if not FeatureService.get_system_features().is_allow_create_workspace and not is_setup: raise WorkSpaceNotAllowedCreateError() available_ta = ( TenantAccountJoin.query.filter_by(account_id=account.id).order_by(TenantAccountJoin.id.asc()).first() @@ -804,7 +804,7 @@ class RegisterService: if open_id is not None or provider is not None: AccountService.link_account_integrate(provider, open_id, account) - if FeatureService.system_features.is_allow_create_workspace: + if FeatureService.get_system_features().is_allow_create_workspace: tenant = TenantService.create_tenant(f"{account.name}'s Workspace") TenantService.create_tenant_member(tenant, account, role="owner") account.current_tenant = tenant diff --git a/api/services/feature_service.py b/api/services/feature_service.py index e85b15919f..c321393bc5 100644 --- a/api/services/feature_service.py +++ b/api/services/feature_service.py @@ -50,8 +50,6 @@ class SystemFeatureModel(BaseModel): class FeatureService: - system_features = SystemFeatureModel() - @classmethod def get_features(cls, tenant_id: str) -> FeatureModel: features = FeatureModel() @@ -65,20 +63,24 @@ class FeatureService: @classmethod def get_system_features(cls) -> SystemFeatureModel: - cls._fulfill_login_params_from_env(cls.system_features) - if dify_config.ENTERPRISE_ENABLED: - cls.system_features.enable_web_sso_switch_component = True - cls._fulfill_params_from_enterprise(cls.system_features) + system_features = SystemFeatureModel() - return cls.system_features + cls._fulfill_system_params_from_env(system_features) + + if dify_config.ENTERPRISE_ENABLED: + system_features.enable_web_sso_switch_component = True + + cls._fulfill_params_from_enterprise(system_features) + + return system_features @classmethod - def _fulfill_login_params_from_env(cls, features: SystemFeatureModel): - features.enable_email_code_login = dify_config.ENABLE_EMAIL_CODE_LOGIN - features.enable_email_password_login = dify_config.ENABLE_EMAIL_PASSWORD_LOGIN - features.enable_social_oauth_login = dify_config.ENABLE_SOCIAL_OAUTH_LOGIN - features.is_allow_register = dify_config.ALLOW_REGISTER - features.is_allow_create_workspace = dify_config.ALLOW_CREATE_WORKSPACE + def _fulfill_system_params_from_env(cls, system_features: SystemFeatureModel): + system_features.enable_email_code_login = dify_config.ENABLE_EMAIL_CODE_LOGIN + system_features.enable_email_password_login = dify_config.ENABLE_EMAIL_PASSWORD_LOGIN + system_features.enable_social_oauth_login = dify_config.ENABLE_SOCIAL_OAUTH_LOGIN + system_features.is_allow_register = dify_config.ALLOW_REGISTER + system_features.is_allow_create_workspace = dify_config.ALLOW_CREATE_WORKSPACE @classmethod def _fulfill_params_from_env(cls, features: FeatureModel):