chore: Remove fetching users without organizationId

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

@@ -1,6 +1,7 @@
 package com.appsmith.server.authentication.handlers;
 
 import com.appsmith.server.authentication.handlers.ce.CustomFormLoginServiceCEImpl;
+import com.appsmith.server.helpers.UserOrganizationHelper;
 import com.appsmith.server.repositories.UserRepository;
 import lombok.extern.slf4j.Slf4j;
 import org.springframework.beans.factory.annotation.Autowired;
@@ -11,7 +12,7 @@
 public class CustomFormLoginServiceImpl extends CustomFormLoginServiceCEImpl {
 
     @Autowired
-    public CustomFormLoginServiceImpl(UserRepository repository) {
-        super(repository);
+    public CustomFormLoginServiceImpl(UserRepository repository, UserOrganizationHelper userOrganizationHelper) {
+        super(repository, userOrganizationHelper);
     }
 }

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

@@ -1,6 +1,7 @@
 package com.appsmith.server.authentication.handlers;
 
 import com.appsmith.server.authentication.handlers.ce.CustomOAuth2UserServiceCEImpl;
+import com.appsmith.server.helpers.UserOrganizationHelper;
 import com.appsmith.server.repositories.UserRepository;
 import com.appsmith.server.services.UserService;
 import lombok.extern.slf4j.Slf4j;
@@ -19,8 +20,9 @@ public class CustomOAuth2UserServiceImpl extends CustomOAuth2UserServiceCEImpl
     private UserService userService;
 
     @Autowired
-    public CustomOAuth2UserServiceImpl(UserRepository repository, UserService userService) {
-        super(repository, userService);
+    public CustomOAuth2UserServiceImpl(
+            UserRepository repository, UserService userService, UserOrganizationHelper userOrganizationHelper) {
+        super(repository, userService, userOrganizationHelper);
         this.repository = repository;
         this.userService = userService;
     }

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

@@ -1,6 +1,7 @@
 package com.appsmith.server.authentication.handlers;
 
 import com.appsmith.server.authentication.handlers.ce.CustomOidcUserServiceCEImpl;
+import com.appsmith.server.helpers.UserOrganizationHelper;
 import com.appsmith.server.repositories.UserRepository;
 import com.appsmith.server.services.UserService;
 import lombok.extern.slf4j.Slf4j;
@@ -17,10 +18,12 @@ public class CustomOidcUserServiceImpl extends CustomOidcUserServiceCEImpl
 
     private UserRepository repository;
     private UserService userService;
+    private UserOrganizationHelper userOrganizationHelper;
 
     @Autowired
-    public CustomOidcUserServiceImpl(UserRepository repository, UserService userService) {
-        super(repository, userService);
+    public CustomOidcUserServiceImpl(
+            UserRepository repository, UserService userService, UserOrganizationHelper userOrganizationHelper) {
+        super(repository, userService, userOrganizationHelper);
         this.repository = repository;
         this.userService = userService;
     }

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

@@ -75,7 +75,10 @@ private Mono<Boolean> isVerificationRequired(String userEmail, String method) {
                 .map(organization -> organization.getOrganizationConfiguration().isEmailVerificationEnabled())
                 .cache();
 
-        Mono<User> userMono = userRepository.findByEmail(userEmail).cache();
+        Mono<User> userMono = organizationService
+                .getCurrentUserOrganizationId()
+                .flatMap(orgId -> userRepository.findByEmailAndOrganizationId(userEmail, orgId))
+                .cache();
         Mono<Boolean> verificationRequiredMono = null;
 
         if ("signup".equals(method)) {
@@ -367,8 +370,9 @@ protected Mono<Application> createDefaultApplication(String defaultWorkspaceId,
                         // In case no workspaces are found for the user, create a new default workspace
                         String email = ((User) authentication.getPrincipal()).getEmail();
 
-                        return userRepository
-                                .findByEmail(email)
+                        return organizationService
+                                .getCurrentUserOrganizationId()
+                                .flatMap(orgId -> userRepository.findByEmailAndOrganizationId(email, orgId))
                                 .flatMap(user -> workspaceService.createDefault(new Workspace(), user))
                                 .map(workspace -> {
                                     application.setWorkspaceId(workspace.getId());

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

@@ -2,6 +2,7 @@
 
 import com.appsmith.server.domains.LoginSource;
 import com.appsmith.server.exceptions.AppsmithError;
+import com.appsmith.server.helpers.UserOrganizationHelper;
 import com.appsmith.server.repositories.UserRepository;
 import lombok.extern.slf4j.Slf4j;
 import org.apache.commons.lang.WordUtils;
@@ -16,10 +17,12 @@
 public class CustomFormLoginServiceCEImpl implements ReactiveUserDetailsService {
 
     private UserRepository repository;
+    private UserOrganizationHelper userOrganizationHelper;
 
     @Autowired
-    public CustomFormLoginServiceCEImpl(UserRepository repository) {
+    public CustomFormLoginServiceCEImpl(UserRepository repository, UserOrganizationHelper userOrganizationHelper) {
         this.repository = repository;
+        this.userOrganizationHelper = userOrganizationHelper;
     }
 
     /**
@@ -31,9 +34,13 @@ public CustomFormLoginServiceCEImpl(UserRepository repository) {
      */
     @Override
     public Mono<UserDetails> findByUsername(String username) {
-        return repository
-                .findByEmail(username)
-                .switchIfEmpty(repository.findFirstByEmailIgnoreCaseOrderByCreatedAtDesc(username))
+
+        return userOrganizationHelper
+                .getCurrentUserOrganizationId()
+                .flatMap(orgId -> repository
+                        .findByEmailAndOrganizationId(username, orgId)
+                        .switchIfEmpty(repository.findFirstByEmailIgnoreCaseAndOrganizationIdOrderByCreatedAtDesc(
+                                username, orgId)))
                 .switchIfEmpty(Mono.error(new UsernameNotFoundException("Unable to find username: " + username)))
                 .onErrorMap(error -> {
                     log.error("Can't find user {}", username);

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

@@ -5,6 +5,7 @@
 import com.appsmith.server.domains.UserState;
 import com.appsmith.server.exceptions.AppsmithException;
 import com.appsmith.server.exceptions.AppsmithOAuth2AuthenticationException;
+import com.appsmith.server.helpers.UserOrganizationHelper;
 import com.appsmith.server.repositories.UserRepository;
 import com.appsmith.server.services.UserService;
 import lombok.extern.slf4j.Slf4j;
@@ -26,11 +27,14 @@ public class CustomOAuth2UserServiceCEImpl extends DefaultReactiveOAuth2UserServ
 
     private UserRepository repository;
     private UserService userService;
+    private UserOrganizationHelper userOrganizationHelper;
 
     @Autowired
-    public CustomOAuth2UserServiceCEImpl(UserRepository repository, UserService userService) {
+    public CustomOAuth2UserServiceCEImpl(
+            UserRepository repository, UserService userService, UserOrganizationHelper userOrganizationHelper) {
         this.repository = repository;
         this.userService = userService;
+        this.userOrganizationHelper = userOrganizationHelper;
     }
 
     @Override
@@ -76,8 +80,12 @@ private Mono<User> checkAndCreateUser(OAuth2User oAuth2User, OAuth2UserRequest u
     }
 
     protected Mono<User> findByUsername(String email) {
-        return repository
-                .findByEmail(email)
-                .switchIfEmpty(repository.findFirstByEmailIgnoreCaseOrderByCreatedAtDesc(email));
+
+        return userOrganizationHelper.getCurrentUserOrganizationId().flatMap(organizationId -> {
+            return repository
+                    .findByEmailAndOrganizationId(email, organizationId)
+                    .switchIfEmpty(repository.findFirstByEmailIgnoreCaseAndOrganizationIdOrderByCreatedAtDesc(
+                            email, organizationId));
+        });
     }
 }

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

@@ -5,6 +5,7 @@
 import com.appsmith.server.domains.UserState;
 import com.appsmith.server.exceptions.AppsmithException;
 import com.appsmith.server.exceptions.AppsmithOAuth2AuthenticationException;
+import com.appsmith.server.helpers.UserOrganizationHelper;
 import com.appsmith.server.repositories.UserRepository;
 import com.appsmith.server.services.UserService;
 import lombok.extern.slf4j.Slf4j;
@@ -28,11 +29,14 @@ public class CustomOidcUserServiceCEImpl extends OidcReactiveOAuth2UserService {
 
     private UserRepository repository;
     private UserService userService;
+    private UserOrganizationHelper userOrganizationHelper;
 
     @Autowired
-    public CustomOidcUserServiceCEImpl(UserRepository repository, UserService userService) {
+    public CustomOidcUserServiceCEImpl(
+            UserRepository repository, UserService userService, UserOrganizationHelper userOrganizationHelper) {
         this.repository = repository;
         this.userService = userService;
+        this.userOrganizationHelper = userOrganizationHelper;
     }
 
     @Override
@@ -82,8 +86,11 @@ public Mono<User> checkAndCreateUser(OidcUser oidcUser, OidcUserRequest userRequ
     }
 
     protected Mono<User> findByUsername(String email) {
-        return repository
-                .findByEmail(email)
-                .switchIfEmpty(repository.findFirstByEmailIgnoreCaseOrderByCreatedAtDesc(email));
+        return userOrganizationHelper.getCurrentUserOrganizationId().flatMap(organizationId -> {
+            return repository
+                    .findByEmailAndOrganizationId(email, organizationId)
+                    .switchIfEmpty(repository.findFirstByEmailIgnoreCaseAndOrganizationIdOrderByCreatedAtDesc(
+                            email, organizationId));
+        });
     }
 }

app/server/appsmith-server/src/main/java/com/appsmith/server/domains/PasswordResetToken.java

@@ -17,6 +17,7 @@
 public class PasswordResetToken extends BaseDomain {
     String tokenHash;
     String email;
+    String organizationId;
     int requestCount; // number of requests in last 24 hours
     Instant firstRequestTime; // when a request was first generated in last 24 hours
 }

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

@@ -0,0 +1,5 @@
+package com.appsmith.server.helpers;
+
+import com.appsmith.server.helpers.ce.UserOrganizationHelperCE;
+
+public interface UserOrganizationHelper extends UserOrganizationHelperCE {}

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

@@ -0,0 +1,14 @@
+package com.appsmith.server.helpers;
+
+import com.appsmith.server.helpers.ce.UserOrganizationHelperCEImpl;
+import com.appsmith.server.repositories.OrganizationRepository;
+import org.springframework.stereotype.Component;
+
+@Component
+public class UserOrganizationHelperImpl extends UserOrganizationHelperCEImpl implements UserOrganizationHelper {
+    public UserOrganizationHelperImpl(
+            OrganizationRepository organizationRepository,
+            InMemoryCacheableRepositoryHelper inMemoryCacheableRepositoryHelper) {
+        super(organizationRepository, inMemoryCacheableRepositoryHelper);
+    }
+}

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

@@ -0,0 +1,11 @@
+package com.appsmith.server.helpers.ce;
+
+import reactor.core.publisher.Mono;
+
+public interface UserOrganizationHelperCE {
+    /**
+     * Gets the organization ID for the current user. If not found in user context, falls back to default organization.
+     * @return Mono containing the organization ID
+     */
+    Mono<String> getCurrentUserOrganizationId();
+}

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

@@ -0,0 +1,44 @@
+package com.appsmith.server.helpers.ce;
+
+import com.appsmith.server.domains.Organization;
+import com.appsmith.server.helpers.InMemoryCacheableRepositoryHelper;
+import com.appsmith.server.repositories.OrganizationRepository;
+import lombok.extern.slf4j.Slf4j;
+import org.springframework.util.StringUtils;
+import reactor.core.publisher.Mono;
+
+import static com.appsmith.server.constants.FieldName.DEFAULT;
+
+@Slf4j
+public class UserOrganizationHelperCEImpl implements UserOrganizationHelperCE {
+
+    private final OrganizationRepository organizationRepository;
+    private final InMemoryCacheableRepositoryHelper inMemoryCacheableRepositoryHelper;
+
+    public UserOrganizationHelperCEImpl(
+            OrganizationRepository organizationRepository,
+            InMemoryCacheableRepositoryHelper inMemoryCacheableRepositoryHelper) {
+        this.organizationRepository = organizationRepository;
+        this.inMemoryCacheableRepositoryHelper = inMemoryCacheableRepositoryHelper;
+    }
+
+    @Override
+    public Mono<String> getCurrentUserOrganizationId() {
+        return getDefaultOrganizationId();
+    }
+
+    protected Mono<String> getDefaultOrganizationId() {
+        // If the value exists in cache, return it as is
+        if (StringUtils.hasLength(inMemoryCacheableRepositoryHelper.getDefaultOrganizationId())) {
+            return Mono.just(inMemoryCacheableRepositoryHelper.getDefaultOrganizationId());
+        }
+        return organizationRepository
+                .findBySlug(DEFAULT)
+                .map(Organization::getId)
+                .map(organizationId -> {
+                    // Set the cache value before returning.
+                    inMemoryCacheableRepositoryHelper.setDefaultOrganizationId(organizationId);
+                    return organizationId;
+                });
+    }
+}

app/server/appsmith-server/src/main/java/com/appsmith/server/migrations/db/ce/Migration068_AddOrganizationIdToPasswordResetToken.java

@@ -0,0 +1,56 @@
+package com.appsmith.server.migrations.db.ce;
+
+import io.mongock.api.annotations.ChangeUnit;
+import io.mongock.api.annotations.Execution;
+import io.mongock.api.annotations.RollbackExecution;
+import lombok.extern.slf4j.Slf4j;
+import org.bson.Document;
+import org.springframework.data.mongodb.core.MongoTemplate;
+import org.springframework.data.mongodb.core.query.Query;
+
+import java.util.Arrays;
+import java.util.List;
+
+import static org.springframework.data.mongodb.core.query.Criteria.where;
+
+@Slf4j
+@ChangeUnit(id = "add-organization-id-to-password-reset-token", order = "068")
+public class Migration068_AddOrganizationIdToPasswordResetToken {
+
+    private final MongoTemplate mongoTemplate;
+
+    public Migration068_AddOrganizationIdToPasswordResetToken(MongoTemplate mongoTemplate) {
+        this.mongoTemplate = mongoTemplate;
+    }
+
+    @RollbackExecution
+    public void rollbackExecution() {}
+
+    @Execution
+    public void execute() {
+        try {
+            // Get the organization ID from the organization collection
+            Document organization = mongoTemplate.findOne(new Query(), Document.class, "organization");
+            if (organization == null) {
+                log.info("No organization found to migrate password reset tokens");
+                return;
+            }
+            String organizationId = organization.getObjectId("_id").toString();
+
+            // Update all password reset tokens that don't have an organizationId
+            Query query = new Query(where("organizationId").exists(false));
+            List<Document> pipeline =
+                    Arrays.asList(new Document("$set", new Document("organizationId", organizationId)));
+
+            long updatedCount = mongoTemplate
+                    .getCollection("passwordResetToken")
+                    .updateMany(query.getQueryObject(), pipeline)
+                    .getModifiedCount();
+
+            log.info("Successfully set organizationId for {} password reset token documents", updatedCount);
+        } catch (Exception e) {
+            log.error("Error while setting organizationId for password reset tokens: ", e);
+            throw e;
+        }
+    }
+}

app/server/appsmith-server/src/main/java/com/appsmith/server/repositories/ce/CustomUserRepositoryCE.java

@@ -1,6 +1,5 @@
 package com.appsmith.server.repositories.ce;
 
-import com.appsmith.server.acl.AclPermission;
 import com.appsmith.server.domains.User;
 import com.appsmith.server.repositories.AppsmithRepository;
 import org.springframework.data.mongodb.core.query.UpdateDefinition;
@@ -10,8 +9,6 @@
 
 public interface CustomUserRepositoryCE extends AppsmithRepository<User> {
 
-    Mono<User> findByEmail(String email, AclPermission aclPermission);
-
     Mono<User> findByEmailAndOrganizationId(String email, String organizationId);
 
     Mono<Boolean> isUsersEmpty();

app/server/appsmith-server/src/main/java/com/appsmith/server/repositories/ce/CustomUserRepositoryCEImpl.java

@@ -1,12 +1,10 @@
 package com.appsmith.server.repositories.ce;
 
-import com.appsmith.server.acl.AclPermission;
 import com.appsmith.server.constants.FieldName;
 import com.appsmith.server.domains.User;
 import com.appsmith.server.exceptions.AppsmithError;
 import com.appsmith.server.exceptions.AppsmithException;
 import com.appsmith.server.helpers.ce.bridge.Bridge;
-import com.appsmith.server.helpers.ce.bridge.BridgeQuery;
 import com.appsmith.server.projections.IdOnly;
 import com.appsmith.server.repositories.BaseAppsmithRepositoryImpl;
 import lombok.extern.slf4j.Slf4j;
@@ -19,12 +17,6 @@
 @Slf4j
 public class CustomUserRepositoryCEImpl extends BaseAppsmithRepositoryImpl<User> implements CustomUserRepositoryCE {
 
-    @Override
-    public Mono<User> findByEmail(String email, AclPermission aclPermission) {
-        BridgeQuery<User> emailCriteria = Bridge.equal(User.Fields.email, email);
-        return queryBuilder().criteria(emailCriteria).permission(aclPermission).one();
-    }
-
     @Override
     public Mono<User> findByEmailAndOrganizationId(String email, String organizationId) {
         return queryBuilder()

app/server/appsmith-server/src/main/java/com/appsmith/server/repositories/ce/PasswordResetTokenRepositoryCE.java

@@ -6,5 +6,5 @@
 
 public interface PasswordResetTokenRepositoryCE extends BaseRepository<PasswordResetToken, String> {
 
-    Mono<PasswordResetToken> findByEmail(String email);
+    Mono<PasswordResetToken> findByEmailAndOrganizationId(String email, String organizationId);
 }