Store all albums.

[Sone.git] / src / main / java / net / pterodactylus / sone / database / memory / MemoryDatabase.java

diff --git a/src/main/java/net/pterodactylus/sone/database/memory/MemoryDatabase.java b/src/main/java/net/pterodactylus/sone/database/memory/MemoryDatabase.java
index 9edab57..e311a5e 100644 (file)
--- a/src/main/java/net/pterodactylus/sone/database/memory/MemoryDatabase.java
+++ b/src/main/java/net/pterodactylus/sone/database/memory/MemoryDatabase.java
@@ -33,8 +33,6 @@ import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
-import java.util.SortedSet;
-import java.util.TreeSet;
 import java.util.concurrent.locks.ReadWriteLock;
 import java.util.concurrent.locks.ReentrantReadWriteLock;
 
@@ -42,8 +40,8 @@ import net.pterodactylus.sone.data.Album;
 import net.pterodactylus.sone.data.Image;
 import net.pterodactylus.sone.data.Post;
 import net.pterodactylus.sone.data.PostReply;
-import net.pterodactylus.sone.data.Reply;
 import net.pterodactylus.sone.data.Sone;
+import net.pterodactylus.sone.data.impl.DefaultSoneBuilder;
 import net.pterodactylus.sone.database.Database;
 import net.pterodactylus.sone.database.DatabaseException;
 import net.pterodactylus.sone.database.PostDatabase;
@@ -216,7 +214,19 @@ public class MemoryDatabase extends AbstractService implements Database {
 
        @Override
        public SoneBuilder newSoneBuilder() {
-               return null;
+               return new DefaultSoneBuilder(this) {
+                       @Override
+                       public Sone build(Optional<SoneCreated> soneCreated) throws IllegalStateException {
+                               Sone sone = super.build(soneCreated);
+                               lock.writeLock().lock();
+                               try {
+                                       sones.put(sone.getId(), sone);
+                               } finally {
+                                       lock.writeLock().unlock();
+                               }
+                               return sone;
+                       }
+               };
        }
 
        //
@@ -526,7 +536,9 @@ public class MemoryDatabase extends AbstractService implements Database {
                lock.writeLock().lock();
                try {
                        allAlbums.put(album.getId(), album);
-                       albumChildren.put(album.getParent().getId(), album.getId());
+                       if (!album.isRoot()) {
+                               albumChildren.put(album.getParent().getId(), album.getId());
+                       }
                } finally {
                        lock.writeLock().unlock();
                }