diff --git a/api/controllers/console/auth/forgot_password.py b/api/controllers/console/auth/forgot_password.py index 2b12a77a2b..36f217087b 100644 --- a/api/controllers/console/auth/forgot_password.py +++ b/api/controllers/console/auth/forgot_password.py @@ -2,7 +2,7 @@ import base64 import logging import secrets -from flask import request +from flask import redirect, request from flask_restful import Resource, reqparse from configs import dify_config @@ -17,11 +17,12 @@ from controllers.console.auth.error import ( ) from controllers.console.error import NotAllowedRegister from controllers.console.setup import setup_required +from events.tenant_event import tenant_was_created from extensions.ext_database import db from libs.helper import email, get_remote_ip from libs.password import hash_password, valid_password from models.account import Account -from services.account_service import AccountService +from services.account_service import AccountService, TenantService from services.errors.account import RateLimitExceededError @@ -107,6 +108,17 @@ class ForgotPasswordResetApi(Resource): account.password = base64_password_hashed account.password_salt = base64_salt db.session.commit() + tenant = TenantService.get_join_tenants(account) + if not tenant: + if not dify_config.ALLOW_CREATE_WORKSPACE: + return redirect( + f"{dify_config.CONSOLE_WEB_URL}/signin?message=Workspace not found, please contact system admin to invite you to join in a workspace." + ) + else: + tenant = TenantService.create_tenant(f"{account.name}'s Workspace") + TenantService.create_tenant_member(tenant, account, role="owner") + account.current_tenant = tenant + tenant_was_created.send(tenant) else: account = AccountService.create_account_and_tenant( email=reset_data.get("email"), diff --git a/api/controllers/console/auth/login.py b/api/controllers/console/auth/login.py index 29bf92de44..8c2112e544 100644 --- a/api/controllers/console/auth/login.py +++ b/api/controllers/console/auth/login.py @@ -1,7 +1,7 @@ from typing import cast import flask_login -from flask import request +from flask import redirect, request from flask_restful import Resource, reqparse import services @@ -16,10 +16,12 @@ from controllers.console.auth.error import ( ) from controllers.console.error import NotAllowedRegister from controllers.console.setup import setup_required +from events.tenant_event import tenant_was_created from libs.helper import email, get_remote_ip from libs.password import valid_password from models.account import Account from services.account_service import AccountService, TenantService +from services.errors.workspace import WorkSpaceNotAllowedCreateError class LoginApi(Resource): @@ -130,11 +132,27 @@ class EmailCodeLoginApi(Resource): AccountService.revoke_email_code_login_token(args["token"]) account = AccountService.get_user_through_email(user_email) - if account is None: - account = AccountService.create_account_and_tenant( - email=user_email, name=user_email, interface_language=languages[0] - ) + tenant = TenantService.get_join_tenants(account) + if not tenant: + if not dify_config.ALLOW_CREATE_WORKSPACE: + return redirect( + f"{dify_config.CONSOLE_WEB_URL}/signin?message=Workspace not found, please contact system admin to invite you to join in a workspace." + ) + else: + tenant = TenantService.create_tenant(f"{account.name}'s Workspace") + TenantService.create_tenant_member(tenant, account, role="owner") + account.current_tenant = tenant + tenant_was_created.send(tenant) + if account is None: + try: + account = AccountService.create_account_and_tenant( + email=user_email, name=user_email, interface_language=languages[0] + ) + except WorkSpaceNotAllowedCreateError: + return redirect( + f"{dify_config.CONSOLE_WEB_URL}/signin?message=Workspace not found, please contact system admin to invite you to join in a workspace." + ) token = AccountService.login(account, ip_address=get_remote_ip(request)) return {"result": "success", "data": token} diff --git a/api/controllers/console/auth/oauth.py b/api/controllers/console/auth/oauth.py index 03e2318416..ff0407d5fb 100644 --- a/api/controllers/console/auth/oauth.py +++ b/api/controllers/console/auth/oauth.py @@ -9,6 +9,7 @@ from werkzeug.exceptions import Unauthorized from configs import dify_config from constants.languages import languages +from events.tenant_event import tenant_was_created from extensions.ext_database import db from libs.helper import get_remote_ip from libs.oauth import GitHubOAuth, GoogleOAuth, OAuthUserInfo @@ -133,7 +134,13 @@ def _generate_account(provider: str, user_info: OAuthUserInfo): if account: tenant = TenantService.get_join_tenants(account) if not tenant: - raise WorkSpaceNotFound() + if not dify_config.ALLOW_CREATE_WORKSPACE: + raise WorkSpaceNotAllowedCreateError() + else: + tenant = TenantService.create_tenant(f"{account.name}'s Workspace") + TenantService.create_tenant_member(tenant, account, role="owner") + account.current_tenant = tenant + tenant_was_created.send(tenant) if not account: if not dify_config.ALLOW_REGISTER: diff --git a/api/services/account_service.py b/api/services/account_service.py index 3ee24f0a51..d291e4a6c7 100644 --- a/api/services/account_service.py +++ b/api/services/account_service.py @@ -334,8 +334,10 @@ class TenantService: return tenant @staticmethod - def create_owner_tenant_if_not_exist(account: Account): + def create_owner_tenant_if_not_exist(account: Account, name: Optional[str] = None): """Create owner tenant if not exist""" + if not dify_config.ALLOW_CREATE_WORKSPACE: + raise WorkSpaceNotAllowedCreateError() available_ta = ( TenantAccountJoin.query.filter_by(account_id=account.id).order_by(TenantAccountJoin.id.asc()).first() )