fix: couldn't log in or resetup after a failed setup (#5739)

api/controllers/console/setup.py

@@ -3,11 +3,10 @@ from functools import wraps
 from flask import current_app, request
 from flask_restful import Resource, reqparse
 
-from extensions.ext_database import db
 from libs.helper import email, get_remote_ip, str_len
 from libs.password import valid_password
 from models.model import DifySetup
-from services.account_service import AccountService, RegisterService, TenantService
+from services.account_service import RegisterService, TenantService
 
 from . import api
 from .error import AlreadySetupError, NotInitValidateError, NotSetupError
@@ -51,28 +50,17 @@ class SetupApi(Resource):
                             required=True, location='json')
         args = parser.parse_args()
 
-        # Register
-        account = RegisterService.register(
+        # setup
+        RegisterService.setup(
             email=args['email'],
             name=args['name'],
-            password=args['password']
+            password=args['password'],
+            ip_address=get_remote_ip(request)
         )
 
-        TenantService.create_owner_tenant_if_not_exist(account)
-
-        setup()
-        AccountService.update_last_login(account, ip_address=get_remote_ip(request))
-
         return {'result': 'success'}, 201
 
 
-def setup():
-    dify_setup = DifySetup(
-        version=current_app.config['CURRENT_VERSION']
-    )
-    db.session.add(dify_setup)
-
-
 def setup_required(view):
     @wraps(view)
     def decorated(*args, **kwargs):

api/services/account_service.py

@@ -17,6 +17,7 @@ from libs.passport import PassportService
 from libs.password import compare_password, hash_password, valid_password
 from libs.rsa import generate_key_pair
 from models.account import *
+from models.model import DifySetup
 from services.errors.account import (
     AccountAlreadyInTenantError,
     AccountLoginError,
@@ -119,10 +120,11 @@ class AccountService:
         return account
 
     @staticmethod
-    def create_account(email: str, name: str, interface_language: str,
-                       password: str = None,
-                       interface_theme: str = 'light',
-                       timezone: str = 'America/New_York', ) -> Account:
+    def create_account(email: str,
+                       name: str,
+                       interface_language: str,
+                       password: Optional[str] = None,
+                       interface_theme: str = 'light') -> Account:
         """create account"""
         account = Account()
         account.email = email
@@ -200,7 +202,6 @@ class AccountService:
         account.last_login_ip = ip_address
         db.session.add(account)
         db.session.commit()
-        logging.info(f'Account {account.id} logged in successfully.')
 
     @staticmethod
     def login(account: Account, *, ip_address: Optional[str] = None):
@@ -444,8 +445,51 @@ class RegisterService:
         return f'member_invite:token:{token}'
 
     @classmethod
-    def register(cls, email, name, password: str = None, open_id: str = None, provider: str = None,
-                 language: str = None, status: AccountStatus = None) -> Account:
+    def setup(cls, email: str, name: str, password: str, ip_address: str) -> None:
+        """
+        Setup dify
+
+        :param email: email
+        :param name: username
+        :param password: password
+        :param ip_address: ip address
+        """
+        try:
+            # Register
+            account = AccountService.create_account(
+                email=email,
+                name=name,
+                interface_language=languages[0],
+                password=password,
+            )
+
+            account.last_login_ip = ip_address
+            account.initialized_at = datetime.now(timezone.utc).replace(tzinfo=None)
+
+            TenantService.create_owner_tenant_if_not_exist(account)
+
+            dify_setup = DifySetup(
+                version=current_app.config['CURRENT_VERSION']
+            )
+            db.session.add(dify_setup)
+            db.session.commit()
+        except Exception as e:
+            db.session.query(DifySetup).delete()
+            db.session.query(TenantAccountJoin).delete()
+            db.session.query(Account).delete()
+            db.session.query(Tenant).delete()
+            db.session.commit()
+
+            logging.exception(f'Setup failed: {e}')
+            raise ValueError(f'Setup failed: {e}')
+
+    @classmethod
+    def register(cls, email, name,
+                 password: Optional[str] = None,
+                 open_id: Optional[str] = None,
+                 provider: Optional[str] = None,
+                 language: Optional[str] = None,
+                 status: Optional[AccountStatus] = None) -> Account:
         db.session.begin_nested()
         """Register account"""
         try: