Merge branch 'feat/new-login' into deploy/dev

* feat/new-login: fix: oauth error when not allowed create workspace
langgenius · Sep 9, 2024 · c2b3179 · c2b3179
2 parents d09a2a5 + 4893631
commit c2b3179
Show file tree

Hide file tree

Showing 3 changed files with 18 additions and 10 deletions.
diff --git a/api/controllers/console/auth/oauth.py b/api/controllers/console/auth/oauth.py
@@ -15,6 +15,7 @@
 from models.account import Account, AccountStatus
 from services.account_service import AccountService, RegisterService, TenantService
 from services.errors.account import AccountNotFound
+from services.errors.workspace import WorkSpaceNotAllowedCreateError
 
 from .. import api
 
@@ -90,6 +91,10 @@ def get(self, provider: str):
             account = _generate_account(provider, user_info)
         except AccountNotFound:
             return redirect(f"{dify_config.CONSOLE_WEB_URL}/signin?message=AccountNotFound")
+        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."
+            )
 
         # Check account status
         if account.status == AccountStatus.BANNED.value or account.status == AccountStatus.CLOSED.value:

diff --git a/api/services/account_service.py b/api/services/account_service.py
@@ -37,6 +37,7 @@
     RoleAlreadyAssignedError,
     TenantNotFound,
 )
+from services.errors.workspace import WorkSpaceNotAllowedCreateError
 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
@@ -607,7 +608,6 @@ def register(
         provider: Optional[str] = None,
         language: Optional[str] = None,
         status: Optional[AccountStatus] = None,
-        is_invite_member: Optional[bool] = False,
     ) -> Account:
         db.session.begin_nested()
         """Register account"""
@@ -621,13 +621,13 @@ def register(
             if open_id is not None or provider is not None:
                 AccountService.link_account_integrate(provider, open_id, account)
 
-            should_create_workspace = not is_invite_member or (is_invite_member and dify_config.ALLOW_CREATE_WORKSPACE)
+            if not dify_config.ALLOW_CREATE_WORKSPACE:
+                raise WorkSpaceNotAllowedCreateError()
 
-            if should_create_workspace:
-                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)
+            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)
 
             db.session.commit()
         except Exception as e:
@@ -648,9 +648,7 @@ def invite_new_member(
             TenantService.check_member_permission(tenant, inviter, None, "add")
             name = email.split("@")[0]
 
-            account = cls.register(
-                email=email, name=name, language=language, status=AccountStatus.PENDING, is_invite_member=True
-            )
+            account = cls.register(email=email, name=name, language=language, status=AccountStatus.PENDING)
             # Create new tenant member for invited tenant
             TenantService.create_tenant_member(tenant, account, role)
             TenantService.switch_tenant(account, tenant.id)

diff --git a/api/services/errors/workspace.py b/api/services/errors/workspace.py
@@ -0,0 +1,5 @@
+from services.errors.base import BaseServiceError
+
+
+class WorkSpaceNotAllowedCreateError(BaseServiceError):
+    pass