Add count method to ModelCatalog

Author: FlorentinD

model-catalog-api/src/main/java/org/neo4j/gds/core/model/ModelCatalog.java

@@ -44,6 +44,8 @@ <D, C extends ModelConfig, I extends ToMapConvertible> Model<D, C, I> get(
 
     Stream<Model<?, ?, ?>> getAllModels();
 
+    long modelsCount();
+
     boolean exists(String username, String modelName);
 
     Model<?, ?, ?> dropOrThrow(String username, String modelName);

model-catalog-api/src/main/java/org/neo4j/gds/core/model/UserCatalog.java

@@ -56,4 +56,6 @@ <D, C extends ModelConfig, I extends ToMapConvertible> Model<D, C, I> get(
     void removeAllLoadedModels();
 
     void verifyModelCanBeStored(String modelName, String modelType);
+
+    long size();
 }

open-model-catalog/src/main/java/org/neo4j/gds/core/model/OpenModelCatalog.java

@@ -92,6 +92,11 @@ public <D, C extends ModelConfig, I extends ToMapConvertible> Model<D, C, I> get
             .flatMap(entry -> entry.getValue().streamModels());
     }
 
+    @Override
+    public long modelsCount() {
+        return userCatalogs.values().stream().mapToLong(OpenUserCatalog::size).sum();
+    }
+
     @Override
     public boolean exists(String username, String modelName) {
         return getUserCatalog(username).exists(modelName);

open-model-catalog/src/main/java/org/neo4j/gds/core/model/OpenUserCatalog.java

@@ -120,6 +120,11 @@ public void verifyModelCanBeStored(String modelName, String modelType) {
         verifyModelsLimit(modelType);
     }
 
+    @Override
+    public long size() {
+        return userModels.size();
+    }
+
     private <D, C extends ModelConfig, I extends ToMapConvertible> Model<D, C, I> get(
         Model<?, ?, ?> model,
         Class<D> dataClass,

open-model-catalog/src/test/java/org/neo4j/gds/core/model/OpenModelCatalogTest.java

@@ -33,6 +33,7 @@
 import java.util.Map;
 import java.util.NoSuchElementException;
 
+import static org.assertj.core.api.Assertions.assertThat;
 import static org.assertj.core.api.Assertions.assertThatThrownBy;
 import static org.junit.jupiter.api.Assertions.assertDoesNotThrow;
 import static org.junit.jupiter.api.Assertions.assertEquals;
@@ -156,6 +157,47 @@ void shouldStoreModels() {
         );
     }
 
+    @Test
+    void shouldCountModels() {
+        var model1 = Model.of(
+            "testAlgo1",
+            GRAPH_SCHEMA,
+            "modelData1",
+            TestTrainConfig.of(USERNAME, "testModel1"),
+            Map::of
+        );
+
+        var model2 = Model.of(
+            "testAlgo2",
+            GRAPH_SCHEMA,
+            1337L,
+            TestTrainConfig.of(USERNAME, "testModel2"),
+            Map::of
+        );
+
+        var publicModel = Model.of(
+            "testAlgo2",
+            GRAPH_SCHEMA,
+            1337L,
+            TestTrainConfig.of("anotherUser", "testModel2"),
+            Map::of
+        );
+
+        assertThat(modelCatalog.modelsCount()).isEqualTo(0);
+
+        modelCatalog.set(model1);
+
+        assertThat(modelCatalog.modelsCount()).isEqualTo(1);
+
+        modelCatalog.set(model2);
+
+        assertThat(modelCatalog.modelsCount()).isEqualTo(2);
+
+        modelCatalog.set(publicModel);
+
+        assertThat(modelCatalog.modelsCount()).isEqualTo(3);
+    }
+
     @Test
     void shouldThrowWhenTryingToGetOtherUsersModel() {
         modelCatalog.set(TEST_MODEL);