From cfd664cfd476a2d5f31f57ee2f244c8c9ce478d5 Mon Sep 17 00:00:00 2001 From: Veselin Nikolov Date: Thu, 5 Dec 2024 10:55:18 +0000 Subject: [PATCH] Use UserCatalog#remove method to clean graph stores Co-authored-by: Ioannis Panagiotas --- .../gds/core/loading/GraphStoreCatalog.java | 18 ++++++++++++------ 1 file changed, 12 insertions(+), 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 index 5fb80d948e..e128d3e47a 100644 --- 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 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 streamGraphStores(String userName) { return graphsByName .values()