Use UserCatalog#remove method to clean graph stores

Co-authored-by: Ioannis Panagiotas <[email protected]>
neo4j · Dec 5, 2024 · cfd664c · cfd664c
1 parent aaf8e52
commit cfd664c
Showing 1 changed file with 12 additions and 6 deletions.
diff --git a/core/src/main/java/org/neo4j/gds/core/loading/GraphStoreCatalog.java b/core/src/main/java/org/neo4j/gds/core/loading/GraphStoreCatalog.java
@@ -267,11 +267,21 @@ public static void setDegreeDistribution(
     }
 
     public static void removeAllLoadedGraphs() {
-        userCatalogs.clear();
+        userCatalogs.forEach((username, userCatalog) -> {
+            userCatalog.graphsByName.forEach(((userCatalogKey, graphStoreCatalogEntry) -> {
+                userCatalog.remove(userCatalogKey, __ -> {}, false);
+            }));
+        });
     }
 
     public static void removeAllLoadedGraphs(DatabaseId databaseId) {
-        userCatalogs.forEach((user, userCatalog) -> userCatalog.remove(databaseId.databaseName()));
+        userCatalogs.forEach((user, userCatalog) -> {
+            userCatalog.graphsByName.forEach(((userCatalogKey, graphStoreCatalogEntry) -> {
+                if (databaseId.databaseName().equals(userCatalogKey.databaseName())) {
+                    userCatalog.remove(userCatalogKey, __ -> {}, false);
+                }
+            }));
+        });
     }
 
     public static Collection<GraphStoreCatalogEntry> getGraphStores(String username) {
@@ -430,10 +440,6 @@ private boolean remove(
                 .orElse(Boolean.FALSE);
         }
 
-        private void remove(String databaseName) {
-            graphsByName.keySet().removeIf(userCatalogKey -> userCatalogKey.databaseName().equals(databaseName));
-        }
-
         private Stream<GraphStoreCatalogEntryWithUsername> streamGraphStores(String userName) {
             return graphsByName
                 .values()