fix: Fix tests (#240)

* fix: fixing tests

* fix: fixing tests

* fix: review fix

Author: tamassoltesz

src/main/java/io/supertokens/storage/postgresql/Start.java

@@ -3590,22 +3590,33 @@ public int getDbActivityCount(String dbname) throws SQLException, StorageQueryEx
     @Override
     public void addBulkImportUsers(AppIdentifier appIdentifier, List<BulkImportUser> users)
             throws StorageQueryException,
-            TenantOrAppNotFoundException,
-            io.supertokens.pluginInterface.bulkimport.exceptions.DuplicateUserIdException {
+            TenantOrAppNotFoundException {
         try {
-            BulkImportQueries.insertBulkImportUsers(this, appIdentifier, users);
-        } catch (SQLException e) {
-            if (e instanceof PSQLException) {
-                ServerErrorMessage serverErrorMessage = ((PSQLException) e).getServerErrorMessage();
-                if (isPrimaryKeyError(serverErrorMessage, Config.getConfig(this).getBulkImportUsersTable())) {
-                    throw new io.supertokens.pluginInterface.bulkimport.exceptions.DuplicateUserIdException();
-                }
-                if (isForeignKeyConstraintError(serverErrorMessage, Config.getConfig(this).getBulkImportUsersTable(),
-                        "app_id")) {
-                    throw new TenantOrAppNotFoundException(appIdentifier);
+            this.startTransaction(con -> {
+                try {
+                    BulkImportQueries.insertBulkImportUsers_Transaction(this, (Connection) con.getConnection(), appIdentifier, users);
+                } catch (SQLException e) {
+                    if (e instanceof PSQLException) {
+                        ServerErrorMessage serverErrorMessage = ((PSQLException) e).getServerErrorMessage();
+                        if (isPrimaryKeyError(serverErrorMessage, Config.getConfig(this).getBulkImportUsersTable())) {
+                            throw new StorageTransactionLogicException(new io.supertokens.pluginInterface.bulkimport.exceptions.DuplicateUserIdException());
+                        }
+                        if (isForeignKeyConstraintError(serverErrorMessage, Config.getConfig(this).getBulkImportUsersTable(),
+                                "app_id")) {
+                            throw new TenantOrAppNotFoundException(appIdentifier);
+                        }
+                    }
                 }
+                return null;
+            });
+        } catch (StorageTransactionLogicException e) {
+            if(e.actualException instanceof StorageQueryException) {
+                throw (StorageQueryException) e.actualException;
+            } else {
+                throw new StorageQueryException(e.actualException);
             }
         }
+
     }
 
     @Override

src/main/java/io/supertokens/storage/postgresql/queries/BulkImportQueries.java

@@ -74,31 +74,24 @@ public static String getQueryToCreatePaginationIndex2(Start start) {
                 + Config.getConfig(start).getBulkImportUsersTable() + " (app_id, created_at DESC, id DESC)";
     }
 
-    public static void insertBulkImportUsers(Start start, AppIdentifier appIdentifier, List<BulkImportUser> users)
+    public static void insertBulkImportUsers_Transaction(Start start, Connection connection, AppIdentifier appIdentifier, List<BulkImportUser> users)
             throws SQLException, StorageQueryException {
-        StringBuilder queryBuilder = new StringBuilder(
-                "INSERT INTO " + Config.getConfig(start).getBulkImportUsersTable() + " (id, app_id, raw_data, created_at, updated_at) VALUES ");
-
-        int userCount = users.size();
-
-        for (int i = 0; i < userCount; i++) {
-            queryBuilder.append(" (?, ?, ?, ?, ?)");
-
-            if (i < userCount - 1) {
-                queryBuilder.append(",");
-            }
+        String queryBuilder = "INSERT INTO " + Config.getConfig(start).getBulkImportUsersTable() +
+                " (id, app_id, raw_data, created_at, updated_at) VALUES "
+                + " (?, ?, ?, ?, ?)";
+
+        List<PreparedStatementValueSetter> valueSetters = new ArrayList<>();
+        for (BulkImportUser user : users) {
+            valueSetters.add(pst -> {
+                pst.setString(1, user.id);
+                pst.setString(2, appIdentifier.getAppId());
+                pst.setString(3, user.toRawDataForDbStorage());
+                pst.setLong(4, System.currentTimeMillis());
+                pst.setLong(5, System.currentTimeMillis());
+            });
         }
 
-        update(start, queryBuilder.toString(), pst -> {
-            int parameterIndex = 1;
-            for (BulkImportUser user : users) {
-                pst.setString(parameterIndex++, user.id);
-                pst.setString(parameterIndex++, appIdentifier.getAppId());
-                pst.setString(parameterIndex++, user.toRawDataForDbStorage());
-                pst.setLong(parameterIndex++, System.currentTimeMillis());
-                pst.setLong(parameterIndex++, System.currentTimeMillis());
-            }
-        });
+        executeBatch(connection, queryBuilder, valueSetters);
     }
 
     public static void updateBulkImportUserStatus_Transaction(Start start, Connection con, AppIdentifier appIdentifier,
@@ -178,18 +171,20 @@ public static List<BulkImportUser> getBulkImportUsersAndChangeStatusToProcessing
                 }
 
                 String updateQuery = "UPDATE " + Config.getConfig(start).getBulkImportUsersTable()
-                        + " SET status = ?, updated_at = ? WHERE app_id = ? AND id IN (" + Utils
-                                .generateCommaSeperatedQuestionMarks(bulkImportUsers.size()) + ")";
-
-                update(sqlCon, updateQuery, pst -> {
-                    int index = 1;
-                    pst.setString(index++, BULK_IMPORT_USER_STATUS.PROCESSING.toString());
-                    pst.setLong(index++, System.currentTimeMillis());
-                    pst.setString(index++, appIdentifier.getAppId());
-                    for (BulkImportUser user : bulkImportUsers) {
-                        pst.setObject(index++, user.id);
-                    }
-                });
+                        + " SET status = ?, updated_at = ? WHERE app_id = ? AND id = ?";
+
+                List<PreparedStatementValueSetter> updateSetters = new ArrayList<>();
+                for(BulkImportUser user : bulkImportUsers){
+                    updateSetters.add(pst -> {
+                        pst.setString(1, BULK_IMPORT_USER_STATUS.PROCESSING.toString());
+                        pst.setLong(2, System.currentTimeMillis());
+                        pst.setString(3, appIdentifier.getAppId());
+                        pst.setObject(4, user.id);
+                    });
+                }
+
+                executeBatch(sqlCon, updateQuery, updateSetters);
+
                 return bulkImportUsers;
             } catch (SQLException throwables) {
                 throw new StorageTransactionLogicException(throwables);

src/test/java/io/supertokens/storage/postgresql/test/OneMillionUsersTest.java

@@ -988,9 +988,9 @@ public void testWithOneMillionUsers() throws Exception {
             return;
         }
 
-        Main main = startCronProcess(String.valueOf(NUM_THREADS));
+        Main main = startCronProcess(String.valueOf(4)); // ci uses instances with 4 cores..
 
-        int NUMBER_OF_USERS_TO_UPLOAD = 1000000; // million
+        int NUMBER_OF_USERS_TO_UPLOAD = 500000; // half million
 
         if (StorageLayer.getBaseStorage(main).getType() != STORAGE_TYPE.SQL || StorageLayer.isInMemDb(main)) {
             return;