From 639ef0db4d0b20b3cb4857f488c207690d02af56 Mon Sep 17 00:00:00 2001
From: Ankit Tiwari <ankucodes@gmail.com>
Date: Tue, 27 Feb 2024 16:12:34 +0530
Subject: [PATCH] fix: PR changes

---
 .../storage/postgresql/QueryExecutorTemplate.java  | 12 ++++++++++++
 .../io/supertokens/storage/postgresql/Start.java   |  4 ++--
 .../postgresql/queries/BulkImportQueries.java      | 14 ++++++++++----
 3 files changed, 24 insertions(+), 6 deletions(-)

diff --git a/src/main/java/io/supertokens/storage/postgresql/QueryExecutorTemplate.java b/src/main/java/io/supertokens/storage/postgresql/QueryExecutorTemplate.java
index db0c9785..098e010f 100644
--- a/src/main/java/io/supertokens/storage/postgresql/QueryExecutorTemplate.java
+++ b/src/main/java/io/supertokens/storage/postgresql/QueryExecutorTemplate.java
@@ -51,6 +51,18 @@ static int update(Start start, String QUERY, PreparedStatementValueSetter setter
         }
     }
 
+    static <T> T update(Start start, String QUERY, PreparedStatementValueSetter setter, ResultSetValueExtractor<T> mapper)
+            throws SQLException, StorageQueryException {
+        try (Connection con = ConnectionPool.getConnection(start)) {
+            try (PreparedStatement pst = con.prepareStatement(QUERY)) {
+                setter.setValues(pst);
+                try (ResultSet result = pst.executeQuery()) {
+                    return mapper.extract(result);
+                }
+            }
+        }
+    }
+
     static int update(Connection con, String QUERY, PreparedStatementValueSetter setter)
             throws SQLException, StorageQueryException {
         try (PreparedStatement pst = con.prepareStatement(QUERY)) {
diff --git a/src/main/java/io/supertokens/storage/postgresql/Start.java b/src/main/java/io/supertokens/storage/postgresql/Start.java
index 0e5bd493..e0bcd61f 100644
--- a/src/main/java/io/supertokens/storage/postgresql/Start.java
+++ b/src/main/java/io/supertokens/storage/postgresql/Start.java
@@ -3086,9 +3086,9 @@ public void updateBulkImportUserStatus_Transaction(AppIdentifier appIdentifier,
     }
 
     @Override
-    public void deleteBulkImportUsers(AppIdentifier appIdentifier, @Nonnull String[] bulkImportUserIds) throws StorageQueryException {
+    public List<String> deleteBulkImportUsers(AppIdentifier appIdentifier, @Nonnull String[] bulkImportUserIds) throws StorageQueryException {
         try {
-            BulkImportQueries.deleteFailedBulkImportUsers(this, appIdentifier, bulkImportUserIds);
+            return BulkImportQueries.deleteBulkImportUsers(this, appIdentifier, bulkImportUserIds);
         } catch (SQLException e) {
             throw new StorageQueryException(e);
         }
diff --git a/src/main/java/io/supertokens/storage/postgresql/queries/BulkImportQueries.java b/src/main/java/io/supertokens/storage/postgresql/queries/BulkImportQueries.java
index 1c893f72..51231dc2 100644
--- a/src/main/java/io/supertokens/storage/postgresql/queries/BulkImportQueries.java
+++ b/src/main/java/io/supertokens/storage/postgresql/queries/BulkImportQueries.java
@@ -169,9 +169,9 @@ public static List<BulkImportUserInfo> getBulkImportUsers(Start start, AppIdenti
         });
     }
 
-    public static void deleteFailedBulkImportUsers(Start start, AppIdentifier appIdentifier, @Nonnull String[] bulkImportUserIds) throws SQLException, StorageQueryException {
+    public static List<String> deleteBulkImportUsers(Start start, AppIdentifier appIdentifier, @Nonnull String[] bulkImportUserIds) throws SQLException, StorageQueryException {
         if (bulkImportUserIds.length == 0) {
-            return;
+            return new ArrayList<>();
         }
 
         String baseQuery =  "DELETE FROM " + Config.getConfig(start).getBulkImportUsersTable();
@@ -190,14 +190,20 @@ public static void deleteFailedBulkImportUsers(Start start, AppIdentifier appIde
             queryBuilder.append("?");
             parameters.add(bulkImportUserIds[i]);
         }
-        queryBuilder.append(")");
+        queryBuilder.append(") RETURNING id");
 
         String query = queryBuilder.toString();
 
-        update(start, query, pst -> {
+        return update(start, query, pst -> {
             for (int i = 0; i < parameters.size(); i++) {
                 pst.setObject(i + 1, parameters.get(i));
             }
+        }, result -> {
+            List<String> deletedUserIds = new ArrayList<>();
+            while (result.next()) {
+                deletedUserIds.add(result.getString("id"));
+            }
+            return deletedUserIds;
         });
     }
     private static class BulkImportUserInfoRowMapper implements RowMapper<BulkImportUserInfo, ResultSet> {