chore: Minor refactor for sign up workspace create flow.

app/server/appsmith-server/src/main/java/com/appsmith/server/authentication/handlers/AuthenticationSuccessHandler.java

@@ -2,19 +2,17 @@
 
 import com.appsmith.server.authentication.handlers.ce.AuthenticationSuccessHandlerCE;
 import com.appsmith.server.helpers.RedirectHelper;
+import com.appsmith.server.helpers.UserSignupHelper;
 import com.appsmith.server.helpers.WorkspaceServiceHelper;
 import com.appsmith.server.instanceconfigs.helpers.InstanceVariablesHelper;
 import com.appsmith.server.ratelimiting.RateLimitService;
 import com.appsmith.server.repositories.UserRepository;
-import com.appsmith.server.repositories.WorkspaceRepository;
 import com.appsmith.server.services.AnalyticsService;
 import com.appsmith.server.services.ApplicationPageService;
 import com.appsmith.server.services.OrganizationService;
 import com.appsmith.server.services.SessionUserService;
 import com.appsmith.server.services.UserDataService;
 import com.appsmith.server.services.UserService;
-import com.appsmith.server.services.WorkspaceService;
-import com.appsmith.server.solutions.WorkspacePermission;
 import lombok.extern.slf4j.Slf4j;
 import org.springframework.stereotype.Component;
 
@@ -28,29 +26,25 @@ public AuthenticationSuccessHandler(
             AnalyticsService analyticsService,
             UserDataService userDataService,
             UserRepository userRepository,
-            WorkspaceRepository workspaceRepository,
-            WorkspaceService workspaceService,
             ApplicationPageService applicationPageService,
-            WorkspacePermission workspacePermission,
             RateLimitService rateLimitService,
             OrganizationService organizationService,
             UserService userService,
             WorkspaceServiceHelper workspaceServiceHelper,
-            InstanceVariablesHelper instanceVariablesHelper) {
+            InstanceVariablesHelper instanceVariablesHelper,
+            UserSignupHelper userSignupHelper) {
         super(
                 redirectHelper,
                 sessionUserService,
                 analyticsService,
                 userDataService,
                 userRepository,
-                workspaceRepository,
-                workspaceService,
                 applicationPageService,
-                workspacePermission,
                 rateLimitService,
                 organizationService,
                 userService,
                 workspaceServiceHelper,
-                instanceVariablesHelper);
+                instanceVariablesHelper,
+                userSignupHelper);
     }
 }

app/server/appsmith-server/src/main/java/com/appsmith/server/authentication/handlers/ce/AuthenticationSuccessHandlerCE.java

@@ -8,22 +8,19 @@
 import com.appsmith.server.domains.Application;
 import com.appsmith.server.domains.LoginSource;
 import com.appsmith.server.domains.User;
-import com.appsmith.server.domains.Workspace;
 import com.appsmith.server.dtos.ResendEmailVerificationDTO;
 import com.appsmith.server.helpers.RedirectHelper;
+import com.appsmith.server.helpers.UserSignupHelper;
 import com.appsmith.server.helpers.WorkspaceServiceHelper;
 import com.appsmith.server.instanceconfigs.helpers.InstanceVariablesHelper;
 import com.appsmith.server.ratelimiting.RateLimitService;
 import com.appsmith.server.repositories.UserRepository;
-import com.appsmith.server.repositories.WorkspaceRepository;
 import com.appsmith.server.services.AnalyticsService;
 import com.appsmith.server.services.ApplicationPageService;
 import com.appsmith.server.services.OrganizationService;
 import com.appsmith.server.services.SessionUserService;
 import com.appsmith.server.services.UserDataService;
 import com.appsmith.server.services.UserService;
-import com.appsmith.server.services.WorkspaceService;
-import com.appsmith.server.solutions.WorkspacePermission;
 import lombok.RequiredArgsConstructor;
 import lombok.extern.slf4j.Slf4j;
 import org.springframework.http.HttpHeaders;
@@ -60,15 +57,13 @@ public class AuthenticationSuccessHandlerCE implements ServerAuthenticationSucce
     private final AnalyticsService analyticsService;
     private final UserDataService userDataService;
     private final UserRepository userRepository;
-    private final WorkspaceRepository workspaceRepository;
-    private final WorkspaceService workspaceService;
     private final ApplicationPageService applicationPageService;
-    private final WorkspacePermission workspacePermission;
     private final RateLimitService rateLimitService;
     private final OrganizationService organizationService;
     private final UserService userService;
     private final WorkspaceServiceHelper workspaceServiceHelper;
     private final InstanceVariablesHelper instanceVariablesHelper;
+    private final UserSignupHelper userSignupHelper;
 
     private Mono<Boolean> isVerificationRequired(String userEmail, String method) {
         Mono<Boolean> emailVerificationEnabledMono =
@@ -194,7 +189,7 @@ private Mono<Void> postVerificationRequiredHandler(
      * then redirects the user to /verificationPending and sends the magic link with the user's redirectUrl
      * in the email.
      */
-    private Mono<Void> formEmailVerificationRedirectionHandler(
+    public Mono<Void> formEmailVerificationRedirectionHandler(
             WebFilterExchange webFilterExchange,
             String defaultWorkspaceId,
             Authentication authentication,
@@ -350,45 +345,9 @@ public Mono<Void> onAuthenticationSuccess(
 
     protected Mono<Application> createDefaultApplication(String defaultWorkspaceId, Authentication authentication) {
         // need to create default application
-        return createWorkspaceIfNotExistsAndGetId(defaultWorkspaceId, authentication)
-                .flatMap(this::createFirstApplication);
-    }
-
-    protected Mono<String> createWorkspaceIfNotExistsAndGetId(
-            String defaultWorkspaceId, Authentication authentication) {
-        if (defaultWorkspaceId == null) {
-            return workspaceRepository
-                    .findAll(workspacePermission.getEditPermission())
-                    .take(1, true)
-                    .collectList()
-                    .flatMap(workspaces -> {
-                        // Since this is the first application creation, the first workspace would be the only
-                        // workspace user has access to, and would be user's default workspace. Hence, we use this
-                        // workspace to create the application.
-                        if (workspaces.size() == 1) {
-                            return Mono.just(workspaces.get(0));
-                        }
-
-                        // In case no workspaces are found for the user, create a new default workspace
-                        String email = ((User) authentication.getPrincipal()).getEmail();
-
-                        return organizationService
-                                .getCurrentUserOrganizationId()
-                                .flatMap(orgId -> userRepository.findByEmailAndOrganizationId(email, orgId))
-                                .flatMap(user -> workspaceService.createDefault(new Workspace(), user));
-                    })
-                    .map(Workspace::getId);
-        }
-
-        return Mono.just(defaultWorkspaceId);
-    }
-
-    protected Mono<Application> createFirstApplication(String workspaceId) {
-        // need to create default application
-        Application application = new Application();
-        application.setWorkspaceId(workspaceId);
-        application.setName("My first application");
-        return applicationPageService.createApplication(application);
+        return userSignupHelper
+                .createWorkspaceIfNotExistsAndGetId(defaultWorkspaceId, authentication)
+                .flatMap(userSignupHelper::createDefaultApplication);
     }
 
     /**

app/server/appsmith-server/src/main/java/com/appsmith/server/helpers/UserSignupHelper.java

@@ -0,0 +1,30 @@
+package com.appsmith.server.helpers;
+
+import com.appsmith.server.helpers.ce.UserSignupHelperCE;
+import com.appsmith.server.repositories.UserRepository;
+import com.appsmith.server.repositories.WorkspaceRepository;
+import com.appsmith.server.services.ApplicationPageService;
+import com.appsmith.server.services.OrganizationService;
+import com.appsmith.server.services.WorkspaceService;
+import com.appsmith.server.solutions.WorkspacePermission;
+import org.springframework.stereotype.Component;
+
+@Component
+public class UserSignupHelper extends UserSignupHelperCE {
+
+    public UserSignupHelper(
+            WorkspaceRepository workspaceRepository,
+            WorkspaceService workspaceService,
+            ApplicationPageService applicationPageService,
+            UserRepository userRepository,
+            OrganizationService organizationService,
+            WorkspacePermission workspacePermission) {
+        super(
+                workspaceRepository,
+                workspaceService,
+                applicationPageService,
+                userRepository,
+                organizationService,
+                workspacePermission);
+    }
+}

app/server/appsmith-server/src/main/java/com/appsmith/server/helpers/ce/UserSignupHelperCE.java

@@ -0,0 +1,122 @@
+package com.appsmith.server.helpers.ce;
+
+import com.appsmith.server.domains.Application;
+import com.appsmith.server.domains.User;
+import com.appsmith.server.domains.Workspace;
+import com.appsmith.server.repositories.UserRepository;
+import com.appsmith.server.repositories.WorkspaceRepository;
+import com.appsmith.server.services.ApplicationPageService;
+import com.appsmith.server.services.OrganizationService;
+import com.appsmith.server.services.WorkspaceService;
+import com.appsmith.server.solutions.WorkspacePermission;
+import lombok.extern.slf4j.Slf4j;
+import org.springframework.security.core.Authentication;
+import reactor.core.publisher.Mono;
+
+@Slf4j
+public class UserSignupHelperCE {
+    private final WorkspaceRepository workspaceRepository;
+    private final WorkspaceService workspaceService;
+    private final ApplicationPageService applicationPageService;
+    private final UserRepository userRepository;
+    private final OrganizationService organizationService;
+    private final WorkspacePermission workspacePermission;
+
+    public UserSignupHelperCE(
+            WorkspaceRepository workspaceRepository,
+            WorkspaceService workspaceService,
+            ApplicationPageService applicationPageService,
+            UserRepository userRepository,
+            OrganizationService organizationService,
+            WorkspacePermission workspacePermission) {
+        this.workspaceRepository = workspaceRepository;
+        this.workspaceService = workspaceService;
+        this.applicationPageService = applicationPageService;
+        this.userRepository = userRepository;
+        this.organizationService = organizationService;
+        this.workspacePermission = workspacePermission;
+    }
+
+    /**
+     * Creates a default workspace and application for a user.
+     * Basic CE implementation that creates a simple workspace and application.
+     *
+     * @param user The user for whom to create the workspace and application
+     * @return A Mono that completes when the workspace and application are created
+     */
+    public Mono<Void> createDefaultWorkspaceAndApplication(User user) {
+        log.debug("Creating default workspace and application for user: {}", user.getEmail());
+
+        // Create a default workspace
+        Workspace workspace = new Workspace();
+        workspace.setName(user.getName() != null ? user.getName() + "'s workspace" : "Default workspace");
+
+        return workspaceService
+                .create(workspace, user, Boolean.FALSE)
+                .flatMap(createdWorkspace ->
+                        createDefaultApplication(createdWorkspace.getId()).then())
+                .doOnError(error -> log.error("Error creating workspace or application: {}", error.getMessage()))
+                .then();
+    }
+
+    /**
+     * Creates a default application in the specified workspace.
+     *
+     * @param workspaceId ID of the workspace to create the application in
+     * @return A Mono containing the created Application
+     */
+    public Mono<Application> createDefaultApplication(String workspaceId) {
+        log.debug("Creating default application in workspace: {}", workspaceId);
+        Application application = new Application();
+        application.setWorkspaceId(workspaceId);
+        application.setName("My first application");
+        return applicationPageService.createApplication(application);
+    }
+
+    /**
+     * Gets an existing workspace ID or creates a new workspace if needed.
+     *
+     * @param defaultWorkspaceId The workspace ID to use if provided
+     * @param authentication The authentication object containing the user principal
+     * @return A Mono containing the workspace ID to use
+     */
+    public Mono<String> createWorkspaceIfNotExistsAndGetId(String defaultWorkspaceId, Authentication authentication) {
+        if (defaultWorkspaceId != null) {
+            return Mono.just(defaultWorkspaceId);
+        }
+
+        return workspaceRepository
+                .findAll(workspacePermission.getEditPermission())
+                .take(1, true)
+                .collectList()
+                .flatMap(workspaces -> {
+                    // Since this is the first application creation, the first workspace would be the only
+                    // workspace user has access to, and would be user's default workspace. Hence, we use this
+                    // workspace to create the application.
+                    if (workspaces.size() == 1) {
+                        return Mono.just(workspaces.get(0));
+                    }
+
+                    // In case no workspaces are found for the user, create a new default workspace
+                    User user = (User) authentication.getPrincipal();
+
+                    // Use the protected method that can be overridden in EE version
+                    return createDefaultWorkspaceForUser(user);
+                })
+                .map(Workspace::getId);
+    }
+
+    /**
+     * Creates a default workspace for a user. This method can be overridden in the EE version
+     * to add additional checks like multi-org feature flag.
+     *
+     * @param user User for whom to create the workspace
+     * @return Mono containing the created workspace
+     */
+    public Mono<Workspace> createDefaultWorkspaceForUser(User user) {
+        return organizationService
+                .getCurrentUserOrganizationId()
+                .flatMap(orgId -> userRepository.findByEmailAndOrganizationId(user.getEmail(), orgId))
+                .flatMap(user1 -> workspaceService.createDefault(new Workspace(), user1));
+    }
+}

app/server/appsmith-server/src/main/java/com/appsmith/server/services/ce/UserServiceCEImpl.java

@@ -445,6 +445,18 @@ protected Mono<Boolean> isSignupAllowed(User user) {
         return Mono.just(TRUE);
     }
 
+    /**
+     * Checks if a workspace should be created for a user during signup.
+     * This method can be overridden in EE to add checks for multi-org settings.
+     *
+     * @param user The user for whom to check workspace creation
+     * @return Mono<Boolean> true if workspace should be created, false otherwise
+     */
+    protected Mono<Boolean> shouldCreateWorkspaceForUser(User user) {
+        // In CE, always create workspace
+        return Mono.just(TRUE);
+    }
+
     @Override
     public Mono<UserSignupDTO> createUser(User user) {
         // Only encode the password if it's a form signup. For OAuth signups, we don't need password
@@ -491,29 +503,39 @@ public Mono<UserSignupDTO> createUser(User user) {
                                 final UserSignupDTO userSignupDTO = new UserSignupDTO();
                                 userSignupDTO.setUser(savedUser);
 
-                                return workspaceService
-                                        .createDefault(new Workspace(), savedUser)
-                                        .elapsed()
-                                        .map(pair -> {
-                                            log.debug(
-                                                    "UserServiceCEImpl::Time taken to create default workspace: {} ms",
-                                                    pair.getT1());
-                                            return pair.getT2();
-                                        })
-                                        .map(workspace -> {
-                                            log.debug(
-                                                    "Created blank default workspace for user '{}'.",
-                                                    savedUser.getEmail());
-                                            userSignupDTO.setDefaultWorkspaceId(workspace.getId());
-                                            return userSignupDTO;
-                                        })
-                                        .onErrorResume(e -> {
-                                            log.debug(
-                                                    "Error creating default workspace for user '{}'.",
-                                                    savedUser.getEmail(),
-                                                    e);
-                                            return Mono.just(userSignupDTO);
-                                        });
+                                // Check if we should create a workspace for this user
+                                return shouldCreateWorkspaceForUser(savedUser).flatMap(shouldCreateWorkspace -> {
+                                    if (Boolean.TRUE.equals(shouldCreateWorkspace)) {
+                                        // Create workspace as normal
+                                        return workspaceService
+                                                .createDefault(new Workspace(), savedUser)
+                                                .elapsed()
+                                                .map(pair -> {
+                                                    log.debug(
+                                                            "UserServiceCEImpl::Time taken to create default workspace: {} ms",
+                                                            pair.getT1());
+                                                    return pair.getT2();
+                                                })
+                                                .map(workspace -> {
+                                                    log.debug(
+                                                            "Created blank default workspace for user '{}'.",
+                                                            savedUser.getEmail());
+                                                    userSignupDTO.setDefaultWorkspaceId(workspace.getId());
+                                                    return userSignupDTO;
+                                                })
+                                                .onErrorResume(e -> {
+                                                    log.debug(
+                                                            "Error creating default workspace for user '{}'.",
+                                                            savedUser.getEmail(),
+                                                            e);
+                                                    return Mono.just(userSignupDTO);
+                                                });
+                                    } else {
+                                        // Skip workspace creation
+                                        log.debug("Skipping workspace creation for user: {}", savedUser.getEmail());
+                                        return Mono.just(userSignupDTO);
+                                    }
+                                });
                             })
                             .flatMap(userSignupDTO -> findByEmail(
                                             userSignupDTO.getUser().getEmail())