fix: fixes storage handling for non-auth recipes

ee/src/main/java/io/supertokens/ee/EEFeatureFlag.java

@@ -271,7 +271,7 @@ private JsonObject getMultiTenancyStats()
         return stats;
     }
 
-    private JsonObject getAccountLinkingStats() throws StorageQueryException {
+    private JsonObject getAccountLinkingStats() throws StorageQueryException, TenantOrAppNotFoundException {
         JsonObject result = new JsonObject();
         Storage[] storages = StorageLayer.getStoragesForApp(main, this.appIdentifier);
         boolean usesAccountLinking = false;

src/main/java/io/supertokens/authRecipe/AuthRecipe.java

@@ -37,6 +37,7 @@
 import io.supertokens.pluginInterface.exceptions.StorageQueryException;
 import io.supertokens.pluginInterface.exceptions.StorageTransactionLogicException;
 import io.supertokens.pluginInterface.multitenancy.AppIdentifierWithStorage;
+import io.supertokens.pluginInterface.multitenancy.AppIdentifierWithStorages;
 import io.supertokens.pluginInterface.multitenancy.TenantIdentifier;
 import io.supertokens.pluginInterface.multitenancy.TenantIdentifierWithStorage;
 import io.supertokens.pluginInterface.multitenancy.exceptions.TenantOrAppNotFoundException;
@@ -661,20 +662,20 @@ public static long getUsersCountForTenant(TenantIdentifierWithStorage tenantIden
                 tenantIdentifier, includeRecipeIds);
     }
 
-    public static long getUsersCountAcrossAllTenants(AppIdentifierWithStorage appIdentifierWithStorage,
+    public static long getUsersCountAcrossAllTenants(AppIdentifierWithStorages appIdentifierWithStorages,
                                                      RECIPE_ID[] includeRecipeIds)
             throws StorageQueryException,
             TenantOrAppNotFoundException, BadPermissionException {
         long count = 0;
 
-        for (Storage storage : appIdentifierWithStorage.getStorages()) {
+        for (Storage storage : appIdentifierWithStorages.getStorages()) {
             if (storage.getType() != STORAGE_TYPE.SQL) {
                 // we only support SQL for now
                 throw new UnsupportedOperationException("");
             }
 
             count += ((AuthRecipeStorage) storage).getUsersCount(
-                    appIdentifierWithStorage, includeRecipeIds);
+                    appIdentifierWithStorages, includeRecipeIds);
         }
 
         return count;

src/main/java/io/supertokens/cronjobs/telemetry/Telemetry.java

@@ -36,6 +36,7 @@
 import io.supertokens.pluginInterface.exceptions.StorageQueryException;
 import io.supertokens.pluginInterface.multitenancy.AppIdentifier;
 import io.supertokens.pluginInterface.multitenancy.AppIdentifierWithStorage;
+import io.supertokens.pluginInterface.multitenancy.AppIdentifierWithStorages;
 import io.supertokens.pluginInterface.multitenancy.TenantIdentifier;
 import io.supertokens.pluginInterface.multitenancy.exceptions.TenantOrAppNotFoundException;
 import io.supertokens.storageLayer.StorageLayer;
@@ -102,13 +103,12 @@ protected void doTaskPerApp(AppIdentifier app) throws Exception {
         if (StorageLayer.getBaseStorage(main).getType() == STORAGE_TYPE.SQL) {
             { // Users count across all tenants
                 Storage[] storages = StorageLayer.getStoragesForApp(main, app);
-                AppIdentifierWithStorage appIdentifierWithAllTenantStorages = new AppIdentifierWithStorage(
-                        app.getConnectionUriDomain(), app.getAppId(),
-                        StorageLayer.getStorage(app.getAsPublicTenantIdentifier(), main), storages
+                AppIdentifierWithStorages appIdentifierWithStorages = new AppIdentifierWithStorages(
+                        app.getConnectionUriDomain(), app.getAppId(), storages
                 );
 
                 json.addProperty("usersCount",
-                        AuthRecipe.getUsersCountAcrossAllTenants(appIdentifierWithAllTenantStorages, null));
+                        AuthRecipe.getUsersCountAcrossAllTenants(appIdentifierWithStorages, null));
             }
 
             { // Dashboard user emails

src/main/java/io/supertokens/storageLayer/StorageLayer.java

@@ -385,7 +385,8 @@ public static List<List<TenantIdentifier>> getTenantsWithUniqueUserPoolId(Main m
         return result;
     }
 
-    public static Storage[] getStoragesForApp(Main main, AppIdentifier appIdentifier) {
+    public static Storage[] getStoragesForApp(Main main, AppIdentifier appIdentifier)
+            throws TenantOrAppNotFoundException {
         Map<String, Storage> userPoolToStorage = new HashMap<>();
 
         Map<ResourceDistributor.KeyClass, ResourceDistributor.SingletonResource> resources =
@@ -397,7 +398,11 @@ public static Storage[] getStoragesForApp(Main main, AppIdentifier appIdentifier
                 userPoolToStorage.put(storage.getUserPoolId(), storage);
             }
         }
-        return userPoolToStorage.values().toArray(new Storage[0]);
+        Storage[] storages = userPoolToStorage.values().toArray(new Storage[0]);
+        if (storages.length == 0) {
+            throw new TenantOrAppNotFoundException(appIdentifier);
+        }
+        return storages;
     }
 
     public static TenantIdentifierWithStorageAndUserIdMapping getTenantIdentifierWithStorageAndUserIdMappingForUser(
@@ -432,64 +437,42 @@ public static TenantIdentifierWithStorageAndUserIdMapping getTenantIdentifierWit
         throw new UnknownUserIdException();
     }
 
-    public static AppIdentifierWithStorageAndUserIdMapping getAppIdentifierWithStorageAndUserIdMappingForUserWithPriorityForTenantStorage(
-            Main main, AppIdentifier appIdentifier, Storage priorityStorage, String userId,
+    public static AppIdentifierWithStorageAndUserIdMapping getAppIdentifierWithStorageAndUserIdMappingForUser(
+            AppIdentifierWithStorages appIdentifierWithStorages, String userId,
             UserIdType userIdType) throws StorageQueryException, TenantOrAppNotFoundException, UnknownUserIdException {
 
-        Storage[] storages = getStoragesForApp(main, appIdentifier);
+        Storage[] storages = appIdentifierWithStorages.getStorages();
 
         if (storages.length == 0) {
-            throw new TenantOrAppNotFoundException(appIdentifier);
+            throw new TenantOrAppNotFoundException(appIdentifierWithStorages);
         }
 
-        // We look for userId in the priorityStorage first just in case multiple storages have the mapping, we
-        // return the mapping from the storage of the tenant from which the request came from.
-        {
-            UserIdMapping mapping = io.supertokens.useridmapping.UserIdMapping.getUserIdMapping(
-                    appIdentifier.withStorage(priorityStorage),
-                    userId, userIdType);
-
-            if (mapping != null) {
-                AppIdentifierWithStorage appIdentifierWithStorage = appIdentifier.withStorage(priorityStorage);
-                return new AppIdentifierWithStorageAndUserIdMapping(appIdentifierWithStorage, mapping);
-            }
+        AppIdentifier appIdentifier = appIdentifierWithStorages;
 
+        // First look in auth recipes
+        for (Storage storage : storages) {
             if (userIdType != UserIdType.EXTERNAL
-                    && ((AuthRecipeStorage) priorityStorage).doesUserIdExist(appIdentifier, userId)) {
-                AppIdentifierWithStorage appIdentifierWithStorage = appIdentifier.withStorage(priorityStorage);
-                return new AppIdentifierWithStorageAndUserIdMapping(appIdentifierWithStorage, null);
-            }
-            if (userIdType != UserIdType.SUPERTOKENS) {
-                AppIdentifierWithStorage appIdentifierWithStorage = appIdentifier.withStorage(priorityStorage);
-                try {
-                    io.supertokens.useridmapping.UserIdMapping.findNonAuthStoragesWhereUserIdIsUsedOrAssertIfUsed(
-                            appIdentifierWithStorage, userId, true);
-                } catch (ServletException e) {
-                    // this means that the userId is being used for a non auth recipe.
-                    return new AppIdentifierWithStorageAndUserIdMapping(appIdentifierWithStorage, null);
-                }
+                    && ((AuthRecipeStorage) storage).doesUserIdExist(appIdentifier, userId)) {
+                AppIdentifierWithStorage appIdentifierWithStorage = appIdentifier.withStorage(storage);
+
+                UserIdMapping mapping = io.supertokens.useridmapping.UserIdMapping.getUserIdMapping(
+                        appIdentifierWithStorages.withStorage(storage),
+                        userId, userIdType);
+
+                return new AppIdentifierWithStorageAndUserIdMapping(appIdentifierWithStorage, mapping);
             }
         }
 
         for (Storage storage : storages) {
-            if (storage == priorityStorage) {
-                continue; // Already checked previously
-            }
-
             UserIdMapping mapping = io.supertokens.useridmapping.UserIdMapping.getUserIdMapping(
-                    appIdentifier.withStorage(storage),
+                    appIdentifierWithStorages.withStorage(storage),
                     userId, userIdType);
 
             if (mapping != null) {
                 AppIdentifierWithStorage appIdentifierWithStorage = appIdentifier.withStorage(storage);
                 return new AppIdentifierWithStorageAndUserIdMapping(appIdentifierWithStorage, mapping);
             }
 
-            if (userIdType != UserIdType.EXTERNAL
-                    && ((AuthRecipeStorage) storage).doesUserIdExist(appIdentifier, userId)) {
-                AppIdentifierWithStorage appIdentifierWithStorage = appIdentifier.withStorage(storage);
-                return new AppIdentifierWithStorageAndUserIdMapping(appIdentifierWithStorage, null);
-            }
             if (userIdType != UserIdType.SUPERTOKENS) {
                 AppIdentifierWithStorage appIdentifierWithStorage = appIdentifier.withStorage(storage);
                 try {

src/main/java/io/supertokens/useridmapping/UserIdMapping.java

@@ -28,6 +28,7 @@
 import io.supertokens.pluginInterface.exceptions.StorageTransactionLogicException;
 import io.supertokens.pluginInterface.jwt.JWTRecipeStorage;
 import io.supertokens.pluginInterface.multitenancy.AppIdentifierWithStorage;
+import io.supertokens.pluginInterface.multitenancy.AppIdentifierWithStorages;
 import io.supertokens.pluginInterface.multitenancy.TenantIdentifierWithStorage;
 import io.supertokens.pluginInterface.multitenancy.exceptions.TenantOrAppNotFoundException;
 import io.supertokens.pluginInterface.session.SessionStorage;
@@ -50,16 +51,26 @@
 public class UserIdMapping {
 
     @TestOnly
-    public static void createUserIdMapping(Main main, AppIdentifierWithStorage appIdentifierWithStorage,
+    public static void createUserIdMapping(AppIdentifierWithStorages appIdentifierWithStorages,
                                            String superTokensUserId, String externalUserId,
                                            String externalUserIdInfo, boolean force)
             throws ServletException, UnknownSuperTokensUserIdException, UserIdMappingAlreadyExistsException,
             StorageQueryException, TenantOrAppNotFoundException {
-        createUserIdMapping(main, appIdentifierWithStorage, superTokensUserId, externalUserId, externalUserIdInfo,
+        createUserIdMapping(appIdentifierWithStorages, superTokensUserId, externalUserId, externalUserIdInfo,
                 force, false);
     }
 
-    public static void createUserIdMapping(Main main, AppIdentifierWithStorage appIdentifierWithStorage,
+    @TestOnly
+    public static void createUserIdMapping(Main main, AppIdentifierWithStorage appIdentifierWithStorage, String supertokensUserId, String externalUserId, String externalUserIdInfo, boolean force)
+            throws ServletException, UnknownSuperTokensUserIdException, UserIdMappingAlreadyExistsException,
+            StorageQueryException, TenantOrAppNotFoundException {
+        createUserIdMapping(
+                new AppIdentifierWithStorages(appIdentifierWithStorage.getConnectionUriDomain(), appIdentifierWithStorage.getAppId(), new Storage[]{appIdentifierWithStorage.getStorage()}),
+                supertokensUserId, externalUserId, externalUserIdInfo, force
+        );
+    }
+
+    public static void createUserIdMapping(AppIdentifierWithStorages appIdentifierWithStorages,
                                            String superTokensUserId, String externalUserId,
                                            String externalUserIdInfo, boolean force, boolean makeExceptionForEmailVerification)
             throws UnknownSuperTokensUserIdException,
@@ -75,9 +86,8 @@ public static void createUserIdMapping(Main main, AppIdentifierWithStorage appId
         // race condition is fixed.
         try { // with external id
             AppIdentifierWithStorageAndUserIdMapping mappingAndStorage =
-                    StorageLayer.getAppIdentifierWithStorageAndUserIdMappingForUserWithPriorityForTenantStorage(
-                            main, appIdentifierWithStorage, appIdentifierWithStorage.getStorage(), externalUserId,
-                            UserIdType.EXTERNAL);
+                    StorageLayer.getAppIdentifierWithStorageAndUserIdMappingForUser(
+                            appIdentifierWithStorages, externalUserId, UserIdType.EXTERNAL);
 
             if (mappingAndStorage.userIdMapping != null) {
                 throw new UserIdMappingAlreadyExistsException(
@@ -89,6 +99,16 @@ public static void createUserIdMapping(Main main, AppIdentifierWithStorage appId
             // ignore this as we do not want external user id to exist
         }
 
+        AppIdentifierWithStorageAndUserIdMapping mappingAndStorage;
+        try {
+            mappingAndStorage = StorageLayer.getAppIdentifierWithStorageAndUserIdMappingForUser(
+                    appIdentifierWithStorages, superTokensUserId, UserIdType.SUPERTOKENS);
+        } catch (UnknownUserIdException e) {
+            throw new UnknownSuperTokensUserIdException();
+        }
+
+        AppIdentifierWithStorage appIdentifierWithStorage = mappingAndStorage.appIdentifierWithStorage;
+
         // if a userIdMapping is created with force, then we skip the following checks
         if (!force) {
             // We do not allow for a UserIdMapping to be created when the externalUserId is a SuperTokens userId.
@@ -99,7 +119,7 @@ public static void createUserIdMapping(Main main, AppIdentifierWithStorage appId
 
             {
                 if (((AuthRecipeStorage) appIdentifierWithStorage.getStorage()).doesUserIdExist(
-                        appIdentifierWithStorage, externalUserId)) {
+                        appIdentifierWithStorages, externalUserId)) {
                     throw new ServletException(new WebserverAPI.BadRequestException(
                             "Cannot create a userId mapping where the externalId is also a SuperTokens userID"));
                 }
@@ -115,7 +135,7 @@ public static void createUserIdMapping(Main main, AppIdentifierWithStorage appId
                     // an exception, then the creation of userIdMapping for the user will be blocked. And, to overcome that the
                     // email will have to be unverified first, then the userIdMapping should be created and then the email must be
                     // verified again on the externalUserId, which is not a good user experience.
-                    appIdentifierWithStorage.getEmailVerificationStorage().updateIsEmailVerifiedToExternalUserId(appIdentifierWithStorage, superTokensUserId, externalUserId);
+                    appIdentifierWithStorage.getEmailVerificationStorage().updateIsEmailVerifiedToExternalUserId(appIdentifierWithStorages, superTokensUserId, externalUserId);
                 } else  if (storageClasses.size() > 0) {
                     String recipeName = storageClasses.get(0);
                     String[] parts = recipeName.split("[.]");
@@ -127,12 +147,10 @@ public static void createUserIdMapping(Main main, AppIdentifierWithStorage appId
             } else {
                 findNonAuthStoragesWhereUserIdIsUsedOrAssertIfUsed(appIdentifierWithStorage, superTokensUserId, true);
             }
-
-
         }
 
         appIdentifierWithStorage.getUserIdMappingStorage()
-                .createUserIdMapping(appIdentifierWithStorage, superTokensUserId,
+                .createUserIdMapping(appIdentifierWithStorages, superTokensUserId,
                         externalUserId, externalUserIdInfo);
     }
     @TestOnly
@@ -152,9 +170,8 @@ public static void createUserIdMapping(Main main,
             UserIdMappingAlreadyExistsException, StorageQueryException, ServletException, UnknownUserIdException {
         try {
             Storage storage = StorageLayer.getStorage(main);
-            createUserIdMapping(main, new AppIdentifierWithStorage(null, null, storage), superTokensUserId,
-                    externalUserId,
-                    externalUserIdInfo, force, makeExceptionForEmailVerification);
+            createUserIdMapping(new AppIdentifierWithStorages(null, null, new Storage[]{storage}), superTokensUserId,
+                    externalUserId, externalUserIdInfo, force, makeExceptionForEmailVerification);
         } catch (TenantOrAppNotFoundException e) {
             throw new IllegalStateException(e);
         }