Add Reference#onUpdate (#17)

Co-authored-by: Sammy Aknan <sammy.aknan@gmail.com>

Author: Leguan16

src/main/java/net/staticstudios/data/Reference.java

@@ -7,6 +7,9 @@
 import net.staticstudios.data.key.DataKey;
 import net.staticstudios.data.util.DataDoesNotExistException;
 import net.staticstudios.data.util.DeletionStrategy;
+import net.staticstudios.data.util.ValueUpdate;
+import net.staticstudios.data.util.ValueUpdateHandler;
+import net.staticstudios.utils.ThreadUtils;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 
@@ -57,6 +60,18 @@ public static <T extends UniqueData> Reference<T> foreign(UniqueData holder, Cla
         return new Reference<>(holder, clazz, PersistentValue.foreign(holder, UUID.class, schemaTableForeignIdColumn, thisForeignIdColumn));
     }
 
+    /**
+     * Adds an update handler for the backing id.
+     * This is useful if logic is needed when the Reference is set or unlinked.
+     *
+     * @param updateHandler the update handler
+     * @return this
+     */
+    public Reference<T> onUpdate(ValueUpdateHandler<UUID> updateHandler) {
+        getBackingValue().onUpdate(updateHandler);
+        return this;
+    }
+
     @Override
     public Reference<T> deletionStrategy(DeletionStrategy strategy) {
         this.deletionStrategy = strategy;
@@ -165,6 +180,8 @@ public UniqueData getRootHolder() {
         return holder.getRootHolder();
     }
 
+
+
     /**
      * Get the backing {@link PersistentValue} object that stores the foreign id.
      * This is for internal use only.

src/main/java/net/staticstudios/data/impl/PersistentValueManager.java

@@ -198,7 +198,7 @@ public void insertInDatabase(Connection connection, UniqueData holder, List<Init
         for (String idSchemaTable : initialDataMap.keySet()) {
             String idColumn = idSchemaTable.split("\\.", 2)[0];
             String schemaTable = idSchemaTable.split("\\.", 2)[1];
-            Collection<InitialPersistentValue> initialDataValues = initialDataMap.get(idSchemaTable);
+            List<InitialPersistentValue> initialDataValues = new ArrayList<>(initialDataMap.get(idSchemaTable));
             initialDataValues.removeIf(i -> i.getValue().getColumn().equals(idColumn));
             List<InitialPersistentValue> overwriteExisting = new ArrayList<>();
             for (InitialPersistentValue initial : initialDataValues) {

src/test/java/net/staticstudios/data/ReferenceTest.java

@@ -24,6 +24,10 @@ public void init() {
                         id uuid primary key,
                         favorite_user_id uuid
                     );
+                    create table if not exists snapchat.user_meta (
+                        id uuid primary key,
+                        update_called integer
+                    );
                     create table if not exists snapchat.user_settings (
                         user_id uuid primary key,
                         enable_friend_requests boolean not null
@@ -101,4 +105,23 @@ public void testSetReference() {
 
         assertEquals(favoriteUser, user.getFavoriteUser());
     }
+
+    @RetryingTest(5)
+    public void testReferenceUpdateHandler() {
+        MockEnvironment environment = getMockEnvironments().getFirst();
+        DataManager dataManager = environment.dataManager();
+
+        SnapchatUser user = SnapchatUser.createSync(dataManager);
+        SnapchatUser favoriteUser = SnapchatUser.createSync(dataManager);
+
+        assertEquals(0, user.getUpdateCalled());
+        user.setFavoriteUser(favoriteUser);
+
+        waitForDataPropagation();
+        assertEquals(1, user.getUpdateCalled());
+
+        user.setFavoriteUser(null);
+        waitForDataPropagation();
+        assertEquals(2, user.getUpdateCalled());
+    }
 }

src/test/java/net/staticstudios/data/mock/reference/SnapchatUser.java

@@ -1,15 +1,21 @@
 package net.staticstudios.data.mock.reference;
 
 import net.staticstudios.data.DataManager;
+import net.staticstudios.data.PersistentValue;
 import net.staticstudios.data.Reference;
 import net.staticstudios.data.UniqueData;
 import org.jetbrains.annotations.Nullable;
 
 import java.util.UUID;
 
 public class SnapchatUser extends UniqueData {
+    private final PersistentValue<Integer> updateCalled = PersistentValue.foreign(this, Integer.class, "snapchat.user_meta.update_called", "id")
+            .withDefault(0);
     private final Reference<SnapchatUserSettings> settings = Reference.of(this, SnapchatUserSettings.class, "id");
-    private final Reference<SnapchatUser> favoriteUser = Reference.of(this, SnapchatUser.class, "favorite_user_id");
+    private final Reference<SnapchatUser> favoriteUser = Reference.of(this, SnapchatUser.class, "favorite_user_id")
+            .onUpdate(update -> {
+                updateCalled.set(updateCalled.get() + 1);
+            });
 
     private SnapchatUser(DataManager dataManager, UUID id) {
         super(dataManager, "snapchat", "users", id);
@@ -34,4 +40,8 @@ public SnapchatUserSettings getSettings() {
     public void setFavoriteUser(@Nullable SnapchatUser favoriteUser) {
         this.favoriteUser.set(favoriteUser);
     }
+
+    public Integer getUpdateCalled() {
+        return updateCalled.get();
+    }
 }