Add #addAll() replacing & ifNotExists methods for auto whitelist

voidpointer0x00 · Apr 15, 2023 · a72b349 · a72b349
1 parent ea47bb9
commit a72b349
Showing 1 changed file with 46 additions and 5 deletions.
diff --git a/src/main/java/voidpointer/spigot/voidwhitelist/storage/db/OrmliteWhitelistService.java b/src/main/java/voidpointer/spigot/voidwhitelist/storage/db/OrmliteWhitelistService.java
@@ -222,9 +222,9 @@ private Optional<Whitelistable> onFindException(final Throwable thrown) {
     public CompletableFuture<Integer> addAllIfNotExists(final Collection<Whitelistable> all) {
         if (all.isEmpty())
             return completedFuture(0);
-        return supplyAsync(() -> addAll((model, addedInTotal) -> {
-            if (!ormliteDatabase.getWhitelistDao().idExists(model.getUniqueId())) {
-                ormliteDatabase.getWhitelistDao().create(model);
+        return supplyAsync(() -> addAll((whitelistableModel, addedInTotal) -> {
+            if (!ormliteDatabase.getWhitelistDao().idExists(whitelistableModel.getUniqueId())) {
+                ormliteDatabase.getWhitelistDao().create(whitelistableModel);
                 addedInTotal.increment();
             }
         }, all));
@@ -233,8 +233,8 @@ public CompletableFuture<Integer> addAllIfNotExists(final Collection<Whitelistab
     public CompletableFuture<Integer> addAllReplacing(final Collection<Whitelistable> all) {
         if (all.isEmpty())
             return completedFuture(0);
-        return supplyAsync(() -> addAll((model, addedInTotal) -> {
-            ormliteDatabase.getWhitelistDao().createOrUpdate(model);
+        return supplyAsync(() -> addAll((whitelistableModel, addedInTotal) -> {
+            ormliteDatabase.getWhitelistDao().createOrUpdate(whitelistableModel);
             addedInTotal.increment();
         }, all));
     }
@@ -260,6 +260,47 @@ private int addAll(
         }
     }
 
+    public CompletableFuture<Integer> addAllAutoIfNotExists(final Collection<TimesAutoWhitelistedNumber> all) {
+        if (all.isEmpty())
+            return completedFuture(0);
+        return supplyAsync(() -> addAllAuto((timesAutoWhitelistedModel, addedInTotal) -> {
+            if (!ormliteDatabase.getAutoWhitelistDao().idExists(timesAutoWhitelistedModel.getUniqueId())) {
+                ormliteDatabase.getAutoWhitelistDao().create(timesAutoWhitelistedModel);
+                addedInTotal.increment();
+            }
+        }, all));
+    }
+
+    public CompletableFuture<Integer> addAllAutoReplacing(final Collection<TimesAutoWhitelistedNumber> all) {
+        if (all.isEmpty())
+            return completedFuture(0);
+        return supplyAsync(() -> addAllAuto((timesAutoWhitelistedModel, addedInTotal) -> {
+            ormliteDatabase.getAutoWhitelistDao().createOrUpdate(timesAutoWhitelistedModel);
+            addedInTotal.increment();
+        }, all));
+    }
+
+    private int addAllAuto(
+            final CheckedBiConsumer<TimesAutoWhitelistedNumberModel, MutableInt> addFunction,
+            final Collection<TimesAutoWhitelistedNumber> allAuto) {
+        final MutableInt addedInTotal = new MutableInt();
+        try {
+            requireConnection();
+            return ormliteDatabase.getAutoWhitelistDao().callBatchTasks(() -> {
+                for (final TimesAutoWhitelistedNumber timesAutoWhitelisted : allAuto) {
+                    if (timesAutoWhitelisted instanceof TimesAutoWhitelistedNumberModel)
+                        addFunction.consume((TimesAutoWhitelistedNumberModel) timesAutoWhitelisted, addedInTotal);
+                    else
+                        addFunction.consume(TimesAutoWhitelistedNumberModel.copyOf(timesAutoWhitelisted), addedInTotal);
+                }
+                return addedInTotal.intValue();
+            });
+        } catch (final Exception exception) {
+            log.warn("Unable to add all whitelistable entities: {0}", exception.getMessage());
+            return addedInTotal.intValue();
+        }
+    }
+
     @Override public CompletableFuture<Optional<Whitelistable>> add(
             final UUID uuid, final String name, final Date expiresAt, final int timesAutoWhitelisted) {
         return supplyAsync(() -> {