fix: bulk migration logging, tests, related query optimizations (#1131)

tamassoltesz · web-flow · commit 06c176d24370 · 2025-03-26T15:55:55.000+05:30
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -7,6 +7,27 @@ to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
 ## [Unreleased]
 
+## [10.1.2]
+
+
+- Adds user_id index to the user roles table
+- Adds more debug logging to bulk migration
+- Adds more tests to bulk migration
+
+### Migration
+
+If using PostgreSQL, run the following SQL script:
+
+```sql
+CREATE INDEX IF NOT EXISTS user_roles_app_id_user_id_index ON user_roles (app_id, user_id);
+```
+
+If using MySQL, run the following SQL script:
+```sql
+CREATE INDEX user_roles_app_id_user_id_index ON user_roles (app_id, user_id);
+```
+
+
 ## [10.1.1]
 
 - Adds debug logging for the bulk migration process
diff --git a/build.gradle b/build.gradle
@@ -19,7 +19,7 @@ compileTestJava { options.encoding = "UTF-8" }
 //    }
 //}
 
-version = "10.1.1"
+version = "10.1.2"
 
 repositories {
     mavenCentral()
diff --git a/src/main/java/io/supertokens/bulkimport/BulkImport.java b/src/main/java/io/supertokens/bulkimport/BulkImport.java
@@ -411,7 +411,7 @@ private static void associateUserToTenants(Main main, AppIdentifier appIdentifie
 
                 TenantIdentifier tenantIdentifier = new TenantIdentifier(appIdentifier.getConnectionUriDomain(),
                         appIdentifier.getAppId(), tenantId);
-                Multitenancy.addUserIdToTenant(main, tenantIdentifier, storage, lm.getSuperTokenOrExternalUserId());
+                Multitenancy.addUserIdToTenant(main, tenantIdentifier, storage, lm.superTokensUserId);
             } catch (TenantOrAppNotFoundException e) {
                 throw new StorageTransactionLogicException(new Exception("E009: " + e.getMessage()));
             } catch (StorageQueryException e) {
diff --git a/src/main/java/io/supertokens/cronjobs/bulkimport/ProcessBulkImportUsers.java b/src/main/java/io/supertokens/cronjobs/bulkimport/ProcessBulkImportUsers.java
@@ -78,6 +78,7 @@ protected void doTaskPerApp(AppIdentifier app)
         int bulkMigrationBatchSize = Config.getConfig(app.getAsPublicTenantIdentifier(), main)
                 .getBulkMigrationBatchSize();
 
+        Logging.debug(main, app.getAsPublicTenantIdentifier(), "CronTask starts. Instance: " + this);
         Logging.debug(main, app.getAsPublicTenantIdentifier(), "CronTask starts. Processing bulk import users with " + bulkMigrationBatchSize
                 + " batch size, one batch split into " + numberOfBatchChunks + " chunks");
 
@@ -119,8 +120,10 @@ protected void doTaskPerApp(AppIdentifier app)
 
                 for (Future<?> task : tasks) {
                     while (!task.isDone()) {
+                        Logging.debug(main, app.getAsPublicTenantIdentifier(), "Waiting for task " + task + " to finish");
                         Thread.sleep(1000);
                     }
+                    Logging.debug(main, app.getAsPublicTenantIdentifier(), "Task " + task + " finished");
                     Void result = (Void) task.get(); //to know if there were any errors while executing and for waiting in this thread for all the other threads to finish up
                     usersProcessed += loadedUsersChunks.get(tasks.indexOf(task)).size();
                     failedUsers = bulkImportSQLStorage.getBulkImportUsersCount(app, BulkImportStorage.BULK_IMPORT_USER_STATUS.FAILED);
diff --git a/src/main/java/io/supertokens/cronjobs/bulkimport/ProcessBulkUsersImportWorker.java b/src/main/java/io/supertokens/cronjobs/bulkimport/ProcessBulkUsersImportWorker.java
@@ -113,7 +113,8 @@ private void processMultipleUsers(AppIdentifier appIdentifier, List<BulkImportUs
                 while (shouldRetryImmediatley) {
                     shouldRetryImmediatley = bulkImportProxyStorage.startTransaction(con -> {
                         try {
-                            BulkImport.processUsersImportSteps(main, appIdentifier, bulkImportProxyStorage, partitionedUsers.get(bulkImportProxyStorage),
+                            BulkImport.processUsersImportSteps(main, appIdentifier, bulkImportProxyStorage,
+                                    partitionedUsers.get(bulkImportProxyStorage),
                                     allStoragesForApp);
 
                             bulkImportProxyStorage.commitTransactionForBulkImportProxyStorage();
@@ -123,8 +124,16 @@ private void processMultipleUsers(AppIdentifier appIdentifier, List<BulkImportUs
                                 toDelete[i] = validUsers.get(i).id;
                             }
 
-                            baseTenantStorage.deleteBulkImportUsers(appIdentifier, toDelete);
-                        } catch (StorageTransactionLogicException e) {
+                            while (true){
+                                try {
+                                    baseTenantStorage.deleteBulkImportUsers(appIdentifier, toDelete);
+                                    break;
+                                } catch (Exception e) {
+                                    // ignore and retry delete. The import transaction is already committed, the delete should happen no matter what
+                                    Logging.debug(main, app.getAsPublicTenantIdentifier(), "Exception while deleting bulk import users: " + e.getMessage());
+                                }
+                            }
+                        } catch (StorageTransactionLogicException | StorageQueryException e) {
                             // We need to rollback the transaction manually because we have overridden that in the proxy
                             // storage
                             bulkImportProxyStorage.rollbackTransactionForBulkImportProxyStorage();
diff --git a/src/main/java/io/supertokens/inmemorydb/queries/GeneralQueries.java b/src/main/java/io/supertokens/inmemorydb/queries/GeneralQueries.java
@@ -402,6 +402,7 @@ public static void createTablesIfNotExists(Start start, Main main) throws SQLExc
             update(start, UserRolesQueries.getQueryToCreateRolePermissionsTable(start), NO_OP_SETTER);
             // index
             update(start, UserRolesQueries.getQueryToCreateRolePermissionsPermissionIndex(start), NO_OP_SETTER);
+            update(start, UserRolesQueries.getQueryToCreateUserRolesUserIdAppIdIndex(start), NO_OP_SETTER);
         }
 
         if (!doesTableExists(start, Config.getConfig(start).getUserRolesTable())) {
diff --git a/src/main/java/io/supertokens/inmemorydb/queries/UserRolesQueries.java b/src/main/java/io/supertokens/inmemorydb/queries/UserRolesQueries.java
@@ -66,6 +66,11 @@ static String getQueryToCreateRolePermissionsPermissionIndex(Start start) {
                 + Config.getConfig(start).getUserRolesPermissionsTable() + "(app_id, permission);";
     }
 
+    public static String getQueryToCreateUserRolesUserIdAppIdIndex(Start start) {
+        return "CREATE INDEX user_roles_app_id_user_id_index ON " + Config.getConfig(start).getUserRolesTable() +
+                "(app_id, user_id)";
+    }
+
     public static String getQueryToCreateUserRolesTable(Start start) {
         String tableName = Config.getConfig(start).getUserRolesTable();
         // @formatter:off
diff --git a/src/test/java/io/supertokens/test/bulkimport/BulkImportFlowTest.java b/src/test/java/io/supertokens/test/bulkimport/BulkImportFlowTest.java
@@ -127,7 +127,7 @@ public void testWithALotOfUsers() throws Exception {
                         throw e;
                     }
                 }
-                Thread.sleep(5000);
+                Thread.sleep(1000);
             }
         }
 
@@ -145,6 +145,120 @@ public void testWithALotOfUsers() throws Exception {
 
     }
 
+
+    @Test
+    public void testCoreRestartMidImportShouldResultInSuccessfulImport() throws Exception {
+        String[] args = { "../" };
+
+        // set processing thread number
+        Utils.setValueInConfig("bulk_migration_parallelism", "8");
+        Utils.setValueInConfig("bulk_migration_batch_size", "1000");
+        Utils.setValueInConfig("log_level", "DEBUG");
+
+        TestingProcessManager.TestingProcess process = TestingProcessManager.start(args, true);
+        Main main = process.getProcess();
+        setFeatureFlags(main, new EE_FEATURES[] {
+                EE_FEATURES.ACCOUNT_LINKING, EE_FEATURES.MULTI_TENANCY, EE_FEATURES.MFA});
+        // We are setting a non-zero initial wait for tests to avoid race condition with the beforeTest process that deletes data in the storage layer
+        CronTaskTest.getInstance(main).setInitialWaitTimeInSeconds(ProcessBulkImportUsers.RESOURCE_KEY, 5);
+        CronTaskTest.getInstance(main).setIntervalInSeconds(ProcessBulkImportUsers.RESOURCE_KEY, 60);
+
+        assertNotNull(process.checkOrWaitForEvent(ProcessState.PROCESS_STATE.STARTED));
+
+        Cronjobs.addCronjob(main, (ProcessBulkImportUsers) main.getResourceDistributor().getResource(new TenantIdentifier(null, null, null), ProcessBulkImportUsers.RESOURCE_KEY));
+
+        int NUMBER_OF_USERS_TO_UPLOAD = 10000;
+
+        if (StorageLayer.getBaseStorage(main).getType() != STORAGE_TYPE.SQL || StorageLayer.isInMemDb(main)) {
+            return;
+        }
+
+        // Create user roles before inserting bulk users
+        {
+            UserRoles.createNewRoleOrModifyItsPermissions(main, "role1", null);
+            UserRoles.createNewRoleOrModifyItsPermissions(main, "role2", null);
+        }
+
+        // upload a bunch of users through the API
+        {
+            for (int i = 0; i < (NUMBER_OF_USERS_TO_UPLOAD / 1000); i++) {
+                JsonObject request = generateUsersJson(1000, i * 1000); // API allows 10k users upload at once
+                JsonObject response = uploadBulkImportUsersJson(main, request);
+                assertEquals("OK", response.get("status").getAsString());
+            }
+
+        }
+
+        long processingStarted = System.currentTimeMillis();
+        boolean restartHappened = false;
+        // wait for the cron job to process them
+        // periodically check the remaining unprocessed users
+        // Note1: the cronjob starts the processing automatically
+        // Note2: the successfully processed users get deleted from the bulk_import_users table
+        {
+            long count = NUMBER_OF_USERS_TO_UPLOAD;
+            while(true) {
+                try {
+                    JsonObject response = loadBulkImportUsersCountWithStatus(main, null);
+                    assertEquals("OK", response.get("status").getAsString());
+                    count = response.get("count").getAsLong();
+                    int newUsersNumber = loadBulkImportUsersCountWithStatus(main,
+                            BulkImportStorage.BULK_IMPORT_USER_STATUS.NEW).get("count").getAsInt();
+                    int processingUsersNumber = loadBulkImportUsersCountWithStatus(main,
+                            BulkImportStorage.BULK_IMPORT_USER_STATUS.PROCESSING).get("count").getAsInt();
+                    int failedUsersNumber = loadBulkImportUsersCountWithStatus(main,
+                            BulkImportStorage.BULK_IMPORT_USER_STATUS.FAILED).get("count").getAsInt();
+                    count = newUsersNumber + processingUsersNumber;
+                    System.out.println("Remaining users: " + count);
+
+                    if (count == 0) {
+                        break;
+                    }
+                    if((System.currentTimeMillis() - processingStarted > 10000) && !restartHappened) {
+                        System.out.println("Killing core");
+                        process.kill(false);
+                        Utils.setValueInConfig("bulk_migration_parallelism", "14");
+                        Utils.setValueInConfig("bulk_migration_batch_size", "4000");
+                        System.out.println("Started new core");
+                        process = TestingProcessManager.start(args, true);
+                        main = process.getProcess();
+                        setFeatureFlags(main, new EE_FEATURES[] {
+                                EE_FEATURES.ACCOUNT_LINKING, EE_FEATURES.MULTI_TENANCY, EE_FEATURES.MFA});
+                        // We are setting a non-zero initial wait for tests to avoid race condition with the beforeTest process that deletes data in the storage layer
+                        CronTaskTest.getInstance(main).setInitialWaitTimeInSeconds(ProcessBulkImportUsers.RESOURCE_KEY, 5);
+                        CronTaskTest.getInstance(main).setIntervalInSeconds(ProcessBulkImportUsers.RESOURCE_KEY, 60);
+
+                        assertNotNull(process.checkOrWaitForEvent(ProcessState.PROCESS_STATE.STARTED));
+
+                        Cronjobs.addCronjob(main, (ProcessBulkImportUsers) main.getResourceDistributor().getResource(new TenantIdentifier(null, null, null), ProcessBulkImportUsers.RESOURCE_KEY));
+                        restartHappened = true;
+                    }
+                } catch (Exception e) {
+                    if(e instanceof SocketTimeoutException)  {
+                        //ignore
+                    } else {
+                        throw e;
+                    }
+                }
+                Thread.sleep(1000);
+            }
+        }
+
+        long processingFinished = System.currentTimeMillis();
+        System.out.println("Processed " + NUMBER_OF_USERS_TO_UPLOAD + " users in " + (processingFinished - processingStarted) / 1000
+                + " seconds ( or " + (processingFinished - processingStarted) / 60000 + " minutes)");
+
+        // after processing finished, make sure every user got processed correctly
+        {
+            int failedImportedUsersNumber = loadBulkImportUsersCountWithStatus(main, BulkImportStorage.BULK_IMPORT_USER_STATUS.FAILED).get("count").getAsInt();
+            int usersInCore = loadUsersCount(main).get("count").getAsInt();
+            assertEquals(NUMBER_OF_USERS_TO_UPLOAD, usersInCore + failedImportedUsersNumber);
+            assertEquals(NUMBER_OF_USERS_TO_UPLOAD, usersInCore);
+        }
+
+    }
+
+
     @Test
     public void testBatchWithOneUser() throws Exception {
         Main main = startCronProcess("14");
@@ -801,7 +915,7 @@ private static JsonObject uploadBulkImportUsersJson(Main main, JsonObject reques
 
     @NotNull
     private Main startCronProcess(String parallelism) throws IOException, InterruptedException, TenantOrAppNotFoundException {
-        return startCronProcess(parallelism, 5*60);
+        return startCronProcess(parallelism, 5 * 60);
     }
 
 
@@ -814,15 +928,14 @@ private Main startCronProcess(String parallelism, int intervalInSeconds) throws
         //Utils.setValueInConfig("bulk_migration_batch_size", "1000");
         Utils.setValueInConfig("log_level", "DEBUG");
 
-        TestingProcessManager.TestingProcess process = TestingProcessManager.start(args, false);
+        TestingProcessManager.TestingProcess process = TestingProcessManager.start(args, true);
         Main main = process.getProcess();
         setFeatureFlags(main, new EE_FEATURES[] {
                 EE_FEATURES.ACCOUNT_LINKING, EE_FEATURES.MULTI_TENANCY, EE_FEATURES.MFA});
         // We are setting a non-zero initial wait for tests to avoid race condition with the beforeTest process that deletes data in the storage layer
         CronTaskTest.getInstance(main).setInitialWaitTimeInSeconds(ProcessBulkImportUsers.RESOURCE_KEY, 5);
         CronTaskTest.getInstance(main).setIntervalInSeconds(ProcessBulkImportUsers.RESOURCE_KEY, intervalInSeconds);
 
-        process.startProcess();
         assertNotNull(process.checkOrWaitForEvent(ProcessState.PROCESS_STATE.STARTED));
 
         Cronjobs.addCronjob(main, (ProcessBulkImportUsers) main.getResourceDistributor().getResource(new TenantIdentifier(null, null, null), ProcessBulkImportUsers.RESOURCE_KEY));
diff --git a/src/test/java/io/supertokens/test/bulkimport/ProcessBulkImportUsersCronJobTest.java b/src/test/java/io/supertokens/test/bulkimport/ProcessBulkImportUsersCronJobTest.java
@@ -531,6 +531,7 @@ public void shouldDeleteEverythingFromTheDBIfAnythingFailsOnMultipleThreads() th
         Main main = process.getProcess();
 
         BulkImportTestUtils.createTenants(main);
+        Thread.sleep(1000);
 
         BulkImportSQLStorage storage = (BulkImportSQLStorage) StorageLayer.getStorage(main);
         AppIdentifier appIdentifier = new AppIdentifier(null, null);

Original file line number	Diff line number	Diff line change
`@@ -402,6 +402,7 @@ public static void createTablesIfNotExists(Start start, Main main) throws SQLExc`
`402`	`402`	`update(start, UserRolesQueries.getQueryToCreateRolePermissionsTable(start), NO_OP_SETTER);`
`403`	`403`	`// index`
`404`	`404`	`update(start, UserRolesQueries.getQueryToCreateRolePermissionsPermissionIndex(start), NO_OP_SETTER);`
	`405`	`+ update(start, UserRolesQueries.getQueryToCreateUserRolesUserIdAppIdIndex(start), NO_OP_SETTER);`
`405`	`406`	`}`
`406`	`407`
`407`	`408`	`if (!doesTableExists(start, Config.getConfig(start).getUserRolesTable())) {`