refactor.

refactor.
This commit is contained in:
Garfield Dai 2024-10-19 00:54:01 +08:00
parent 9f72fdf6f7
commit 85ad6a3f22
6 changed files with 32 additions and 30 deletions

View File

@ -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

View File

@ -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")

View File

@ -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(

View File

@ -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

View File

@ -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

View File

@ -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):